diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md index f630f698a141..b1bfabeb578f 100644 --- a/.agents/skills/model-microbenchmark/SKILL.md +++ b/.agents/skills/model-microbenchmark/SKILL.md @@ -7,8 +7,9 @@ description: Time a vLLM model's forward pass on exact batch shapes (decode, pre Builds the batch without the scheduler, runs it through MRV2's real `execute_model`, and times the forward as a one-shot CUDA graph (prefill -included). Sampling, the speculator, and CPU step overhead are excluded, so -expect results somewhat below a serving step. +included). The configured drafter is included by default; use +`--exclude-drafter` for target-only timing. All draft tokens are accepted. +Target sampling, rejection sampling, and CPU step overhead are excluded. ## Usage @@ -44,7 +45,8 @@ Key flags (plus all `EngineArgs`): - `--kv-len-jitter 0.5` (default): prefixes drawn within ±50%, same on every DP rank, capped at `max_model_len`. Use `0` for exact lengths. -- `--routing auto|real|balanced`: `auto` = balanced with dummy weights. +- `--routing auto|real|balanced|zipf`: `auto` = Zipf with dummy weights; + `zipf` uses fixed alpha=1.25 skew (target only, no PCP or grouped top-k). - `--iters 20 --samples 5`, `--output-json out.json` (per-rank latency stats, `samples_us`, graph-vs-eager check). - Defaults: `--load-format dummy`, prefix caching off, no torch.compile. @@ -94,7 +96,7 @@ that the JSON summary excludes. ## Caveats - **Dummy weights lie for data-dependent kernels.** Dummy gates route every - token to the same experts (hence balanced routing). Cross-check suspicious + token to the same experts (hence synthetic routing). Cross-check suspicious kernels with real weights. - **torch.compile models (e.g. Qwen3/3.5) look slow**: no compile here. Use them for correctness only; trust perf for models under `vllm/models/`. diff --git a/.agents/skills/model-recipes/SKILL.md b/.agents/skills/model-recipes/SKILL.md index f6ee22b1946f..5100b3938a00 100644 --- a/.agents/skills/model-recipes/SKILL.md +++ b/.agents/skills/model-recipes/SKILL.md @@ -7,9 +7,9 @@ description: Find vLLM deployment recipes and configs for LLMs and Hugging Face Model serving recipes are maintained in this repository's [recipes/models/ directory](../../../recipes/models). Omni (image, video, -audio) recipes stay in vllm-project/recipes for now (moving to vllm-omni), -and Ascend NPU content is inline in these recipes (moving to vllm-ascend -overlays). The published catalog, rendered by +audio) recipes live in vllm-project/vllm-omni, and Ascend NPU content lives in +vllm-project/vllm-ascend as overlays merged into these recipes. The published +catalog, rendered by [vllm-project/recipes](https://github.com/vllm-project/recipes), is at [recipes.vllm.ai](https://recipes.vllm.ai). diff --git a/.buildkite/amd-disagg/cluster.sh b/.buildkite/amd-disagg/cluster.sh index a8906e6e578b..d71c9eb410df 100644 --- a/.buildkite/amd-disagg/cluster.sh +++ b/.buildkite/amd-disagg/cluster.sh @@ -156,6 +156,21 @@ export MORI_GPU_ARCHS="gfx950" # path fails to come up. export MORI_ENABLE_HOST_PROXY="${MORI_ENABLE_HOST_PROXY:-1}" +# NIXL UCX defaults: RC on ionic_0, GPU-direct RDMA. +if [[ "${KV_CONNECTOR:-}" == Nixl* ]]; then + export UCX_TLS="${UCX_TLS:-rc_verbs,ud_verbs,tcp,self,rocm}" + export UCX_NET_DEVICES="${UCX_NET_DEVICES:-ionic_0:1,lo}" + export UCX_IB_GID_INDEX="${UCX_IB_GID_INDEX:-${_IB_GID_INDEX}}" + export UCX_UD_VERBS_ETH_DGID_CHECK="${UCX_UD_VERBS_ETH_DGID_CHECK:-n}" + export UCX_IB_GPU_DIRECT_RDMA="${UCX_IB_GPU_DIRECT_RDMA:-yes}" + export UCX_RMA_PPLN_ENABLE="${UCX_RMA_PPLN_ENABLE:-y}" + export UCX_PROTO_EMULATION_ENABLE="${UCX_PROTO_EMULATION_ENABLE:-n}" + export UCX_KEEPALIVE_INTERVAL="${UCX_KEEPALIVE_INTERVAL:-inf}" + export UCX_UD_VERBS_RX_QUEUE_LEN="${UCX_UD_VERBS_RX_QUEUE_LEN:-1024}" + export UCX_RCACHE_MAX_UNRELEASED="${UCX_RCACHE_MAX_UNRELEASED:-1024}" + export ACCURACY_NUM_CONCURRENT="${ACCURACY_NUM_CONCURRENT:-1}" +fi + # ----------------------------------------------------------------- benchmark export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/128 2048/128}" export BENCHMARK_CON="${BENCHMARK_CON:-32 64}" diff --git a/.buildkite/amd-disagg/models.yaml b/.buildkite/amd-disagg/models.yaml index 114de22ecd8d..873e2c14f444 100644 --- a/.buildkite/amd-disagg/models.yaml +++ b/.buildkite/amd-disagg/models.yaml @@ -105,6 +105,20 @@ models: tp: "--gpu-memory-utilization 0.9" ep: '--gpu-memory-utilization 0.75 --compilation-config {"cudagraph_mode":"PIECEWISE","custom_ops":["+quant_fp8"]}' + - model: DeepSeek-V4-Pro + env: + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + VLLM_ROCM_USE_AITER_MOE: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + base_flags: "--trust-remote-code --kv-cache-dtype fp8" + prefill: + tp: "--gpu-memory-utilization 0.85" + ep: "--gpu-memory-utilization 0.85 --enforce-eager" + decode: + tp: "--gpu-memory-utilization 0.85" + ep: '--gpu-memory-utilization 0.75 --compilation-config {"cudagraph_mode":"PIECEWISE","custom_ops":["+quant_fp8"]}' + - model: GLM-5.2-FP8 env: VLLM_ROCM_USE_AITER: "1" diff --git a/.buildkite/amd-disagg/pipeline-disagg.yaml b/.buildkite/amd-disagg/pipeline-disagg.yaml index bf1a3a74281c..71c7b8387d86 100644 --- a/.buildkite/amd-disagg/pipeline-disagg.yaml +++ b/.buildkite/amd-disagg/pipeline-disagg.yaml @@ -102,6 +102,15 @@ steps: commands: - MODEL_NAME=MiniMax-M3-MXFP8 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh +- label: "DeepSeek-V4-Pro-PD-1P1D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V4-Pro NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + - label: "GLM-5.2-FP8-PD-1P1D-TP8-MoRIIO-proxy" agents: queue: amd_mi350_ainic @@ -310,6 +319,15 @@ steps: commands: - MODEL_NAME=MiniMax-M3-MXFP8 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh +- label: "DeepSeek-V4-Pro-PD-1P1D-EP8/DP8-WideEP-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V4-Pro NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + - label: "GLM-5.2-FP8-PD-1P1D-EP8/DP8-WideEP-MoRIIO-proxy" agents: queue: amd_mi350_ainic @@ -378,6 +396,15 @@ steps: commands: - MODEL_NAME=MiniMax-M3-MXFP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy WORKLOAD_TIMEOUT_S=10800 ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh +- label: "DeepSeek-V4-Pro-PD-2P2D-EP16/DP16-WideEP-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 1 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-V4-Pro xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + - label: "GLM-5.2-FP8-PD-2P2D-EP16/DP16-WideEP-MoRIIO-vllm-router" agents: queue: amd_mi350_ainic @@ -387,3 +414,58 @@ steps: commands: - MODEL_NAME=GLM-5.2-FP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh +#-------------------------------------- mi350 · disagg NIXL (UCX RDMA) · SLURM --------------------------------------# + +- label: "DeepSeek-V3-PD-1P1D-TP8-NIXL-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 KV_CONNECTOR=NixlConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-2P2D-EP16/DP16-WideEP-NIXL-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 1 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-NIXL-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-1P1D-TP8-NIXL-PUSH-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 KV_CONNECTOR=NixlPushConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-NIXL-PUSH-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlPushConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-2P2D-EP16/DP16-WideEP-NIXL-PUSH-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 270 + concurrency: 1 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=1 KV_CONNECTOR=NixlPushConnector RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh diff --git a/.buildkite/amd-disagg/run-slurm-disagg-test.sh b/.buildkite/amd-disagg/run-slurm-disagg-test.sh index 641133b9ff45..720b11573159 100644 --- a/.buildkite/amd-disagg/run-slurm-disagg-test.sh +++ b/.buildkite/amd-disagg/run-slurm-disagg-test.sh @@ -47,6 +47,7 @@ SHARED_MOUNT="${SHARED_MOUNT:-/data}" LOG_ROOT="${LOG_ROOT:-${SHARED_MOUNT}/${USER:-$(whoami)}/disagg_logs}" DRY_RUN="${DRY_RUN:-0}" MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" +KV_CONNECTOR="${KV_CONNECTOR:-MoRIIOConnector}" WAIT="${WAIT:-0}" @@ -69,7 +70,7 @@ export DISAGG_SCRIPTS_DIR="${STAGED_DIR}" echo "[slurm-submit] staged scripts for compute nodes: ${DISAGG_SCRIPTS_DIR}" >&2 export IMAGE MODEL_NAME WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S -export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE +export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE KV_CONNECTOR export ROUTER_TYPE ROUTER_PORT VLLM_ROUTER_IMAGE # Leave the image's own ionic userspace in place instead of mounting the host's. export SKIP_IONIC_MOUNTS diff --git a/.buildkite/amd-disagg/run_xPyD_disagg.slurm b/.buildkite/amd-disagg/run_xPyD_disagg.slurm index e9bf8014784f..19fbb54856b0 100644 --- a/.buildkite/amd-disagg/run_xPyD_disagg.slurm +++ b/.buildkite/amd-disagg/run_xPyD_disagg.slurm @@ -23,6 +23,7 @@ LOG_ROOT="${LOG_ROOT:-/data/${USER:-$(id -un)}/disagg_logs}" RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy DRY_RUN="${DRY_RUN:-0}" MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" # 0 -> prefill writes (default); 1 -> decode reads KV (read_mode) +KV_CONNECTOR="${KV_CONNECTOR:-MoRIIOConnector}" # MoRIIOConnector | NixlConnector | NixlPushConnector # ---- router / gateway (host-side; also forwarded into the container) -------- # vllm-router -> external `vllm/vllm-router` container started on the rank-0 node (default) @@ -232,7 +233,7 @@ echo "[disagg-pd] LOG_PATH=$LOG_PATH" # rank xP .. end -> decode export IPADDRS -export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE +export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE KV_CONNECTOR export IMAGE SHARED_MOUNT CONTAINER_SCRIPTS DISAGG_SCRIPTS_DIR export CLUSTER_ENV="${CLUSTER_ENV_IN}" export DOCKER_CONT_NAME="vllm_disagg_${MODEL_NAME:-model}_${SLURM_JOB_ID}" @@ -247,6 +248,21 @@ export PREFILL_PORT DECODE_PORT PROXY_PORT ROUTER_METRICS_PORT export SERVE_PORT RPC_PORT KV_PORT LOCAL_PING_PORT export PREFLIGHT PREFLIGHT_MAX_VRAM_PCT PREFLIGHT_GPU_SETTLE_S +# NIXL has no discovery service, so pass the worker URLs. +ROUTER_KV_ARGS="--kv-connector moriio --vllm-discovery-address 0.0.0.0:${PROXY_PING_PORT}" +if [[ "${KV_CONNECTOR}" == Nixl* ]]; then + nixl_startup_timeout=1800 + [[ "${WIDE_EP_MODE}" == "1" && "${xP}" -gt 1 ]] && nixl_startup_timeout=7200 + ROUTER_KV_ARGS="--kv-connector nixl --worker-startup-timeout-secs ${nixl_startup_timeout}" + if [[ "${WIDE_EP_MODE}" == "1" ]]; then + ROUTER_KV_ARGS+=" --prefill http://${IPS[0]}:${SERVE_PORT} --decode http://${IPS[$xP]}:${SERVE_PORT}" + else + for ((i = 0; i < xP; i++)); do ROUTER_KV_ARGS+=" --prefill http://${IPS[$i]}:${PREFILL_PORT}"; done + for ((i = xP; i < xP + yD; i++)); do ROUTER_KV_ARGS+=" --decode http://${IPS[$i]}:${DECODE_PORT}"; done + fi +fi +export ROUTER_KV_ARGS + echo "[disagg-pd] IPADDRS=${IPADDRS}" echo "[disagg-pd] launching ${NUM_NODES} containers (1/node) -> vllm_disagg.sh node" @@ -479,7 +495,7 @@ run_node_container() { -e USER="$(id -un)" \ -e LOGNAME="$(id -un)" \ "$VLLM_ROUTER_IMAGE" \ - bash -lc "exec vllm-router --vllm-pd-disaggregation --kv-connector moriio --vllm-discovery-address 0.0.0.0:$PROXY_PING_PORT --port $ROUTER_PORT --host 0.0.0.0 --intra-node-data-parallel-size $ROUTER_DP_LOCAL --policy $ROUTER_POLICY --prefill-policy $ROUTER_POLICY --decode-policy $ROUTER_POLICY --log-level info >> $LOG_PATH/vllm_router.log 2>&1" \ + bash -lc "exec vllm-router --vllm-pd-disaggregation $ROUTER_KV_ARGS --port $ROUTER_PORT --host 0.0.0.0 --intra-node-data-parallel-size $ROUTER_DP_LOCAL --policy $ROUTER_POLICY --prefill-policy $ROUTER_POLICY --decode-policy $ROUTER_POLICY --log-level info >> $LOG_PATH/vllm_router.log 2>&1" \ && echo "vllm-router container started ok" \ || echo "WARN: failed to start vllm-router container on $(hostname)" } 2>&1 | tee -a "$LOG_PATH/vllm_router_start.log" @@ -528,6 +544,7 @@ run_node_container() { -e HEALTH_TIMEOUT_S="$HEALTH_TIMEOUT_S" \ -e DRY_RUN="$DRY_RUN" \ -e MORIIO_READ_MODE="$MORIIO_READ_MODE" \ + -e KV_CONNECTOR="$KV_CONNECTOR" \ -e ROUTER_TYPE="$ROUTER_TYPE" \ -e ROUTER_PORT="$ROUTER_PORT" \ -e ROUTER_POLICY="$ROUTER_POLICY" \ diff --git a/.buildkite/amd-disagg/vllm_disagg.sh b/.buildkite/amd-disagg/vllm_disagg.sh index 23320aaf6b1b..37f8915c421e 100644 --- a/.buildkite/amd-disagg/vllm_disagg.sh +++ b/.buildkite/amd-disagg/vllm_disagg.sh @@ -102,6 +102,18 @@ moriio_read_mode_kv() { esac } +# NIXL KV config. UCX backend. +nixl_kv_cfg() { + cat <. A cache source only (see ci.hcl), unlike + # the VLLM_MERGE_BASE_COMMIT build arg, which would change the cache key. + if [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]]; then + return 0 + fi + if [[ -n "${VLLM_CACHE_MERGE_BASE_COMMIT:-}" ]]; then + echo "Using provided VLLM_CACHE_MERGE_BASE_COMMIT: ${VLLM_CACHE_MERGE_BASE_COMMIT}" + return 0 + fi + local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" + local base_ref="refs/remotes/vllm-cache-base/${base_branch}" + git fetch --quiet --no-tags --depth=200 origin \ + "+refs/heads/${base_branch}:${base_ref}" 2>/dev/null || true + VLLM_CACHE_MERGE_BASE_COMMIT=$(git merge-base HEAD "${base_ref}" 2>/dev/null || echo "") + if [[ -z "${VLLM_CACHE_MERGE_BASE_COMMIT}" ]]; then + git fetch --quiet --no-tags --deepen=1000 origin \ + "+refs/heads/${base_branch}:${base_ref}" 2>/dev/null || true + VLLM_CACHE_MERGE_BASE_COMMIT=$(git merge-base HEAD "${base_ref}" 2>/dev/null || echo "") + fi + if [[ -n "${VLLM_CACHE_MERGE_BASE_COMMIT}" ]]; then + echo "Computed merge base with ${base_branch} for cache fallback: ${VLLM_CACHE_MERGE_BASE_COMMIT}" + export VLLM_CACHE_MERGE_BASE_COMMIT + else + echo "Could not determine merge base with ${base_branch}; skipping that cache source" + fi +} + print_bake_config() { echo "--- :page_facing_up: Resolved bake configuration" # Write to a temp directory to avoid polluting the repo root (which is the @@ -353,6 +382,7 @@ setup_buildx_builder resolve_parent_commit export PARENT_COMMIT +resolve_merge_base_cache_commit print_bake_config diff --git a/.buildkite/intel_jobs/basic_correctness_intel.yaml b/.buildkite/intel_jobs/basic_correctness_intel.yaml index 7ad8de9a22ce..da6d2ec0f41b 100644 --- a/.buildkite/intel_jobs/basic_correctness_intel.yaml +++ b/.buildkite/intel_jobs/basic_correctness_intel.yaml @@ -20,10 +20,15 @@ steps: - vllm/ - tests/basic_correctness/memory/ - tests/basic_correctness/cpu_offload/ + # MRV2 model state remains in the weights pool on XPU (vllm-project/vllm#63548). + # Re-enable Mamba deep sleep once XPU nested memory pools are supported. commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && pytest -v -s basic_correctness/cpu_offload && - pytest -v -s basic_correctness/memory --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem"' + pytest -v -s + basic_correctness/memory + --deselect="tests/basic_correctness/memory/cumem/test_cumem.py::test_basic_cumem" + --deselect="tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py::test_deep_sleep[yujiepan/mamba2-codestral-v0.1-tiny-random]"' diff --git a/.buildkite/intel_jobs/engine_intel.yaml b/.buildkite/intel_jobs/engine_intel.yaml index a9fa467caee6..7c7c0f02c7e6 100644 --- a/.buildkite/intel_jobs/engine_intel.yaml +++ b/.buildkite/intel_jobs/engine_intel.yaml @@ -38,7 +38,12 @@ steps: - >- 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 jit_monitor/test_hooks.py' + pytest -v -s + engine/test_arg_utils.py + test_sequence.py + test_logger.py + test_vllm_port.py + jit_monitor/test_hooks.py' - label: ":intel: (BMG) Engine (1 GPU)" key: engine-1-gpu @@ -62,7 +67,10 @@ 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 --ignore v1/engine/test_core_engine_actor_manager.py && + pytest -v -s + v1/engine + --ignore v1/engine/test_preprocess_error_handling.py + --ignore v1/engine/test_core_engine_actor_manager.py && pytest -v -s test_config.py' - label: ":intel: (BMG) E2E Core" @@ -88,12 +96,15 @@ steps: - vllm/model_executor/models/qwen3_next - vllm/config/cache.py - vllm/platforms/xpu.py + - vllm/v1/engine/ - tests/v1/e2e/general/test_mamba_prefix_cache.py + - tests/v1/e2e/general/test_streaming_input.py commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py' + pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py && + pytest -v -s v1/e2e/general/test_streaming_input.py' - label: ":intel: (BMG) V1 e2e (2 GPUs)" key: v1-e2e-2-gpus diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index c05e93f38aa5..8b9885583a81 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -77,8 +77,12 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc && - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer + pytest -v -s + entrypoints/serve + --ignore=entrypoints/serve/dev/rpc && + pytest -v -s + entrypoints/rl + --ignore=entrypoints/rl/weight_transfer --ignore=entrypoints/rl/return_routed_experts/test_openai_api.py --ignore=entrypoints/rl/weight_transfer_metrics --deselect="tests/entrypoints/rl/state_transitions/test_pause_resume.py::TestPauseResume::test_clear_cache_preserves_output_and_controls_prefix_cache[MRV1]" && diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index f3a9d58f4357..fdb087971991 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -48,3 +48,29 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s kernels/core/test_uva.py' + +- label: ":intel: (BMG) XPU Sparse MLA Tests" + key: xpu-sparse-mla-tests + timeout_in_minutes: 20 + 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/v1/attention/backends/mla/xpu_mla_sparse.py + - tests/kernels/attention/test_xpu_mla_sparse.py + - tests/v1/attention/test_xpu_mla_sparse_backend.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s + kernels/attention/test_xpu_mla_sparse.py + v1/attention/test_xpu_mla_sparse_backend.py' diff --git a/.buildkite/intel_jobs/lora_intel.yaml b/.buildkite/intel_jobs/lora_intel.yaml index ae5e2d0165cc..3ff650c99752 100644 --- a/.buildkite/intel_jobs/lora_intel.yaml +++ b/.buildkite/intel_jobs/lora_intel.yaml @@ -57,8 +57,10 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && - pytest -v -s lora/test_fused_moe_lora_kernel.py && - pytest -v -s lora/test_moe_lora_align_sum.py --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"' + pytest -v -s lora/test_fused_moe_lora_kernel.py --deselect="lora/test_fused_moe_lora_kernel.py::test_fused_moe_lora_kernel_large_rank_fallback[xpu:0-dtype0-1-True-256]" && + pytest -v -s + lora/test_moe_lora_align_sum.py + --deselect="tests/lora/test_moe_lora_align_sum.py::test_moe_lora_align_block_size_mixed_base_and_lora[1]"' - label: ":intel: (BMG) LoRA Punica Kernels" key: lora-punica-kernels @@ -134,7 +136,11 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && export VLLM_WORKER_MULTIPROC_METHOD=spawn && - pytest -v -s lora/test_quant_model.py --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" && + pytest -v -s + lora/test_quant_model.py + --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model0]" + --deselect="tests/lora/test_quant_model.py::test_quant_model_lora[model1]" + --deselect="tests/lora/test_quant_model.py::test_quant_model_tp_equality[model0]" && pytest -v -s lora/test_transformers_model.py && pytest -v -s lora/test_chatglm3_tp.py && pytest -v -s lora/test_llama_tp.py::test_llama_lora' diff --git a/.buildkite/intel_jobs/misc_intel.yaml b/.buildkite/intel_jobs/misc_intel.yaml index 0cdb7d78b663..f0361f4128c8 100644 --- a/.buildkite/intel_jobs/misc_intel.yaml +++ b/.buildkite/intel_jobs/misc_intel.yaml @@ -30,71 +30,73 @@ steps: pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py && bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh' -- label: ":intel: (BMG) V1 Core + KV + Metrics" - key: v1-core-kv-metrics-intel - timeout_in_minutes: 45 - 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/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/ec_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - commands: - - >- - bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install -r requirements/kv_connectors.txt && - export VLLM_WORKER_MULTIPROC_METHOD=spawn && - cd tests && - pytest -v -s v1/executor && - VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit - --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py' +# Disabled until the KV connectors installation is locked. +# - label: ":intel: (BMG) V1 Core + KV + Metrics" +# key: v1-core-kv-metrics-intel +# timeout_in_minutes: 45 +# 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/v1/core +# - tests/v1/executor +# - tests/v1/kv_offload +# - tests/v1/worker +# - tests/v1/kv_connector/unit +# - tests/v1/ec_connector/unit +# - tests/v1/metrics +# - tests/entrypoints/openai/correctness/test_lmeval.py +# commands: +# - >- +# bash .buildkite/scripts/hardware_ci/run-intel-test.sh +# 'pip install -r requirements/kv_connectors.txt && +# export VLLM_WORKER_MULTIPROC_METHOD=spawn && +# cd tests && +# pytest -v -s v1/executor && +# VLLM_BATCH_INVARIANT=1 pytest -v -s -m "not cpu_test" v1/ec_connector/unit +# --ignore=v1/ec_connector/unit/test_ec_cpu_connector.py' -- label: ":intel: (BMG) EC CPU Connector (XPU)" - key: ec-cpu-connector-xpu - 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/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/worker - - tests/v1/kv_connector/unit - - tests/v1/ec_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - commands: - - >- - bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'pip install -r requirements/kv_connectors.txt && - export VLLM_WORKER_MULTIPROC_METHOD=spawn && - cd tests && - pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py' +# Disabled until the KV connectors installation is locked. +# - label: ":intel: (BMG) EC CPU Connector (XPU)" +# key: ec-cpu-connector-xpu +# 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/ +# - tests/v1/core +# - tests/v1/executor +# - tests/v1/kv_offload +# - tests/v1/worker +# - tests/v1/kv_connector/unit +# - tests/v1/ec_connector/unit +# - tests/v1/metrics +# - tests/entrypoints/openai/correctness/test_lmeval.py +# commands: +# - >- +# bash .buildkite/scripts/hardware_ci/run-intel-test.sh +# 'pip install -r requirements/kv_connectors.txt && +# export VLLM_WORKER_MULTIPROC_METHOD=spawn && +# cd tests && +# pytest -v -s v1/ec_connector/unit/test_ec_cpu_connector.py' - label: ":intel: (BMG) V1 Sample + Logits" key: v1-sample-logits-intel @@ -132,8 +134,14 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'export VLLM_WORKER_MULTIPROC_METHOD=spawn && cd tests && - pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py && - pytest -v -s v1/test_oracle.py v1/test_request.py v1/test_outputs.py && + pytest -v -s + v1/logits_processors + --ignore=v1/logits_processors/test_custom_online.py + --ignore=v1/logits_processors/test_custom_offline.py && + pytest -v -s + v1/test_oracle.py + v1/test_request.py + v1/test_outputs.py && pytest -v -s v1/sample' retry: automatic: @@ -325,9 +333,13 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s detokenizer --ignore=detokenizer/test_stop_strings.py && + pytest -v -s + detokenizer + --ignore=detokenizer/test_stop_strings.py && pytest -v -s -m "not cpu_test" ./multimodal && - pytest -v -s utils_ --ignore=utils_/test_mem_utils.py' + pytest -v -s + utils_ + --ignore=utils_/test_mem_utils.py' - label: ":intel: (BMG) Detokenizer Stop Strings (XPU)" key: detokenizer-stop-strings-xpu diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index c6a59ca99ed0..48f84b348443 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -1,59 +1,61 @@ group: Model Executor Intel depends_on: - image-build-xpu -steps: -- label: ":intel: (BMG) Model Executor" - key: model-executor-intel - timeout_in_minutes: 45 - 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/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 && - 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" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py" --ignore="model_executor/model_loader/test_weight_cache.py" --deselect="tests/model_executor/test_enabled_custom_ops.py::test_topk_gating_enabled_by_cdna"' +steps: [] +# Disabled until the apt packages installation is locked. +# - label: ":intel: (BMG) Model Executor" +# key: model-executor-intel +# timeout_in_minutes: 45 +# 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/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 && +# 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" --ignore="model_executor/model_loader/tensorizer_loader/test_tensorizer.py" --ignore="model_executor/model_loader/test_reload.py" --ignore="model_executor/model_loader/test_weight_cache.py" --deselect="tests/model_executor/test_enabled_custom_ops.py::test_topk_gating_enabled_by_cdna"' -- label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)" - key: tensorizer-model-loader-tests-xpu - 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 && - export VLLM_WORKER_MULTIPROC_METHOD=spawn && - export PYTHONFAULTHANDLER=1 && - cd tests && - pytest -v -s model_executor/model_loader/tensorizer_loader/test_tensorizer.py -m "not slow_test"' +# Disabled until the apt packages installation is locked. +# - label: ":intel: (BMG) Tensorizer Model Loader Tests (XPU)" +# key: tensorizer-model-loader-tests-xpu +# 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 && +# export VLLM_WORKER_MULTIPROC_METHOD=spawn && +# export PYTHONFAULTHANDLER=1 && +# cd tests && +# pytest -v -s model_executor/model_loader/tensorizer_loader/test_tensorizer.py -m "not slow_test"' diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 9a9790d5adc7..9864f103a32a 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -73,7 +73,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && - pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model + pytest -v -s models/multimodal/generation/test_gemma_ragged_audio.py -m core_model && pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model' - label: ":intel: (BMG) Multi-Modal Models (Standard) 4: other + whisper" @@ -97,7 +97,17 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_gemma_ragged_audio.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing' + pytest -v -s + models/multimodal + -m core_model + --ignore models/multimodal/generation/test_common.py + --ignore models/multimodal/generation/test_ultravox.py + --ignore models/multimodal/generation/test_qwen2_5_vl.py + --ignore models/multimodal/generation/test_qwen2_vl.py + --ignore models/multimodal/generation/test_whisper.py + --ignore models/multimodal/generation/test_gemma_ragged_audio.py + --ignore models/multimodal/generation/test_memory_leak.py + --ignore models/multimodal/processing' - label: ":intel: (BMG) Multi-Modal Processor" # 44min key: multi-modal-processor diff --git a/.buildkite/intel_jobs/quantization_intel.yaml b/.buildkite/intel_jobs/quantization_intel.yaml index 035ad11beef2..887945e98377 100644 --- a/.buildkite/intel_jobs/quantization_intel.yaml +++ b/.buildkite/intel_jobs/quantization_intel.yaml @@ -24,8 +24,12 @@ steps: # - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_per_token_kv_cache.py --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]" + 'VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s + tests/quantization/test_per_token_kv_cache.py + --deselect="tests/quantization/test_per_token_kv_cache.py::test_triton_unified_attention_per_token_head_scale[int4-16-128-num_heads0-seq_lens1]" && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_auto_awq.py && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/quantization/test_fp8_per_channel.py - && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s tests/models/quantization/test_fp8_per_channel.py --deselect="tests/models/quantization/test_fp8_per_channel.py::test_fp8_per_channel_logprobs[moe]"' + && VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s + tests/models/quantization/test_fp8_per_channel.py + --deselect="tests/models/quantization/test_fp8_per_channel.py::test_fp8_per_channel_logprobs[moe]"' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 74d6eee4900e..635978a4c652 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -76,63 +76,65 @@ steps: python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 && python3 examples/basic/offline_inference/generate.py --linear-backend torch --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8_per_tensor --enforce-eager --max-model-len 4096 ' - - label: ":intel: (BMG) V1 test" - key: xpu-v1-test - depends_on: - - image-build-xpu - timeout_in_minutes: 70 - device: intel_gpu - agent_tags: - label: production - gpu: 1+ - mem: 16+ - no_plugin: true - env: - REGISTRY: "public.ecr.aws/q9t5s3a7" - REPO: "vllm-ci-test-repo" - source_file_dependencies: - - vllm/ - - .buildkite/intel_jobs/test-intel.yaml - commands: - - >- - bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && pip install dpctl && - pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && - pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py --ignore=v1/engine/test_core_engine_actor_manager.py && - pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && - pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && - pytest -v -s v1/structured_output && - pytest -v -s v1/test_serial_utils.py && - pytest -v -s v1/e2e/general/test_correctness_sliding_window.py && - pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' - parallelism: 2 - - label: ":intel: (BMG) Qwen3-30B + V1 KV Connector" - key: xpu-qwen3-30b-v1-kv-connector - depends_on: - - image-build-xpu - timeout_in_minutes: 90 - device: intel_gpu - agent_tags: - label: production - gpu: 2+ - mem: 24+ - no_plugin: true - env: - REGISTRY: "public.ecr.aws/q9t5s3a7" - REPO: "vllm-ci-test-repo" - source_file_dependencies: - - vllm/ - - .buildkite/intel_jobs/test-intel.yaml - commands: - - >- - bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && - pip install dpctl && - pytest -v -s v1/kv_connector/unit/test_multi_connector.py && - pytest -v -s v1/kv_connector/unit/test_example_connector.py && - cd .. && - VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192' + # Disabled until the dpctl installation is locked. + # - label: ":intel: (BMG) V1 test" + # key: xpu-v1-test + # depends_on: + # - image-build-xpu + # timeout_in_minutes: 70 + # device: intel_gpu + # agent_tags: + # label: production + # gpu: 1+ + # mem: 16+ + # no_plugin: true + # env: + # REGISTRY: "public.ecr.aws/q9t5s3a7" + # REPO: "vllm-ci-test-repo" + # source_file_dependencies: + # - vllm/ + # - .buildkite/intel_jobs/test-intel.yaml + # commands: + # - >- + # bash .buildkite/scripts/hardware_ci/run-intel-test.sh + # 'cd tests && pip install dpctl && + # pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py && + # pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py --ignore=v1/engine/test_core_engine_actor_manager.py && + # pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" && + # pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py && + # pytest -v -s v1/structured_output && + # pytest -v -s v1/test_serial_utils.py && + # pytest -v -s v1/e2e/general/test_correctness_sliding_window.py && + # pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py --deselect=tests/v1/spec_decode/test_mtp.py::test_glm_mtp_defers_lm_head && + # pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --deselect=tests/v1/kv_connector/unit/test_nixl_connector.py::test_abort_timeout_on_prefiller[ray]' + # parallelism: 2 + # Disabled until the dpctl installation is locked. + # - label: ":intel: (BMG) Qwen3-30B + V1 KV Connector" + # key: xpu-qwen3-30b-v1-kv-connector + # depends_on: + # - image-build-xpu + # timeout_in_minutes: 90 + # device: intel_gpu + # agent_tags: + # label: production + # gpu: 2+ + # mem: 24+ + # no_plugin: true + # env: + # REGISTRY: "public.ecr.aws/q9t5s3a7" + # REPO: "vllm-ci-test-repo" + # source_file_dependencies: + # - vllm/ + # - .buildkite/intel_jobs/test-intel.yaml + # commands: + # - >- + # bash .buildkite/scripts/hardware_ci/run-intel-test.sh + # 'cd tests && + # pip install dpctl && + # pytest -v -s v1/kv_connector/unit/test_multi_connector.py && + # pytest -v -s v1/kv_connector/unit/test_example_connector.py && + # cd .. && + # VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192' - label: ":intel: (BMG) Server test" key: xpu-server-test depends_on: @@ -178,7 +180,10 @@ steps: bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && pytest -v -s quantization/test_auto_round.py && - pytest -v -s quantization/test_online.py --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]"' + pytest -v -s + quantization/test_online.py + --deselect=tests/quantization/test_online.py::test_online_quantization_loads_real_weights + --deselect="tests/quantization/test_online.py::test_online_quantization[False-requantization_mxfp8_ptcp_fp8]"' - label: ":intel: (BMG) Compressed tensors FP8 test" key: xpu-compressed-tensors-fp8-test depends_on: diff --git a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py index 6b2e293ba981..a3c8b648d777 100644 --- a/.buildkite/lm-eval-harness/test_lm_eval_correctness.py +++ b/.buildkite/lm-eval-harness/test_lm_eval_correctness.py @@ -21,6 +21,7 @@ from lm_eval.utils import simple_parse_args_string from vllm import LLM +from vllm.config import CacheConfig from vllm.distributed import cleanup_dist_env_and_memory from vllm.logger import init_logger from vllm.model_executor.model_loader.weight_utils import drop_checkpoint_cache @@ -89,6 +90,7 @@ def _shutdown_lm(lm) -> None: return gpu_memory_utilization = ( llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization + or CacheConfig.DEFAULT_GPU_MEMORY_UTILIZATION ) try: llm.llm_engine.engine_core.shutdown(timeout=ROCM_ENGINE_SHUTDOWN_TIMEOUT_S) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index face057ee487..d966f3ed65d8 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -272,10 +272,6 @@ steps: queue: arm64_cpu_queue_release commands: - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" - # TODO(arpera): - # Remove BUILD_NIXL=true once NIXL has a wheel built against torch 2.15 - # For details: https://github.com/vllm-project/vllm/pull/60090 - # Tracking issue report: https://github.com/ai-dynamo/nixl/issues/2297 - | DOCKER_BUILDKIT=1 docker buildx build --load \ $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134) \ @@ -287,7 +283,6 @@ steps: --build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \ --build-arg INSTALL_KV_CONNECTORS=true \ --build-arg INSTALL_RUBIN_PRERELEASE=true \ - --build-arg BUILD_NIXL=true \ --build-arg TRITON_INSTALL_FROM_SOURCE_REPO=https://github.com/triton-lang/triton.git \ --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg RUN_WHEEL_CHECK=false \ @@ -300,6 +295,80 @@ steps: - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134" release-images' + + # ---- AWS EFA variants: docker/Dockerfile.efa on the images built above ---- + - label: "Build release image - x86_64 - CUDA 13.0 - EFA" + depends_on: + - build-release-image-x86 + id: build-release-image-x86-efa + agents: + queue: cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh efa) \ + --build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \ + --progress plain \ + -f docker/Dockerfile.efa docker/ + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa" + # re-tag to default image tag and push, just in case arm64 build fails + - "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa" + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa" release-images' + + - label: "Build release image - aarch64 - CUDA 13.0 - EFA" + depends_on: + - build-release-image-arm64 + id: build-release-image-arm64-efa + agents: + queue: arm64_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh efa) \ + --build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \ + --progress plain \ + -f docker/Dockerfile.efa docker/ + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-efa" release-images' + + - label: "Build release image - x86_64 - Rubin CUDA 13.4 - EFA" + depends_on: + - build-release-image-x86-rubin-cuda-13-4 + id: build-release-image-x86-rubin-cuda-13-4-efa + if: build.env("NIGHTLY") == "1" + agents: + queue: cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134-efa) \ + --build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134 \ + --progress plain \ + -f docker/Dockerfile.efa docker/ + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa" release-images' + + - label: "Build release image - aarch64 - Rubin CUDA 13.4 - EFA" + depends_on: + - build-release-image-arm64-rubin-cuda-13-4 + id: build-release-image-arm64-rubin-cuda-13-4-efa + if: build.env("NIGHTLY") == "1" + agents: + queue: arm64_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - | + DOCKER_BUILDKIT=1 docker build \ + $(bash .buildkite/scripts/docker-build-metadata-args.sh cu134-efa) \ + --build-arg BASE_IMAGE=public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134 \ + --progress plain \ + -f docker/Dockerfile.efa docker/ + - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu134-efa" release-images' - label: ":docker: Build release image - x86_64 - XPU" depends_on: ~ id: build-xpu-release-image @@ -394,6 +463,35 @@ steps: - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134" - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: Rubin CUDA 13.4" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134" release-manifests' + + + # ---- AWS EFA variant manifests ---- + - label: "Create multi-arch manifest - CUDA 13.0 - EFA" + depends_on: + - build-release-image-x86-efa + - build-release-image-arm64-efa + id: create-multi-arch-manifest-efa + agents: + queue: small_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-efa --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 - EFA" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-efa" release-manifests' + + - label: "Create multi-arch manifest - Rubin CUDA 13.4 - EFA" + depends_on: + - build-release-image-x86-rubin-cuda-13-4-efa + - build-release-image-arm64-rubin-cuda-13-4-efa + id: create-multi-arch-manifest-rubin-cuda-13-4-efa + if: build.env("NIGHTLY") == "1" + agents: + queue: small_cpu_queue_release + commands: + - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7" + - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu134-efa public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu134-efa --amend" + - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134-efa" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: Rubin CUDA 13.4 - EFA" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu134-efa" release-manifests' - label: "Create manifest - XPU" depends_on: - build-xpu-release-image @@ -494,252 +592,96 @@ steps: # ============================================================================= # ROCm Release Pipelines (x86_64 only) # ============================================================================= - # - # One group per ROCm stack, identified by its base Dockerfile (ROCM_BASE_DOCKERFILE). - # The scripts derive the ROCm version, wheel variant, cache keys and image tags from - # it (.buildkite/scripts/rocm/stack.sh), so both groups share one set of step - # definitions, anchored on the ROCm 10.0 group. - # - # - ROCm 10.0 (TheRock, docker/Dockerfile.rocm{,_base}) is the default: it owns the - # :nightly / :latest images and the root wheel index. - # - ROCm 7.2 (docker/Dockerfile.rocm_72{,_base}) is kept for a transition period as - # rocm72 images and wheels. - # - # vLLM version is determined by the Buildkite checkout (like the CUDA pipeline). + # One group per ROCm stack, on this pipeline's own agents. A serialized step builds the stack and + # runtime from source into ECR (once per content key, with only its own ECR build cache; it never + # reads rocm/vllm-ci); the build step compiles csrc for all release arches and pushes + # :-{,-wheels,-base} (.buildkite/scripts/rocm/build.sh release); the steps + # after it only copy those (.buildkite/scripts/rocm/publish.sh). + # - ROCm 10.0 (TheRock, stack `rocm`) is the default: it owns the :nightly / :latest images and + # the root wheel index. # ============================================================================= - - group: "Build ROCm 10.0 Wheel / Image" - key: "build-rocm100" + - group: ":rocm: ROCm 10.0" + key: build-rocm100 depends_on: ~ steps: - - label: ":rocm: Build ROCm 10.0 Base Image & Wheels" - id: build-rocm100-base + - label: ":rocm: Ensure ROCm 10.0 release stack + runtime" + key: build-rocm100-release-base depends_on: ~ - agents: &rocm_build_agents - queue: cpu_queue_release - timeout_in_minutes: 360 - env: &rocm100_env - ROCM_BASE_DOCKERFILE: "docker/Dockerfile.rocm_base" - ROCM_DEFAULT_ARG: "--default" - DOCKER_BUILDKIT: "1" - S3_BUCKET: "vllm-wheels" - DOCKERHUB_USERNAME: "vllmbot" - commands: &rocm_base_commands - - | - set -euo pipefail - source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" - SCCACHE_ARGS=(--build-arg USE_SCCACHE=1 --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache - --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0) - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - - # Base image, cached in ECR by the base Dockerfile's hash - if docker manifest inspect "$${ROCM_STACK_ECR_BASE}" > /dev/null 2>&1; then - echo "ECR image cache HIT: $${ROCM_STACK_ECR_BASE}" - else - docker buildx build --file "$${ROCM_BASE_DOCKERFILE}" --tag "$${ROCM_STACK_ECR_BASE}" \ - "$${SCCACHE_ARGS[@]}" --push . - fi - - # Base wheels, cached in S3 (see cache-rocm-base-wheels.sh) - CACHE=.buildkite/scripts/cache-rocm-base-wheels.sh - if [ "$$($${CACHE} "$${ROCM_BASE_DOCKERFILE}" check)" = "hit" ]; then - $${CACHE} "$${ROCM_BASE_DOCKERFILE}" download - else - DEBS="rocm-base-debs:$${BUILDKITE_BUILD_NUMBER}-$${ROCM_STACK_VARIANT}" - docker buildx build --file "$${ROCM_BASE_DOCKERFILE}" --tag "$${DEBS}" \ - --target debs_wheel_release "$${SCCACHE_ARGS[@]}" --load . - mkdir -p "$${ROCM_STACK_BASE_WHEELS_DIR}" - cid=$$(docker create "$${DEBS}") - docker cp "$${cid}":/app/debs/. "$${ROCM_STACK_BASE_WHEELS_DIR}/" - docker rm "$${cid}" - $${CACHE} "$${ROCM_BASE_DOCKERFILE}" upload - fi - buildkite-agent artifact upload "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" - - - label: ":python: Build vLLM ROCm 10.0 Wheel - x86_64" - id: build-rocm100-vllm-wheel - depends_on: build-rocm100-base - agents: *rocm_build_agents - timeout_in_minutes: 180 - env: *rocm100_env - commands: &rocm_vllm_wheel_commands - - | - set -euo pipefail - source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" - # Buildkite's default fetch doesn't update moved tags; version detection needs them - git fetch --tags --force origin - buildkite-agent artifact download "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" . - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - docker pull "$${ROCM_STACK_ECR_BASE}" - mkdir -p docker/context/base-wheels - touch docker/context/base-wheels/.keep - cp "$${ROCM_STACK_BASE_WHEELS_DIR}"/*.whl docker/context/base-wheels/ - docker build --file "$${ROCM_STACK_DOCKERFILE}" --target export_vllm_wheel_release \ - --output type=local,dest=rocm-dist \ - --build-arg BASE_IMAGE="$${ROCM_STACK_ECR_BASE}" --build-arg REMOTE_VLLM=0 \ - --build-arg GIT_REPO_CHECK=1 --build-arg USE_SCCACHE=1 \ - --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0 . - mkdir -p "$${ROCM_STACK_VLLM_WHEEL_DIR}" - cp rocm-dist/*.whl "$${ROCM_STACK_VLLM_WHEEL_DIR}/" - buildkite-agent artifact upload "$${ROCM_STACK_VLLM_WHEEL_DIR}/*.whl" - - - label: ":s3: Upload ROCm 10.0 Wheels to S3" - id: upload-rocm100-wheels - depends_on: build-rocm100-vllm-wheel - agents: *rocm_build_agents - timeout_in_minutes: 60 - env: *rocm100_env - commands: &rocm_upload_commands - - | - set -euo pipefail - source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" - buildkite-agent artifact download "$${ROCM_STACK_BASE_WHEELS_DIR}/*.whl" . - buildkite-agent artifact download "$${ROCM_STACK_VLLM_WHEEL_DIR}/*.whl" . - bash .buildkite/scripts/upload-rocm-wheels.sh "$${ROCM_BASE_DOCKERFILE}" - - - label: ":memo: Annotate ROCm 10.0 wheel release" - id: annotate-rocm100-release - depends_on: upload-rocm100-wheels - agents: *rocm_build_agents - env: *rocm100_env - commands: &rocm_annotate_commands - - bash .buildkite/scripts/annotate-rocm-release.sh "$${ROCM_BASE_DOCKERFILE}" $${ROCM_DEFAULT_ARG} - - # Generate https://wheels.vllm.ai/rocm/ (the default stack's variant), so users can - # `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/` (release only) - - block: "Generate Root Index for ROCm Wheels for Release" - key: block-generate-root-index-rocm-wheels - depends_on: - - upload-rocm100-wheels - - upload-rocm72-wheels - - - label: ":package: Generate Root Index for ROCm Wheels for Release" - depends_on: block-generate-root-index-rocm-wheels - id: generate-root-index-rocm-wheels - agents: *rocm_build_agents - env: *rocm100_env + agents: {queue: cpu_queue_release} + timeout_in_minutes: 720 # a new stack key (release never uses the CI cache) takes hours + concurrency_group: vllm/rocm/release/stack-rocm # one run builds a new key; the others then reuse it + concurrency: 1 + env: &rocm_release_env {ROCM_CI_REPO: &rocm_ecr public.ecr.aws/q9t5s3a7/vllm-release-repo, SCCACHE_BUCKET_NAME: vllm-build-sccache} commands: - - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh" - - - label: ":docker: Build release image - x86_64 - ROCm 10.0" - id: build-rocm100-release-image - depends_on: - - step: block-build-release-images - allow_failure: true - - step: build-rocm100-base - allow_failure: false - agents: *rocm_build_agents - timeout_in_minutes: 180 - env: *rocm100_env - commands: &rocm_release_image_commands - - | - set -euo pipefail - source .buildkite/scripts/rocm/stack.sh "$${ROCM_BASE_DOCKERFILE}" - aws ecr-public get-login-password --region us-east-1 | \ - docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 - docker pull "$${ROCM_STACK_ECR_BASE}" - docker build --file "$${ROCM_STACK_DOCKERFILE}" --target vllm-openai \ - --tag "$${ROCM_STACK_ECR_IMAGE}" --progress plain \ - --build-arg max_jobs=16 --build-arg BASE_IMAGE="$${ROCM_STACK_ECR_BASE}" \ - --build-arg INSTALL_LMCACHE=true --build-arg USE_SCCACHE=1 \ - --build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \ - --build-arg SCCACHE_REGION_NAME=us-west-2 --build-arg SCCACHE_S3_NO_CREDENTIALS=0 . - docker push "$${ROCM_STACK_ECR_IMAGE}" - - - label: "Publish nightly ROCm 10.0 image to DockerHub" - id: publish-rocm100-nightly - depends_on: build-rocm100-release-image - if: build.env("NIGHTLY") == "1" - agents: &rocm_publish_agents - queue: small_cpu_queue_release - env: *rocm100_env - commands: &rocm_publish_nightly_commands - - bash .buildkite/scripts/push-nightly-builds-rocm.sh "$${ROCM_BASE_DOCKERFILE}" $${ROCM_DEFAULT_ARG} - plugins: &rocm_dockerhub_login + - aws ecr-public get-login-password --region us-east-1 | docker login -u AWS --password-stdin $${ROCM_CI_REPO%/*} + - bash .buildkite/scripts/rocm/build.sh rocm release-base + retry: &rocm_release_retry {automatic: [{exit_status: 75, limit: 1}]} # registry error (build.sh) + plugins: &rocm_hub_login # authenticated Docker Hub pulls and pushes - docker-login#v3.0.0: username: vllmbot password-env: DOCKERHUB_TOKEN - - group: "Build ROCm 7.2 Wheel / Image" - key: "build-rocm72" - depends_on: ~ - steps: - - label: ":rocm: Build ROCm 7.2 Base Image & Wheels" - id: build-rocm72-base - depends_on: ~ - agents: *rocm_build_agents - timeout_in_minutes: 360 - env: &rocm72_env - ROCM_BASE_DOCKERFILE: "docker/Dockerfile.rocm_72_base" - ROCM_DEFAULT_ARG: "" - DOCKER_BUILDKIT: "1" - S3_BUCKET: "vllm-wheels" - DOCKERHUB_USERNAME: "vllmbot" - commands: *rocm_base_commands - - - label: ":python: Build vLLM ROCm 7.2 Wheel - x86_64" - id: build-rocm72-vllm-wheel - depends_on: build-rocm72-base - agents: *rocm_build_agents - timeout_in_minutes: 180 - env: *rocm72_env - commands: *rocm_vllm_wheel_commands - - - label: ":s3: Upload ROCm 7.2 Wheels to S3" - id: upload-rocm72-wheels - depends_on: build-rocm72-vllm-wheel - agents: *rocm_build_agents + - label: ":rocm: Build ROCm 10.0 release image + wheels" + key: build-rocm100-release + depends_on: build-rocm100-release-base + agents: {queue: cpu_queue_release} + timeout_in_minutes: 720 # an all-arch csrc compile without sccache hits takes ~3h + env: *rocm_release_env + commands: + - aws ecr-public get-login-password --region us-east-1 | docker login -u AWS --password-stdin $${ROCM_CI_REPO%/*} + - bash .buildkite/scripts/rocm/build.sh rocm release + retry: *rocm_release_retry + plugins: *rocm_hub_login + + - label: ":s3: Upload ROCm 10.0 wheels" + key: upload-rocm100-wheels + depends_on: build-rocm100-release + agents: {queue: cpu_queue_release} timeout_in_minutes: 60 - env: *rocm72_env - commands: *rocm_upload_commands - - - label: ":memo: Annotate ROCm 7.2 wheel release" - id: annotate-rocm72-release - depends_on: upload-rocm72-wheels - agents: *rocm_build_agents - env: *rocm72_env - commands: *rocm_annotate_commands - - - label: ":docker: Build release image - x86_64 - ROCm 7.2" - id: build-rocm72-release-image + env: {S3_BUCKET: vllm-wheels, ECR: *rocm_ecr} + command: bash .buildkite/scripts/rocm/publish.sh rocm wheels + plugins: *rocm_hub_login + + - label: ":docker: ROCm 10.0 image: ECR tags (+ nightly Docker Hub)" + key: build-rocm100-release-image depends_on: - step: block-build-release-images allow_failure: true - - step: build-rocm72-base - allow_failure: false - agents: *rocm_build_agents - timeout_in_minutes: 180 - env: *rocm72_env - commands: *rocm_release_image_commands - - - label: "Publish nightly ROCm 7.2 image to DockerHub" - id: publish-rocm72-nightly - depends_on: build-rocm72-release-image - if: build.env("NIGHTLY") == "1" - agents: *rocm_publish_agents - env: *rocm72_env - commands: *rocm_publish_nightly_commands - plugins: *rocm_dockerhub_login + - build-rocm100-release + agents: {queue: small_cpu_queue_release} + env: {ECR: *rocm_ecr} + command: bash .buildkite/scripts/rocm/publish.sh rocm images + plugins: *rocm_hub_login - - label: "Publish nightly XPU image to DockerHub" - depends_on: - - create-manifest-xpu - if: build.env("NIGHTLY") == "1" - agents: - queue: small_cpu_queue_release - commands: - - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh" - - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu" - plugins: - - docker-login#v3.0.0: - username: vllmbot - password-env: DOCKERHUB_TOKEN - env: - DOCKER_BUILDKIT: "1" - DOCKERHUB_USERNAME: "vllmbot" + # Generate https://wheels.vllm.ai/rocm/ (the default stack's variant). + - block: "Generate Root Index for ROCm Wheels for Release" + key: block-generate-root-index-rocm-wheels + depends_on: upload-rocm100-wheels + + - label: ":package: Generate Root Index for ROCm Wheels for Release" + key: generate-root-index-rocm-wheels + depends_on: block-generate-root-index-rocm-wheels + agents: {queue: cpu_queue_release} + env: {S3_BUCKET: vllm-wheels} + command: VARIANT=$$(buildkite-agent meta-data get rocm-rocm-variant) bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh + + - label: "Publish nightly XPU image to DockerHub" + depends_on: + - create-manifest-xpu + if: build.env("NIGHTLY") == "1" + agents: + queue: small_cpu_queue_release + commands: + - "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh" + - "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu" + plugins: + - docker-login#v3.0.0: + username: vllmbot + password-env: DOCKERHUB_TOKEN + env: + DOCKER_BUILDKIT: "1" + DOCKERHUB_USERNAME: "vllmbot" # ============================================================================= # Publish to DockerHub and PyPI (at the end so all builds complete first) @@ -751,7 +693,6 @@ steps: - create-multi-arch-manifest - create-manifest-xpu - build-rocm100-release-image - - build-rocm72-release-image - input-release-version # Wait for CPU builds if their block steps were unblocked, so publish # doesn't race the in-progress CPU build. allow_failure lets publish @@ -761,6 +702,10 @@ steps: allow_failure: true - step: build-cpu-release-image-arm64 allow_failure: true + # Wait for the EFA images, but don't let an EFA failure block publishing + # the others; the cuda-13-0-efa publish job fails on a partial build. + - step: create-multi-arch-manifest-efa + allow_failure: true - label: "Publish {{matrix}} release images to DockerHub" depends_on: @@ -772,8 +717,9 @@ steps: - "bash .buildkite/scripts/publish-release-images.sh {{matrix}}" matrix: - "cuda-13-0" + # AWS EFA variant; skipped if no EFA image was built, fails on a partial build. + - "cuda-13-0-efa" - "rocm" - - "rocm72" - "xpu" - "cpu" plugins: diff --git a/.buildkite/scripts/annotate-rocm-release.sh b/.buildkite/scripts/annotate-rocm-release.sh deleted file mode 100755 index 1889db43b5ad..000000000000 --- a/.buildkite/scripts/annotate-rocm-release.sh +++ /dev/null @@ -1,87 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# -# Generate Buildkite annotation for a ROCm stack's wheel release. -# -# Usage: annotate-rocm-release.sh [--default] -# base-dockerfile e.g. docker/Dockerfile.rocm_base or docker/Dockerfile.rocm_72_base -# --default the stack also owns the plain :latest / :v Docker Hub tags -set -ex - -# shellcheck source=.buildkite/scripts/rocm/stack.sh -source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 [--default]}" -TAG_VARIANTS=("-${ROCM_STACK_VARIANT}") -if [[ "${2:-}" == "--default" ]]; then - TAG_VARIANTS=("" "-${ROCM_STACK_VARIANT}") -fi -ROCM_VERSION="$ROCM_STACK_VERSION" -VARIANT="$ROCM_STACK_VARIANT" -PYTHON_VERSION=$(sed -nE 's/^ARG PYTHON_VERSION="?([^" ]+)"?.*/\1/p' "$ROCM_STACK_BASE_DOCKERFILE") -PYTORCH_ROCM_ARCH=$(sed -nE 's/^ARG PYTORCH_ROCM_ARCH="?([^"]+)"?.*/\1/p' "$ROCM_STACK_BASE_DOCKERFILE") - -# Get release version, default to 1.0.0.dev for nightly/per-commit builds -RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "") -if [ -z "${RELEASE_VERSION}" ]; then - RELEASE_VERSION="1.0.0.dev" -fi - -# S3 URLs -S3_BUCKET="${S3_BUCKET:-vllm-wheels}" -S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}" -S3_HOST="${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com" -S3_URL="http://${S3_HOST}" -WHEEL_PATH="rocm/${BUILDKITE_COMMIT}/${VARIANT}-wheels/" - -DOCKER_CMDS="docker pull ${ROCM_STACK_ECR_BASE} -docker pull ${ROCM_STACK_ECR_IMAGE}" -for suffix in "${TAG_VARIANTS[@]}"; do - for tag in "latest${suffix}" "v${RELEASE_VERSION}${suffix}"; do - DOCKER_CMDS+=" -docker tag ${ROCM_STACK_ECR_BASE} vllm/vllm-openai-rocm:${tag}-base -docker tag ${ROCM_STACK_ECR_IMAGE} vllm/vllm-openai-rocm:${tag} -docker push vllm/vllm-openai-rocm:${tag}-base -docker push vllm/vllm-openai-rocm:${tag}" - done -done - -buildkite-agent annotate --style 'success' --context "rocm-release-workflow-${VARIANT}" << EOF -## ROCm ${ROCM_VERSION} Wheel and Docker Image Releases -### Build Configuration -| Setting | Value | -|---------|-------| -| **ROCm Version** | ${ROCM_VERSION} | -| **Index Variant** | ${VARIANT} | -| **Python Version** | ${PYTHON_VERSION} | -| **GPU Architectures** | ${PYTORCH_ROCM_ARCH} | -| **Branch** | \`${BUILDKITE_BRANCH}\` | -| **Commit** | \`${BUILDKITE_COMMIT}\` | - -### :package: Installation - -\`\`\`bash -# This build (by commit) -pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${VARIANT}/ --trusted-host ${S3_HOST} -# Nightly (if published) -pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/${VARIANT}/ --trusted-host ${S3_HOST} -\`\`\` - -### :floppy_disk: Download Wheels Directly - -\`\`\`bash -aws s3 ls s3://${S3_BUCKET}/${WHEEL_PATH} -aws s3 cp --recursive --exclude "*" --include "*.whl" s3://${S3_BUCKET}/${WHEEL_PATH} . -\`\`\` - -### :warning: Notes -- These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs -- Supported GPU architectures: ${PYTORCH_ROCM_ARCH} -- Platform: Linux x86_64 only - -### :package: Docker Image Release - -\`\`\` -${DOCKER_CMDS} -\`\`\` - -EOF diff --git a/.buildkite/scripts/cache-rocm-base-wheels.sh b/.buildkite/scripts/cache-rocm-base-wheels.sh deleted file mode 100755 index 02fbbf52cf54..000000000000 --- a/.buildkite/scripts/cache-rocm-base-wheels.sh +++ /dev/null @@ -1,125 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# -# Cache helper for ROCm base wheels -# -# This script manages caching of pre-built ROCm base wheels (torch, triton, etc.) -# to avoid rebuilding them when the stack's base Dockerfile hasn't changed. -# -# Usage: cache-rocm-base-wheels.sh -# base-dockerfile e.g. docker/Dockerfile.rocm_base or docker/Dockerfile.rocm_72_base -# Commands: -# check - Check if cache exists, outputs "hit" or "miss" -# upload - Upload wheels to cache -# download - Download wheels from cache -# key - Output the cache key -# path - Output the full S3 cache path -# -# Environment variables: -# S3_BUCKET - S3 bucket name (default: vllm-wheels) -# -# The cache key hashes the base Dockerfile and the repo files it COPYs in, which -# capture the ROCm version and every component pin (see rocm/stack.sh). - -set -euo pipefail - -BUCKET="${S3_BUCKET:-vllm-wheels}" -CACHE_PREFIX="rocm/cache" - -# shellcheck source=.buildkite/scripts/rocm/stack.sh -source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 }" -shift -CACHE_KEY="$ROCM_STACK_WHEEL_KEY" -WHEELS_DIR="$ROCM_STACK_BASE_WHEELS_DIR" -CACHE_PATH="s3://${BUCKET}/${CACHE_PREFIX}/${CACHE_KEY}/" - -case "${1:-}" in - check) - echo "Checking cache for key: ${CACHE_KEY}" >&2 - echo "Cache path: ${CACHE_PATH}" >&2 - - # Check if cache exists by listing objects - # We look for at least one .whl file - echo "Running: aws s3 ls ${CACHE_PATH}" >&2 - S3_OUTPUT=$(aws s3 ls "${CACHE_PATH}" 2>&1) || true - echo "S3 ls output:" >&2 - echo "$S3_OUTPUT" | head -5 >&2 - - if echo "$S3_OUTPUT" | grep -q "\.whl"; then - echo "hit" - else - echo "miss" - fi - ;; - - upload) - echo "========================================" - echo "Uploading wheels to cache" - echo "========================================" - echo "Cache key: ${CACHE_KEY}" - echo "Cache path: ${CACHE_PATH}" - echo "" - - if [[ ! -d "${WHEELS_DIR}" ]]; then - echo "ERROR: ${WHEELS_DIR} directory not found" >&2 - exit 1 - fi - - WHEEL_COUNT=$(find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) - if [[ "$WHEEL_COUNT" -eq 0 ]]; then - echo "ERROR: No wheels found in ${WHEELS_DIR}/" >&2 - exit 1 - fi - - echo "Uploading $WHEEL_COUNT wheels..." - aws s3 cp --recursive "${WHEELS_DIR}/" "${CACHE_PATH}" - - echo "" - echo "Cache upload complete!" - echo "========================================" - ;; - - download) - echo "========================================" - echo "Downloading wheels from cache" - echo "========================================" - echo "Cache key: ${CACHE_KEY}" - echo "Cache path: ${CACHE_PATH}" - echo "" - mkdir -p "${WHEELS_DIR}" - - # Use sync with include/exclude to only download .whl files - aws s3 sync "${CACHE_PATH}" "${WHEELS_DIR}/" \ - --exclude "*" \ - --include "*.whl" - - echo "" - echo "Downloaded wheels:" - find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' -exec ls -lh {} \; - WHEEL_COUNT=$(find "${WHEELS_DIR}" -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) - echo "" - echo "Total: $WHEEL_COUNT wheels" - echo "========================================" - ;; - - key) - echo "${CACHE_KEY}" - ;; - - path) - echo "${CACHE_PATH}" - ;; - - *) - echo "Usage: $0 {check|upload|download|key|path}" >&2 - echo "" >&2 - echo "Commands:" >&2 - echo " check - Check if cache exists, outputs 'hit' or 'miss'" >&2 - echo " upload - Upload wheels from ${WHEELS_DIR}/ to cache" >&2 - echo " download - Download wheels from cache to ${WHEELS_DIR}/" >&2 - echo " key - Output the cache key" >&2 - echo " path - Output the full S3 cache path" >&2 - exit 1 - ;; -esac diff --git a/.buildkite/scripts/check-ray-compatibility.sh b/.buildkite/scripts/check-ray-compatibility.sh index ec42834802d6..8e1eef6948fe 100644 --- a/.buildkite/scripts/check-ray-compatibility.sh +++ b/.buildkite/scripts/check-ray-compatibility.sh @@ -27,6 +27,8 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then echo ">>> WARNING: ROCm ${ROCM_VER} wheel index not found at ${CANDIDATE_URL}" echo ">>> Falling back to default PyPI (resolution may be incomplete)" TORCH_INDEX_URL="" + # The image's own stack builds (TheRock torch, triton, amdsmi) are on no public index. + LOCAL_BUILDS=1 fi else if [ "${TORCH_NIGHTLY:-0}" = "1" ]; then @@ -110,6 +112,14 @@ with open(out_path, "w") as f: print(f"Wrote {len(lines)} constraints to {out_path}") PYEOF +if [[ -n "${LOCAL_BUILDS:-}" ]]; then + # Compare them at their public release. Dropped: amdsmi has none (TheRock ships it as + # sources) and triton is the stack's own, which its torch is repacked to require. + echo ">>> Pinning local builds to their public release:" + grep -E '==[^ ;]*\+|\[device-' "${WORK_DIR}/vllm-constraints.txt" || true + sed -i -E -e '/^(amdsmi|triton)([=<>!~ ;[]|$)|\[device-/d' -e 's/(==[^ ;+]+)\+[^ ;]*/\1/' "${WORK_DIR}/vllm-constraints.txt" +fi + echo ">>> Installed vllm deps (first 20 lines):" head -20 "${WORK_DIR}/vllm-constraints.txt" diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh deleted file mode 100644 index 5935938616bc..000000000000 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ /dev/null @@ -1,2737 +0,0 @@ -#!/bin/bash -# ci-bake-rocm.sh - Docker buildx bake wrapper for ROCm CI builds. -# -# The wrapper keeps three build concerns separate: -# * ci_base builds are content-addressed by vllm.ci_base.content_hash. -# * test image contents record the commit while runtime tags stay build-scoped. -# * ROCm install artifacts are uploaded only for targets that export wheels. -# -# Usage: -# ci-bake-rocm.sh [TARGET] -# -# Set BAKE_PRINT_ONLY=1 to stop after docker buildx bake --print. - -set -euo pipefail - -DEFAULT_REPO_SLUG="vllm-project/vllm" -DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl" -DEFAULT_CI_BASE_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt tools/install_torchcodec_rocm.sh rust-toolchain.toml tests/vllm_test_utils" -DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm" -DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust-toolchain build_nixl lmcache_source build_lmcache build_deepep mori_base ci_base" -DEFAULT_CI_BASE_METADATA_VERSION="3" -# ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected -# local-source stages rather than unreachable remote-fetch alternatives. -DEFAULT_ROCM_CSRC_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt pyproject.toml setup.py CMakeLists.txt cmake csrc vllm/envs.py vllm/__init__.py tools/build_rust.py" -DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies rocm-triton-kernels csrc-build" -DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/build_rust.sh" -DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm vllm-version rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build" -# Docker's 128-character tag limit minus the longest cache prefix -# ("csrc-rocm-branch-" and "rust-rocm-branch-", both 17 characters). -ROCM_CACHE_BRANCH_TAG_MAX_LEN=111 -CI_BASE_WRITE_SCOPE="" - -TARGET="" -CI_HCL_SOURCE="${CI_HCL_SOURCE:-}" -CI_HCL_PATH="" -CI_BASE_LABEL_OVERRIDE_PATH="" -CSRC_CACHE_OVERRIDE_PATH="" -ROCM_ARG_OVERRIDE_PATH="" -BUILD_CONTEXT_OVERRIDE_PATH="" -SCRIPT_TMP_DIR="" -BAKE_CONFIG_FILE="" -ROCM_BUILD_CONTEXT_ROOT="" -ROCM_BUILD_CONTEXT_INDEX="" -ROCM_BUILD_CONTEXT_COMMIT="" -BAKE_FILES=() -BAKE_ALLOW_ARGS=() -BAKE_TARGETS=() -DEPENDENCY_CACHE_TARGETS=() - -cleanup() { - if [[ -n "${SCRIPT_TMP_DIR}" && -d "${SCRIPT_TMP_DIR}" ]]; then - rm -rf "${SCRIPT_TMP_DIR}" - fi -} -trap cleanup EXIT - -clean_docker_tag() { - local input="$1" - echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128 -} - -is_url_like() { - local value="${1:-}" - [[ "${value}" =~ ^[a-zA-Z][a-zA-Z0-9+.-]*:// || "${value}" == git@*:* ]] -} - -is_full_git_sha() { - local value="${1:-}" - [[ "${value}" =~ ^[0-9a-fA-F]{40}$ ]] -} - -select_cache_branch_name() { - local candidate="" - local var="" - - for var in \ - ROCM_CACHE_BRANCH_NAME \ - BUILDKITE_PULL_REQUEST_HEAD_BRANCH \ - BUILDKITE_HEAD_BRANCH \ - BUILDKITE_BRANCH \ - VLLM_BRANCH; do - candidate="${!var:-}" - [[ -n "${candidate}" ]] || continue - is_url_like "${candidate}" && continue - is_full_git_sha "${candidate}" && continue - printf '%s\n' "${candidate}" - return 0 - done -} - -cache_scope_suffix() { - local arch_hash="" - arch_hash=$(printf '%s' "${PYTORCH_ROCM_ARCH:-default}" | sha256sum | cut -c1-12) - printf 'arch-%s\n' "${arch_hash}" -} - -compose_cache_branch_tag() { - local repo_slug="$1" - local branch="$2" - local suffix="" - local prefix="" - local max_prefix_len=0 - local max_tag_len="${ROCM_CACHE_BRANCH_TAG_MAX_LEN}" - - suffix="$(cache_scope_suffix)" - prefix="$(clean_docker_tag "${repo_slug}")-$(clean_docker_tag "${branch}")" - max_prefix_len=$((max_tag_len - ${#suffix} - 1)) - if (( max_prefix_len < 1 )); then - max_prefix_len=1 - fi - printf '%s-%s\n' "${prefix:0:${max_prefix_len}}" "${suffix}" -} - -parse_repo_slug() { - local repo_url="${1:-}" - local repo_slug="" - - if [[ -z "${repo_url}" ]]; then - printf '%s\n' "${DEFAULT_REPO_SLUG}" - return 0 - fi - - repo_slug=$(echo "${repo_url}" | sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#') - if [[ "${repo_slug}" != */* ]]; then - repo_slug="${DEFAULT_REPO_SLUG}" - fi - printf '%s\n' "${repo_slug}" -} - -normalize_repo_slug() { - local repo_slug="${1:-}" - - repo_slug="${repo_slug%/}" - repo_slug="${repo_slug%.git}" - repo_slug="${repo_slug#https://github.com/}" - repo_slug="${repo_slug#http://github.com/}" - repo_slug="${repo_slug#ssh://git@github.com/}" - repo_slug="${repo_slug#git@github.com:}" - repo_slug="${repo_slug#github.com/}" - printf '%s\n' "${repo_slug}" -} - -is_trusted_ci_cache_writer() { - local actual_repo="" - local trusted_repo="" - - [[ "${BUILDKITE:-false}" == "true" ]] || return 1 - [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]] || return 1 - [[ "${BUILDKITE_BRANCH:-}" == "${CI_BASE_STABLE_BRANCH:-main}" ]] || return 1 - actual_repo=$(normalize_repo_slug "${BUILDKITE_REPO:-}") - trusted_repo=$(normalize_repo_slug \ - "${CI_BASE_STABLE_REPO_SLUG:-${DEFAULT_REPO_SLUG}}") - [[ -n "${actual_repo}" && "${actual_repo}" == "${trusted_repo}" ]] -} - -ci_base_write_scope() { - local identity="" - local source_repo="${BUILDKITE_PULL_REQUEST_REPO:-${BUILDKITE_REPO:-local}}" - - if is_trusted_ci_cache_writer; then - return 0 - fi - identity=$(printf '%s\n' "${source_repo}" | sha256sum | cut -c1-12) - printf 'preview-%s\n' "${identity}" -} - -configure_ci_base_write_scope() { - local scope="" - - scope=$(ci_base_write_scope) - if [[ -n "${scope}" ]]; then - CI_BASE_WRITE_SCOPE=$(clean_docker_tag "${scope}") - echo "Non-canonical cache writes use source scope: ${CI_BASE_WRITE_SCOPE}" - else - CI_BASE_WRITE_SCOPE="" - echo "Trusted main build: publishing canonical ci_base refs" - fi - export CI_BASE_WRITE_SCOPE -} - -get_buildkite_repo_slug() { - parse_repo_slug "${BUILDKITE_PULL_REQUEST_REPO:-${BUILDKITE_REPO:-}}" -} - -get_buildkite_target_repo_slug() { - parse_repo_slug "${BUILDKITE_REPO:-}" -} - -get_buildkite_target_repo_url() { - local repo_url="${BUILDKITE_REPO:-}" - - if [[ -n "${repo_url}" ]] && is_url_like "${repo_url}"; then - printf '%s\n' "${repo_url}" - return 0 - fi - - printf 'https://github.com/%s.git\n' "${DEFAULT_REPO_SLUG}" -} - -git_fetch_with_timeout() { - local timeout_secs="${ROCM_CACHE_GIT_FETCH_TIMEOUT:-60}" - local -a fetch_command=(git fetch --no-auto-maintenance) - - # Detached maintenance can race a later shallow fetch on .git/shallow. - if command -v timeout >/dev/null 2>&1; then - timeout "${timeout_secs}s" "${fetch_command[@]}" "$@" - else - "${fetch_command[@]}" "$@" - fi -} - -git_fetch_for_cache() { - git_fetch_with_timeout "$@" 2>/dev/null -} - -hash_string_short() { - printf '%s' "$1" | sha256sum | cut -c1-16 -} - -list_content_files() { - # Hash the checkout inputs Docker can intentionally consume, not ignored - # compiler/test debris left behind on a reused worker. - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" ]]; then - GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \ - git ls-files -z --cached -- "$1" | LC_ALL=C sort -z - else - git ls-files -z --cached --others --exclude-standard -- "$1" \ - | LC_ALL=C sort -z - fi -} - -content_regular_file() { - local file="$1" - local physical_file="" - - physical_file="${file}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${file}" != /* ]]; then - physical_file="${ROCM_BUILD_CONTEXT_ROOT}/${file}" - fi - [[ -f "${physical_file}" ]] -} - -hash_content_file() { - local file="$1" - local physical_file="" - local checksum="" - local file_mode="" - - physical_file="${file}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${file}" != /* ]]; then - physical_file="${ROCM_BUILD_CONTEXT_ROOT}/${file}" - fi - if [[ -L "${physical_file}" ]]; then - printf 'symlink:%s\ntarget:' "${file}" - readlink -n -- "${physical_file}" || return $? - printf '\n' - return - fi - if [[ ! -f "${physical_file}" ]]; then - printf 'missing:%s\n' "${file}" - return - fi - file_mode=$(stat -c '%a' "${physical_file}") || return $? - printf 'file:%s\nmode:%s\n' "${file}" "${file_mode}" - checksum=$(sha256sum < "${physical_file}") || return $? - checksum="${checksum%% *}" - printf '%s %s\n' "${checksum}" "${file}" -} - -hash_content_directory() { - if ! list_content_files "$1" | while IFS= read -r -d '' file; do - hash_content_file "${file}" || exit $? - done; then - echo "Failed to hash content under $1" >&2 - return 1 - fi -} - -compute_content_hash() { - local path="" - local physical_path="" - - for path in "$@"; do - physical_path="${path}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${path}" != /* ]]; then - physical_path="${ROCM_BUILD_CONTEXT_ROOT}/${path}" - fi - if [[ -L "${physical_path}" || -f "${physical_path}" ]]; then - hash_content_file "${path}" || return $? - elif [[ -d "${physical_path}" ]]; then - hash_content_directory "${path}" || return $? - else - printf 'missing:%s\n' "${path}" - fi - done | sha256sum | cut -d' ' -f1 -} - -validate_ci_build_context_source() { - local source_root="$1" - local context_commit="" - local source_head="" - local entry="" - local metadata="" - local mode="" - local stage="" - local path="" - local worktree_diff_status=0 - local staged_diff_status=0 - local index_modes_file="${SCRIPT_TMP_DIR}/git-index-modes" - local info_attributes="" - - if ! context_commit=$(git -C "${source_root}" rev-parse \ - --verify "${BUILDKITE_COMMIT:-HEAD}^{commit}"); then - echo "Failed to resolve the CI Docker context revision" >&2 - return 1 - fi - if ! source_head=$(git -C "${source_root}" rev-parse --verify HEAD); then - echo "Failed to resolve the checked-out CI revision" >&2 - return 1 - fi - if [[ "${context_commit}" != "${source_head}" ]]; then - echo "BUILDKITE_COMMIT does not match the checked-out CI revision" >&2 - return 1 - fi - ROCM_BUILD_CONTEXT_COMMIT="${context_commit}" - # Shared AMD workspaces may present tracked files with inflated executable - # bits. Context modes come from the pinned Git tree, so ignore only that - # filesystem drift while continuing to reject content changes. - git -C "${source_root}" -c core.fileMode=false diff \ - --quiet --no-ext-diff --ignore-submodules=none -- \ - || worktree_diff_status=$? - if (( worktree_diff_status != 0 )); then - if (( worktree_diff_status > 1 )); then - printf 'Failed to inspect tracked CI worktree changes (git diff exited %s)\n' \ - "${worktree_diff_status}" >&2 - return 1 - fi - echo "Tracked worktree changes cannot be omitted from the CI Docker context" \ - >&2 - echo "Tracked worktree diff (first 50 entries):" >&2 - git -C "${source_root}" -c core.fileMode=false diff \ - --no-ext-diff --ignore-submodules=none --name-status -- \ - | sed -n '1,50p' >&2 || true - git -C "${source_root}" -c core.fileMode=false diff \ - --no-ext-diff --ignore-submodules=none --summary -- \ - | sed -n '1,50p' >&2 || true - return 1 - fi - git -C "${source_root}" diff --cached \ - --quiet --no-ext-diff --ignore-submodules=none HEAD -- \ - || staged_diff_status=$? - if (( staged_diff_status != 0 )); then - if (( staged_diff_status > 1 )); then - printf 'Failed to inspect staged CI changes (git diff exited %s)\n' \ - "${staged_diff_status}" >&2 - return 1 - fi - echo "Staged changes cannot be omitted from the CI Docker context" >&2 - echo "Staged diff (first 50 entries):" >&2 - git -C "${source_root}" diff --cached \ - --no-ext-diff --ignore-submodules=none --name-status HEAD -- \ - | sed -n '1,50p' >&2 || true - git -C "${source_root}" diff --cached \ - --no-ext-diff --ignore-submodules=none --summary HEAD -- \ - | sed -n '1,50p' >&2 || true - return 1 - fi - # Untracked and ignored worker outputs are intentionally absent: the - # pinned Git tree, rather than mutable checkout contents, is the contract. - if ! info_attributes=$(git -C "${source_root}" \ - rev-parse --path-format=absolute --git-path info/attributes); then - echo "Failed to locate repository-local Git attributes" >&2 - return 1 - fi - if [[ -s "${info_attributes}" ]]; then - echo "Repository-local Git attributes cannot define the CI Docker context" \ - >&2 - return 1 - fi - - ROCM_BUILD_CONTEXT_INDEX="${SCRIPT_TMP_DIR}/docker-context.index" - if ! GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \ - git -C "${source_root}" -c core.splitIndex=false \ - read-tree "${context_commit}^{tree}"; then - echo "Failed to create the CI Docker context index" >&2 - return 1 - fi - if ! GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \ - git -C "${source_root}" ls-files --stage -z \ - > "${index_modes_file}"; then - echo "Failed to read the Git index for the CI Docker context" >&2 - return 1 - fi - while IFS= read -r -d '' entry; do - metadata="${entry%%$'\t'*}" - mode="${metadata%% *}" - stage="${metadata##* }" - path="${entry#*$'\t'}" - if [[ "${stage}" != "0" ]]; then - echo "Unmerged Git entry cannot be used as Docker context: ${path}" >&2 - return 1 - fi - case "${mode}" in - 100644|100755|120000) - ;; - 160000) - echo "Git submodule cannot be materialized in the CI Docker context: ${path}" \ - >&2 - return 1 - ;; - *) - echo "Unsupported Git mode ${mode} for ${path}" >&2 - return 1 - ;; - esac - done < "${index_modes_file}" -} - -describe_ci_revision() { - git -C "$1" describe --tags --long --abbrev=10 \ - --match 'v[0-9]*' "$2" 2>/dev/null -} - -write_ci_git_archival_metadata() { - local source_root="$1" - local context_root="$2" - local commit="${ROCM_BUILD_CONTEXT_COMMIT}" - local commit_date="" - local describe="" - local is_shallow="false" - - commit_date=$(git -C "${source_root}" show -s --format=%cI "${commit}") \ - || return $? - is_shallow=$(git -C "${source_root}" rev-parse --is-shallow-repository) \ - || return $? - if git -C "${source_root}" remote get-url origin >/dev/null 2>&1; then - # AMD agents use shallow, no-tag clones. Reach a release tag so one - # commit cannot acquire different versions from different depths. - # Versioning only needs tags and commits, not historical source trees. - echo "Synchronizing version tags for the canonical CI Docker context" - (cd "${source_root}" && git_fetch_with_timeout --quiet \ - --filter=tree:0 --prune origin '+refs/tags/*:refs/tags/*') \ - || return $? - if [[ "${is_shallow}" == "true" ]] \ - && ! describe_ci_revision "${source_root}" "${commit}" >/dev/null; then - echo "Deepening history to reach a version tag" - (cd "${source_root}" && git_fetch_with_timeout --quiet \ - --filter=tree:0 --no-tags --deepen=1000 origin "${commit}") \ - || return $? - is_shallow=$(git -C "${source_root}" \ - rev-parse --is-shallow-repository) || return $? - if [[ "${is_shallow}" == "true" ]] \ - && ! describe_ci_revision "${source_root}" "${commit}" >/dev/null; then - echo "No version tag within 1,000 commits; fetching full history" - (cd "${source_root}" && git_fetch_with_timeout --quiet \ - --filter=tree:0 --no-tags --unshallow origin "${commit}") \ - || return $? - fi - fi - fi - if ! describe=$(describe_ci_revision "${source_root}" "${commit}"); then - echo "No numeric version tag is reachable from the CI revision" >&2 - return 1 - fi - if [[ -e "${context_root}/.git_archival.txt" \ - || -L "${context_root}/.git_archival.txt" ]]; then - echo "The source tree already contains .git_archival.txt" >&2 - return 1 - fi - printf 'node: %s\nnode-date: %s\ndescribe-name: %s\n' \ - "${commit}" "${commit_date}" "${describe}" \ - > "${context_root}/.git_archival.txt" \ - && chmod 0644 -- "${context_root}/.git_archival.txt" -} - -write_build_context_override() { - local escaped_context="" - - [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" ]] || return 0 - escaped_context=$(hcl_escape_string "${ROCM_BUILD_CONTEXT_ROOT}") || return $? - if ! { - printf 'target "_common-rocm" {\n' - printf ' context = "%s"\n' "${escaped_context}" - printf '}\n' - } > "${BUILD_CONTEXT_OVERRIDE_PATH}"; then - echo "Failed to write the CI Docker context override" >&2 - return 1 - fi - BAKE_FILES+=(-f "${BUILD_CONTEXT_OVERRIDE_PATH}") -} - -prepare_ci_build_context() { - local source_root="" - local context_root="" - - [[ "${BUILDKITE:-false}" == "true" ]] || return 0 - [[ "${REMOTE_VLLM:-0}" == "0" ]] || return 0 - - # BuildKit includes file modes in cache keys. Export the pinned revision to - # an owned context instead of changing modes in the shared checkout. - if ! source_root=$(git rev-parse --show-toplevel); then - echo "Failed to locate the CI source checkout" >&2 - return 1 - fi - validate_ci_build_context_source "${source_root}" || return $? - context_root="${SCRIPT_TMP_DIR}/docker-context" - if ! mkdir -m 0700 -- "${context_root}"; then - echo "Failed to create the owned CI Docker context" >&2 - return 1 - fi - if ! ( - umask 0022 - unset GIT_LFS_SKIP_SMUDGE - export GIT_ATTR_NOSYSTEM=1 - export GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" - git -C "${source_root}" \ - -c core.attributesFile=/dev/null \ - -c core.autocrlf=false \ - -c core.eol=lf \ - -c core.symlinks=true \ - --work-tree="${context_root}" \ - checkout-index --all - ); then - echo "Failed to materialize the CI Docker context from the Git index" >&2 - return 1 - fi - - # setuptools-scm understands Git's stable archive format, so wheels and - # Rust artifacts retain their exact version without copying Git history. - if ! is_ci_base_target; then - write_ci_git_archival_metadata "${source_root}" "${context_root}" \ - || return $? - fi - if [[ -e "${context_root}/.git" ]]; then - echo "Canonical CI Docker context unexpectedly contains .git" >&2 - return 1 - fi - ROCM_BUILD_CONTEXT_ROOT="${context_root}" - BAKE_ALLOW_ARGS+=(--allow "fs.read=${ROCM_BUILD_CONTEXT_ROOT}") - echo "Using canonical CI Docker context: ${ROCM_BUILD_CONTEXT_ROOT}" -} - -configure_custom_rocm_stages() { - using_custom_rocm_dockerfiles || return 0 - local context_root="${ROCM_BUILD_CONTEXT_ROOT:-.}" - local dockerfile="${context_root}/${CI_BASE_DOCKERFILE}" - local stages="" - - (cd "${context_root}" && \ - validate_rocm_dockerfile "${CI_BASE_DOCKERFILE}" \ - ci_base test export_vllm export_test_smoke csrc-build rust-build) \ - || return $? - # Hash all stages for custom layouts, including additional helper stages. - stages=$(rocm_dockerfile_stages "${dockerfile}" | paste -sd ' ') - CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${stages}}" - ROCM_CSRC_DOCKERFILE_STAGES="${ROCM_CSRC_DOCKERFILE_STAGES:-${stages}}" - ROCM_RUST_DOCKERFILE_STAGES="${ROCM_RUST_DOCKERFILE_STAGES:-${stages}}" -} - -compose_dependency_cache_key() { - local prefix="$1" - local material="$2" - local cleaned_prefix="" - - cleaned_prefix=$(clean_docker_tag "${prefix}" | cut -c1-96) - printf '%s-%s\n' "${cleaned_prefix}" "$(hash_string_short "${material}")" -} - -hash_dockerfile_stages() { - local dockerfile="$1" - local stages="$2" - local physical_dockerfile="" - - physical_dockerfile="${dockerfile}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then - physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" - fi - - awk -v wanted_stages="${stages}" ' - BEGIN { - split(wanted_stages, stage_list, /[[:space:]]+/) - for (idx in stage_list) { - if (stage_list[idx] != "") { - wanted[stage_list[idx]] = 1 - } - } - emit = 0 - } - toupper($1) == "FROM" { - stage = "" - for (idx = 1; idx <= NF; idx++) { - if (tolower($idx) == "as" && idx < NF) { - stage = tolower($(idx + 1)) - } - } - emit = (stage in wanted) - } - emit { - print - } - ' "${physical_dockerfile}" -} - -discover_dockerfile_stage_args() { - local dockerfile="$1" - local stages="$2" - local physical_dockerfile="" - - physical_dockerfile="${dockerfile}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then - physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" - fi - [[ -f "${physical_dockerfile}" ]] || return 0 - - awk -v wanted_stages="${stages}" ' - function add_arg(name) { - if (name != "" && !(name in seen)) { - seen[name] = 1 - args[++arg_count] = name - } - } - BEGIN { - split(wanted_stages, stage_list, /[[:space:]]+/) - for (idx in stage_list) { - if (stage_list[idx] != "") { - wanted[stage_list[idx]] = 1 - } - } - emit = 1 - } - { - line = $0 - if (toupper($1) == "FROM") { - stage = "" - for (idx = 1; idx <= NF; idx++) { - if (tolower($idx) == "as" && idx < NF) { - stage = tolower($(idx + 1)) - } - } - emit = (stage in wanted) - } - if (emit) { - lines[++line_count] = line - } - } - END { - for (idx = 1; idx <= line_count; idx++) { - line = lines[idx] - arg_name = line - sub(/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+/, "", arg_name) - if (arg_name != line) { - sub(/[=[:space:]].*/, "", arg_name) - if (arg_name ~ /^[A-Za-z_][A-Za-z0-9_]*$/) { - add_arg(arg_name) - } - } - } - - for (idx = 1; idx <= line_count; idx++) { - line = lines[idx] - for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) { - name = args[arg_idx] - if (line ~ "\\$\\{" name "([}:][^}]*)?\\}" \ - || line ~ "\\$" name "([^A-Za-z0-9_]|$)") { - used[name] = 1 - } - } - } - - for (arg_idx = 1; arg_idx <= arg_count; arg_idx++) { - name = args[arg_idx] - if (used[name]) { - print name - } - } - } - ' "${physical_dockerfile}" -} - -get_content_arg_names() { - local dockerfile="$1" - local stages="$2" - local explicit_args="${3:-}" - - if [[ -n "${explicit_args}" ]]; then - tr ' ' '\n' <<< "${explicit_args}" - else - discover_dockerfile_stage_args "${dockerfile}" "${stages}" - fi | awk 'NF && !seen[$0]++' -} - -compute_ci_base_content_hash() { - local -a content_paths=() - local -a content_args=() - local content_files_hash="" - local dockerfile="${CI_BASE_DOCKERFILE:-}" - local stages="${CI_BASE_DOCKERFILE_STAGES:-}" - - read -r -a content_paths <<< "${CI_BASE_CONTENT_FILES}" - mapfile -t content_args < <( - get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}" - ) - if ! content_files_hash=$(compute_content_hash "${content_paths[@]}"); then - echo "Failed to hash ci_base content files" >&2 - return 1 - fi - - { - printf 'content-files-hash:%s\n' "${content_files_hash}" - if [[ -n "${dockerfile}" ]]; then - printf 'dockerfile:%s\n' "${dockerfile}" - printf 'resolved-build-args:\n' - hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" \ - || return 1 - if [[ -n "${stages}" ]]; then - printf 'dockerfile-stages:%s\n' "${stages}" - if content_regular_file "${dockerfile}"; then - hash_dockerfile_stages "${dockerfile}" "${stages}" - else - printf 'missing:%s\n' "${dockerfile}" - fi - fi - fi - } | sha256sum | cut -d' ' -f1 -} - -extract_dockerfile_arg_default() { - local dockerfile="$1" - local arg_name="$2" - local physical_dockerfile="" - - physical_dockerfile="${dockerfile}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then - physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" - fi - sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ - "${physical_dockerfile}" | head -1 -} - -resolve_image_digest() { - local image_ref="$1" - local attempts="${ROCM_IMAGE_DIGEST_ATTEMPTS:-4}" - local delay_secs="${ROCM_IMAGE_DIGEST_RETRY_DELAY:-2}" - local attempt=0 - local digest="" - local output="" - local status=0 - - if [[ "${image_ref}" =~ @(sha256:[0-9a-f]{64})$ ]]; then - printf '%s\n' "${BASH_REMATCH[1]}" - return - fi - if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \ - || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then - echo "Invalid image digest retry configuration" >&2 - return 1 - fi - - for ((attempt = 1; attempt <= attempts; attempt++)); do - status=0 - output=$(docker buildx imagetools inspect "${image_ref}" 2>&1) || status=$? - digest=$(awk '$1 == "Digest:" { print $2; exit }' <<< "${output}") - if ((status == 0)) && [[ "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then - printf '%s\n' "${digest}" - return - fi - if ((attempt < attempts)); then - printf \ - 'Image digest lookup failed for %s (%d/%d, status %d); retrying\n' \ - "${image_ref}" "${attempt}" "${attempts}" "${status}" >&2 - sleep "${delay_secs}" - fi - done - - printf 'Failed to resolve digest for %s (status %d)\n%s\n' \ - "${image_ref}" "${status}" "${output:-}" >&2 - return 1 -} - -resolve_dockerfile_arg_value() { - local dockerfile="$1" - local arg_name="$2" - local env_name="${arg_name}" - local value="" - - case "${arg_name}" in - ARG_PYTORCH_ROCM_ARCH) - env_name="PYTORCH_ROCM_ARCH" - ;; - max_jobs) - env_name="CI_MAX_JOBS" - ;; - esac - - value="${!env_name:-}" - if [[ -z "${value}" && "${env_name}" != "${arg_name}" ]]; then - value="${!arg_name:-}" - fi - if [[ -z "${value}" ]] && content_regular_file "${dockerfile}"; then - value=$(extract_dockerfile_arg_default "${dockerfile}" "${arg_name}") - fi - - printf '%s\n' "${value}" -} - -hash_dockerfile_arg_values() { - local dockerfile="$1" - local arg_name="" - local arg_value="" - local digest="" - shift || true - - for arg_name in "$@"; do - [[ -n "${arg_name}" ]] || continue - arg_value=$(resolve_dockerfile_arg_value "${dockerfile}" "${arg_name}") - if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then - if ! digest=$(resolve_image_digest "${arg_value}"); then - echo "Failed to resolve digest for BASE_IMAGE=${arg_value}" >&2 - return 1 - fi - printf 'arg:%s.digest=%s\n' "${arg_name}" "${digest}" - else - printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-}" - fi - done -} - -pin_base_image() { - local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" - local base_image="" - local digest="" - - base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE") - [[ -n "${base_image}" ]] || return 0 - if ! digest=$(resolve_image_digest "${base_image}"); then - echo "Error: could not resolve base image digest for ${base_image}" >&2 - echo "Refusing to compute content-addressed cache keys from a mutable tag." >&2 - return 1 - fi - BASE_IMAGE="${base_image%@*}@${digest}" - export BASE_IMAGE - echo "Pinned base image for this build: ${BASE_IMAGE}" -} - -is_ci_base_target() { - [[ "${TARGET}" == *"ci-base-rocm"* ]] -} - -is_commit_image_target() { - [[ -n "${IMAGE_TAG:-}" && -n "${BUILDKITE_COMMIT:-}" ]] || return 1 - is_ci_base_target && return 1 - return 0 -} - -should_upload_wheel_artifacts() { - [[ "${UPLOAD_ROCM_WHEEL_ARTIFACTS:-0}" == "1" ]] && return 0 - [[ "${TARGET}" == *"with-wheel"* \ - || "${TARGET}" == *"export-wheel"* \ - || "${TARGET}" == *"artifact"* ]] -} - -should_export_rocm_smoke() { - [[ "${TARGET}" == "test-rocm-ci-with-wheel" \ - || "${TARGET}" == "smoke-test-rocm-ci" ]] -} - -verify_rocm_smoke_export() { - local marker="./build/rocm-smoke-export/vllm-smoke-ok" - local expected_smoke_id="${BUILDKITE_BUILD_ID:-local}" - local actual_smoke_id="" - - should_export_rocm_smoke || return 0 - if [[ ! -f "${marker}" ]]; then - echo "ROCm BuildKit smoke marker is missing: ${marker}" >&2 - return 1 - fi - actual_smoke_id="$(< "${marker}")" - if [[ "${actual_smoke_id}" != "${expected_smoke_id}" ]]; then - echo "ROCm BuildKit smoke marker belongs to ${actual_smoke_id}, not ${expected_smoke_id}" \ - >&2 - return 1 - fi -} - -get_remote_image_label() { - local image_ref="$1" - local label_key="$2" - local format="{{ index .Image.Config.Labels \"${label_key}\" }}" - - docker buildx imagetools inspect "${image_ref}" \ - --format "${format}" 2>/dev/null | awk 'NF { print; exit }' || true -} - -remote_ci_base_identity_is_current_with_retry() { - local image_ref="$1" - local attempts="${CI_BASE_LABEL_ATTEMPTS:-4}" - local delay_secs="${CI_BASE_LABEL_RETRY_DELAY:-2}" - local identity="" - local remote_hash="" - local remote_version="" - local content_files_hash="" - local base_digest="" - local attempt=0 - local expected_version="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" - local format='{{ index .Image.Config.Labels "vllm.ci_base.content_hash" }}|{{ index .Image.Config.Labels "vllm.ci_base.metadata_version" }}|{{ index .Image.Config.Labels "vllm.ci_base.content_files_hash" }}|{{ index .Image.Config.Labels "vllm.rocm.base_image_digest" }}' - - if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \ - || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then - echo "Invalid ci_base label retry configuration" >&2 - return 1 - fi - - for ((attempt = 1; attempt <= attempts; attempt++)); do - identity=$(docker buildx imagetools inspect "${image_ref}" \ - --format "${format}" 2>/dev/null || true) - IFS='|' read -r \ - remote_hash remote_version content_files_hash base_digest <<< "${identity}" - if [[ "${remote_hash}" == "${CI_BASE_CONTENT_HASH:-}" \ - && "${remote_version}" == "${expected_version}" \ - && "${content_files_hash}" =~ ^[0-9a-f]{64}$ \ - && "${base_digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then - return 0 - fi - ((attempt == attempts)) || sleep "${delay_secs}" - done - - echo "ci_base identity did not match after ${attempts} attempts: ${image_ref}" >&2 - echo " expected hash/version: ${CI_BASE_CONTENT_HASH:-}/${expected_version}" >&2 - echo " observed hash/version: ${remote_hash:-}/${remote_version:-}" >&2 - return 1 -} - -registry_ref_exists_with_retry() { - local inspect_kind="$1" - local image_ref="$2" - local attempts="${ROCM_REGISTRY_PROBE_ATTEMPTS:-2}" - local delay_secs="${ROCM_REGISTRY_PROBE_RETRY_DELAY:-1}" - local output="" - local status=0 - local attempt=0 - - if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \ - || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then - echo "Invalid registry probe retry configuration" >&2 - return 1 - fi - - for ((attempt = 1; attempt <= attempts; attempt++)); do - status=0 - if [[ "${inspect_kind}" == "manifest" ]]; then - output=$(docker manifest inspect "${image_ref}" 2>&1) || status=$? - else - output=$(docker buildx imagetools inspect "${image_ref}" 2>&1) \ - || status=$? - fi - if ((status == 0)); then - return 0 - fi - # A definitive registry miss is normal for a new content key. Retry - # transport/rate-limit failures, but do not add latency to a real 404. - if grep -Eiq \ - 'manifest unknown|no such manifest|name unknown|not found|does not exist' \ - <<< "${output}"; then - return 1 - fi - if ((attempt < attempts)); then - echo "Registry probe failed for ${image_ref} (${attempt}/${attempts}); retrying" >&2 - sleep "${delay_secs}" - fi - done - - echo "Registry probe failed after ${attempts} attempts: ${image_ref}" >&2 - return 1 -} - -remote_image_exists() { - registry_ref_exists_with_retry manifest "$1" -} - -use_existing_builder() { - echo "Using existing builder: ${BUILDER_NAME}" - docker buildx use "${BUILDER_NAME}" - docker buildx inspect --bootstrap -} - -buildx_driver() { - local builder="${1:-}" - - if [[ -n "${builder}" ]]; then - docker buildx inspect "${builder}" 2>/dev/null - else - docker buildx inspect 2>/dev/null - fi | awk -F': *' '$1 == "Driver" { print $2; exit }' -} - -builder_supports_registry_cache() { - local driver="$1" - - [[ -n "${driver}" && "${driver}" != "docker" ]] -} - -create_and_bootstrap_builder() { - local driver="$1" - local endpoint="${2:-}" - - echo "Creating builder '${BUILDER_NAME}' with ${driver} driver" - if [[ -n "${endpoint}" ]]; then - docker buildx create \ - --name "${BUILDER_NAME}" \ - --driver "${driver}" \ - --use \ - "${endpoint}" - else - docker buildx create --name "${BUILDER_NAME}" --driver "${driver}" --use - fi - docker buildx inspect --bootstrap -} - -init_config() { - # shellcheck source=.buildkite/scripts/rocm/build-config.sh - source "$(dirname "${BASH_SOURCE[0]}")/rocm/build-config.sh" - configure_rocm_build - - TARGET="${1:-test-ci}" - if using_custom_rocm_dockerfiles; then - if [[ -z "${BASE_IMAGE:-}" ]]; then - echo "Custom ROCm builds require BASE_IMAGE from the selected base handoff" >&2 - return 1 - fi - if ! is_ci_base_target && [[ -z "${CI_BASE_IMAGE:-}" ]]; then - echo "Custom ROCm builds require CI_BASE_IMAGE from the selected ci_base handoff" >&2 - return 1 - fi - fi - BAKE_TARGETS=("${TARGET}") - DEPENDENCY_CACHE_TARGETS=() - CI_HCL_SOURCE="${CI_HCL_SOURCE:-${CI_HCL_FILE:-${DEFAULT_CI_HCL_SOURCE}}}" - VLLM_BAKE_FILE="${VLLM_BAKE_FILE:-docker/docker-bake-rocm.hcl}" - BUILDER_NAME="${BUILDER_NAME:-vllm-builder}" - BUILDKIT_SOCKET="${BUILDKIT_SOCKET:-/run/buildkit/buildkitd.sock}" - PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950}" - CI_BASE_CONTENT_FILES="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" - CI_BASE_DOCKERFILE="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" - if ! using_custom_rocm_dockerfiles; then - CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" - fi - CI_BASE_METADATA_VERSION="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" - CI_BASE_IMAGE_TAG="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" - export PYTORCH_ROCM_ARCH CI_BASE_DOCKERFILE - - SCRIPT_TMP_DIR=$(mktemp -d -t ci-bake-rocm.XXXXXX) - CI_HCL_PATH="${SCRIPT_TMP_DIR}/ci.hcl" - CI_BASE_LABEL_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/ci-base-label-override.hcl" - CSRC_CACHE_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-csrc-cache-override.hcl" - ROCM_ARG_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-arg-override.hcl" - BUILD_CONTEXT_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/build-context-override.hcl" - BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json" -} - -print_header() { - echo "--- :docker: Setting up Docker buildx bake" - echo "Target: ${TARGET}" - echo "CI HCL source: ${CI_HCL_SOURCE}" - echo "vLLM bake file: ${VLLM_BAKE_FILE}" - if is_ci_base_target; then - echo "Build mode: ci_base" - elif is_commit_image_target; then - echo "Build mode: build-scoped commit image" - else - echo "Build mode: generic" - fi - if [[ "${USE_SCCACHE:-0}" == "1" ]]; then - echo "Compiler cache: sccache enabled" - fi -} - -validate_inputs() { - if using_custom_rocm_dockerfiles; then - validate_rocm_dockerfile "${ROCM_BASE_DOCKERFILE}" || return $? - validate_rocm_dockerfile "${CI_BASE_DOCKERFILE}" \ - ci_base test export_vllm export_test_smoke csrc-build rust-build || return $? - fi - if [[ ! -f "${VLLM_BAKE_FILE}" ]]; then - echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE}" - echo "Make sure you're running from the vLLM repository root" - exit 1 - fi - - if [[ -n "${CI_HCL_SOURCE:-}" ]] && is_url_like "${CI_HCL_SOURCE}"; then - echo "Error: remote CI HCL sources are not supported: ${CI_HCL_SOURCE}" - echo "Use the vLLM-owned docker/ci-rocm.hcl or set CI_HCL_SOURCE to a local file." - exit 1 - fi - - if [[ -n "${CI_HCL_SOURCE:-}" && ! -f "${CI_HCL_SOURCE}" ]]; then - echo "Error: CI HCL file not found at ${CI_HCL_SOURCE}" - echo "Set CI_HCL_SOURCE to a local file if you need an override." - exit 1 - fi -} - -load_ci_hcl() { - echo "--- :page_facing_up: Loading ci.hcl" - cp "${CI_HCL_SOURCE}" "${CI_HCL_PATH}" - echo "Copied ${CI_HCL_SOURCE} to ${CI_HCL_PATH}" -} - -init_bake_files() { - BAKE_FILES=(-f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}") - if using_custom_rocm_dockerfiles; then - # Custom stacks reuse local cache without replacing the - # standard ROCm registry caches for the same commit or branch. - BAKE_FILES+=(--set '*.cache-to=') - fi -} - -compute_ci_base_hash_if_needed() { - if [[ -z "${CI_BASE_CONTENT_FILES:-}" ]]; then - return 0 - fi - pin_base_image - if ! is_ci_base_target; then - return 0 - fi - if [[ "${REMOTE_VLLM:-0}" != "0" ]]; then - echo "Error: content-addressed ci_base builds require REMOTE_VLLM=0" >&2 - return 1 - fi - - CI_BASE_CONTENT_HASH=$(compute_ci_base_content_hash) - export CI_BASE_CONTENT_HASH - echo "ci_base content hash: ${CI_BASE_CONTENT_HASH:0:16}..." -} - -wants_stable_ci_base_tag() { - if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then - return 1 - fi - - if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then - return 0 - fi - if [[ "${CI_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then - return 1 - fi - - [[ "${NIGHTLY:-0}" == "1" && "${BUILDKITE_BRANCH:-}" == "${CI_BASE_STABLE_BRANCH:-main}" ]] -} - -trusted_ci_base_tip_matches_build() { - local branch="${CI_BASE_STABLE_BRANCH:-main}" - local build_commit="${BUILDKITE_COMMIT:-}" - local remote_tip="" - - is_trusted_ci_cache_writer || return 1 - if [[ ! "${build_commit}" =~ ^[0-9a-fA-F]{40}$ ]]; then - echo "Skipping ci_base stable tag: Buildkite commit is missing or invalid" >&2 - return 1 - fi - remote_tip=$(git ls-remote --exit-code "${BUILDKITE_REPO}" \ - "refs/heads/${branch}" 2>/dev/null | awk 'NR == 1 { print $1 }') - if [[ ! "${remote_tip}" =~ ^[0-9a-fA-F]{40}$ ]]; then - echo "Skipping ci_base stable tag: could not resolve remote ${branch} tip" >&2 - return 1 - fi - if [[ "${remote_tip,,}" != "${build_commit,,}" ]]; then - echo "Skipping ci_base stable tag: ${branch} advanced from ${build_commit} to ${remote_tip}" >&2 - return 1 - fi -} - -should_push_stable_ci_base_tag() { - wants_stable_ci_base_tag \ - && is_trusted_ci_cache_writer \ - && trusted_ci_base_tip_matches_build -} - -ci_base_tag_with_suffix() { - local base_tag="$1" - local suffix="$2" - - printf '%s-%s\n' "${base_tag}" "$(clean_docker_tag "${suffix}")" -} - -configure_ci_base_image_refs() { - local stable_tag="${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" - local metadata_version="${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" - local scope="${CI_BASE_WRITE_SCOPE:-}" - local trusted_content_tag="" - local content_tag="" - local build_tag="" - local primary_tag="" - - if [[ ! "${metadata_version}" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,15}$ ]]; then - echo "Invalid ci_base metadata version: ${metadata_version}" >&2 - return 1 - fi - - if [[ "${BUILDKITE:-false}" == "true" ]] \ - && ! is_full_git_sha "${BUILDKITE_COMMIT:-}"; then - echo "Invalid Buildkite commit for ci_base handoff: ${BUILDKITE_COMMIT:-}" >&2 - return 1 - fi - if [[ "${BUILDKITE:-false}" == "true" \ - && -z "${BUILDKITE_BUILD_ID:-}" ]]; then - echo "Buildkite build ID is required for the ci_base runtime handoff" >&2 - return 1 - fi - if [[ -n "${BUILDKITE_BUILD_ID:-}" ]]; then - build_tag=$(ci_base_tag_with_suffix \ - "${stable_tag}" "build-${BUILDKITE_BUILD_ID}") - fi - CI_BASE_IMAGE_TAG_BUILD_REF="${build_tag}" - export CI_BASE_IMAGE_TAG_BUILD_REF - - if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then - if is_ci_base_target; then - echo "Error: ci_base builds require a content hash" >&2 - return 1 - fi - CI_BASE_IMAGE="${CI_BASE_IMAGE:-${stable_tag}}" - export CI_BASE_IMAGE - return 0 - fi - - trusted_content_tag=$(ci_base_tag_with_suffix \ - "${stable_tag}" "v${metadata_version}-${CI_BASE_CONTENT_HASH}") - if [[ -n "${scope}" ]]; then - content_tag=$(ci_base_tag_with_suffix \ - "${stable_tag}" "v${metadata_version}-${scope}-${CI_BASE_CONTENT_HASH}") - else - content_tag="${trusted_content_tag}" - fi - CI_BASE_IMAGE_TAG_CONTENT_REF="${content_tag}" - CI_BASE_TRUSTED_CONTENT_REF="${trusted_content_tag}" - primary_tag="${build_tag:-${content_tag}}" - - # Main writes canonical content refs. Other builds import those refs - # read-only and publish into a source-scoped preview namespace. - # A fresh CI build first publishes its unique runtime tag. Content and - # stable aliases are created only after that image passes identity checks. - CI_BASE_STABLE_PROMOTION_REF="${stable_tag}" - CI_BASE_IMAGE_TAG="${primary_tag}" - export CI_BASE_IMAGE_TAG - export CI_BASE_IMAGE_TAG_CONTENT_REF - export CI_BASE_TRUSTED_CONTENT_REF - export CI_BASE_STABLE_PROMOTION_REF - - if is_ci_base_target; then - IMAGE_TAG="${primary_tag}" - CI_BASE_IMAGE="${primary_tag}" - export CI_BASE_IMAGE - export IMAGE_TAG - - echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}" - if [[ -n "${build_tag}" ]]; then - echo "ci_base build image tag: ${build_tag}" - fi - echo "ci_base content image tag: ${content_tag}" - if wants_stable_ci_base_tag && is_trusted_ci_cache_writer; then - echo "ci_base stable alias is eligible for post-build promotion: ${stable_tag}" - else - echo "ci_base stable alias will not be pushed for this build" - echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}" - fi - return 0 - fi - - if [[ -z "${CI_BASE_IMAGE:-}" || "${CI_BASE_IMAGE}" == "${stable_tag}" ]]; then - CI_BASE_IMAGE="${content_tag}" - export CI_BASE_IMAGE - echo "Using ci_base image: ${CI_BASE_IMAGE}" - else - echo "Using provided CI_BASE_IMAGE override: ${CI_BASE_IMAGE}" - fi -} - -publish_ci_base_handoff_ref() { - local source_ref="${1:-${CI_BASE_IMAGE_TAG_BUILD_REF:-${CI_BASE_IMAGE_TAG_CONTENT_REF:-}}}" - local content_ref="${CI_BASE_IMAGE_TAG_CONTENT_REF:-}" - local digest="" - local handoff_ref="" - - is_ci_base_target || return 0 - if [[ -z "${source_ref}" || -z "${content_ref}" ]]; then - echo "Cannot publish ci_base handoff without source and content refs" >&2 - return 1 - fi - if ! digest=$(resolve_image_digest "${source_ref}"); then - echo "Could not resolve immutable ci_base handoff: ${source_ref}" >&2 - return 1 - fi - - handoff_ref="${content_ref%@*}@${digest}" - if ! confirm_remote_image_push "${handoff_ref}"; then - echo "Could not validate immutable ci_base handoff: ${handoff_ref}" >&2 - return 1 - fi - if command -v buildkite-agent >/dev/null 2>&1; then - if ! buildkite-agent meta-data set "rocm-ci-base-image" "${handoff_ref}"; then - echo "Could not publish required ci_base handoff metadata" >&2 - return 1 - fi - elif [[ "${BUILDKITE:-false}" == "true" ]]; then - echo "buildkite-agent not found; cannot publish ci_base handoff" >&2 - return 1 - fi - echo "Published immutable ci_base handoff: ${handoff_ref}" -} - -ci_base_output_refs() { - printf '%s\n' \ - "${CI_BASE_IMAGE_TAG_CONTENT_REF:-}" \ - "${CI_BASE_IMAGE_TAG_BUILD_REF:-}" \ - | awk 'NF && !seen[$0]++' -} - -ci_base_candidate_refs() { - printf '%s\n' \ - "${CI_BASE_TRUSTED_CONTENT_REF:-}" \ - "${CI_BASE_IMAGE_TAG_CONTENT_REF:-}" \ - "${CI_BASE_STABLE_PROMOTION_REF:-}" \ - | awk 'NF && !seen[$0]++' -} - -find_matching_ci_base_ref() { - local candidate="" - local candidate_digest="" - local immutable_candidate="" - - while IFS= read -r candidate; do - [[ -n "${candidate}" ]] || continue - # Missing content refs are normal after an input change. Avoid paying - # the full digest retry budget for each definitely absent candidate. - remote_image_exists "${candidate}" || continue - if ! candidate_digest=$(resolve_image_digest "${candidate}"); then - echo "Could not pin ci_base candidate: ${candidate}" >&2 - continue - fi - immutable_candidate="${candidate%@*}@${candidate_digest}" - if remote_ci_base_identity_is_current_with_retry "${immutable_candidate}"; then - printf '%s\n' "${immutable_candidate}" - return 0 - fi - done < <(ci_base_candidate_refs) - - return 1 -} - -refresh_ci_base_tags_from_ref() { - local source_ref="$1" - local source_digest="" - local immutable_source="" - local tag="" - local tag_digest="" - - if ! source_digest=$(resolve_image_digest "${source_ref}"); then - echo "Could not resolve selected ci_base image: ${source_ref}" >&2 - return 1 - fi - immutable_source="${source_ref%@*}@${source_digest}" - - while IFS= read -r tag; do - [[ -n "${tag}" ]] || continue - [[ "${tag}" != "${source_ref}" ]] || continue - tag_digest="" - if remote_image_exists "${tag}"; then - tag_digest=$(resolve_image_digest "${tag}" || true) - fi - if [[ "${tag_digest}" == "${source_digest}" ]]; then - echo "ci_base tag is already current: ${tag}" - continue - fi - echo "Updating ci_base tag ${tag} -> ${immutable_source}" - if ! docker buildx imagetools create --prefer-index=false \ - -t "${tag}" "${immutable_source}"; then - echo "Failed to update ci_base tag ${tag} from ${immutable_source}" >&2 - return 1 - fi - if ! tag_digest=$(resolve_image_digest "${tag}") \ - || [[ "${tag_digest}" != "${source_digest}" ]]; then - echo "Updated ci_base tag does not resolve to the selected digest: ${tag}" >&2 - return 1 - fi - done < <(ci_base_output_refs) -} - -promote_stable_ci_base_tag() { - local source_ref="${1:-${CI_BASE_IMAGE_TAG_BUILD_REF:-${CI_BASE_IMAGE_TAG_CONTENT_REF:-}}}" - local stable_ref="${CI_BASE_STABLE_PROMOTION_REF:-}" - local digest="" - local immutable_source="" - - is_ci_base_target || return 0 - wants_stable_ci_base_tag || return 0 - if ! should_push_stable_ci_base_tag; then - echo "Skipping ci_base stable promotion: build is not the trusted current main tip" - return 0 - fi - if [[ -z "${source_ref}" || -z "${stable_ref}" ]]; then - echo "Cannot promote ci_base stable tag without source and destination refs" >&2 - return 1 - fi - if ! digest=$(resolve_image_digest "${source_ref}"); then - echo "Could not pin ci_base source before stable promotion: ${source_ref}" >&2 - return 1 - fi - immutable_source="${source_ref%@*}@${digest}" - echo "Promoting ci_base stable tag from ${immutable_source}" - if ! docker buildx imagetools create --prefer-index=false \ - -t "${stable_ref}" "${immutable_source}"; then - echo "Failed to promote ci_base stable tag" >&2 - return 1 - fi - if ! confirm_remote_image_push "${stable_ref}"; then - echo "Promoted ci_base stable tag did not become visible: ${stable_ref}" >&2 - return 1 - fi -} - -maybe_reuse_matching_ci_base_ref() { - local matching_ref="" - - matching_ref=$(find_matching_ci_base_ref) || return 1 - - echo "Found existing ci_base image with matching content hash: ${matching_ref}" - if ! refresh_ci_base_tags_from_ref "${matching_ref}"; then - echo "ci_base tag refresh failed after finding an exact image; aborting" >&2 - return 2 - fi - if ! promote_stable_ci_base_tag "${matching_ref}"; then - echo "ci_base stable promotion failed after finding an exact image; aborting" >&2 - return 2 - fi - if ! publish_ci_base_handoff_ref "${matching_ref}"; then - echo "ci_base handoff failed after finding an exact image; aborting" >&2 - return 2 - fi - echo "Content hashes match -- ci_base is current" - return 0 -} - -maybe_skip_existing_image() { - local remote_revision="" - local reuse_status=0 - - if [[ -z "${IMAGE_TAG:-}" ]]; then - return 0 - fi - - if [[ "${FORCE_BUILD:-0}" == "1" ]]; then - echo "FORCE_BUILD=1 set; skipping existing-image check" - return 0 - fi - if ! is_ci_base_target \ - && { should_upload_wheel_artifacts || should_export_rocm_smoke; }; then - echo "Local-output targets always run for the current build" - return 0 - fi - - echo "--- :mag: Checking image tag" - echo "Image tag: ${IMAGE_TAG}" - - if is_ci_base_target && [[ -n "${CI_BASE_CONTENT_HASH:-}" ]]; then - maybe_reuse_matching_ci_base_ref || reuse_status=$? - case "${reuse_status}" in - 0) - echo "Skipping build" - exit 0 - ;; - 1) - echo "No current ci_base image matched the expected content hash" - echo "Proceeding with build" - return 0 - ;; - *) - return "${reuse_status}" - ;; - esac - fi - - if ! remote_image_exists "${IMAGE_TAG}"; then - echo "Image not found, proceeding with build" - return 0 - fi - - if is_ci_base_target; then - echo "ci_base image already exists and no content hash was configured" - echo "Skipping build" - exit 0 - fi - - if is_commit_image_target; then - remote_revision=$(get_remote_image_label "${IMAGE_TAG}" "org.opencontainers.image.revision") - if [[ -n "${remote_revision}" && "${remote_revision}" != "${BUILDKITE_COMMIT}" ]]; then - echo "Existing image revision does not match ${BUILDKITE_COMMIT}" - echo " found revision: ${remote_revision}" - echo "Rebuilding image" - return 0 - fi - - echo "Commit image already exists: ${IMAGE_TAG}" - echo "Skipping build" - exit 0 - fi - - echo "Image already exists: ${IMAGE_TAG}" - echo "Skipping build" - exit 0 -} - -setup_builder() { - echo "--- :buildkite: Setting up buildx builder" - - local setup_mode="${ROCM_SETUP_BUILDX_BUILDER:-auto}" - local current_driver="" - local named_driver="" - - if [[ "${setup_mode}" == "0" || "${setup_mode}" == "false" ]]; then - echo "Using current Docker buildx builder" - echo "ROCM_SETUP_BUILDX_BUILDER=${setup_mode}; cache exporters may fail if the driver is docker" - docker buildx inspect --bootstrap - echo "Active builder:" - docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls - return 0 - fi - - current_driver=$(buildx_driver || true) - if [[ "${setup_mode}" != "1" ]] && builder_supports_registry_cache "${current_driver}"; then - echo "Using current Docker buildx builder with ${current_driver} driver" - docker buildx inspect --bootstrap - echo "Active builder:" - docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls - return 0 - fi - - if [[ "${setup_mode}" != "1" ]]; then - echo "Current buildx driver '${current_driver:-unknown}' cannot export registry caches" - echo "Creating or using a cache-capable builder: ${BUILDER_NAME}" - fi - - if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then - named_driver=$(buildx_driver "${BUILDER_NAME}" || true) - if ! builder_supports_registry_cache "${named_driver}"; then - echo "Builder '${BUILDER_NAME}' uses ${named_driver:-unknown} driver; using ${BUILDER_NAME}-cache instead" - BUILDER_NAME="${BUILDER_NAME}-cache" - fi - fi - - if [[ -S "${BUILDKIT_SOCKET}" ]]; then - echo "Found local buildkitd socket at ${BUILDKIT_SOCKET}" - echo "Using remote driver to connect to buildkitd" - - if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then - use_existing_builder - else - create_and_bootstrap_builder remote "unix://${BUILDKIT_SOCKET}" - fi - elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then - use_existing_builder - else - echo "No local buildkitd found, using docker-container driver" - create_and_bootstrap_builder docker-container - fi - - echo "Active builder:" - docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls -} - -validate_cache_branch_tag() { - local name="$1" - local value="$2" - - if [[ -n "${value}" \ - && (! "${value}" =~ ^[A-Za-z0-9_][A-Za-z0-9_.-]*$ \ - || ${#value} -gt ${ROCM_CACHE_BRANCH_TAG_MAX_LEN}) ]]; then - echo "Invalid ${name}; expected a Docker tag component of at most ${ROCM_CACHE_BRANCH_TAG_MAX_LEN} characters" >&2 - return 1 - fi -} - -prepare_git_cache_metadata() { - local cache_branch_name="" - local cache_base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" - local target_repo_slug="" - local target_repo_url="" - local merge_base_ref="" - - if is_ci_base_target; then - echo "Skipping commit-cache ancestry lookup for content-addressed ci_base" - return 0 - fi - - if [[ -z "${PARENT_COMMIT:-}" || -z "${VLLM_MERGE_BASE_COMMIT:-}" ]] \ - && git rev-parse --is-shallow-repository 2>/dev/null | grep -q "true"; then - echo "Shallow clone detected - deepening for cache key computation" - git_fetch_for_cache --filter=tree:0 --no-tags --deepen=1 \ - origin "$(git rev-parse HEAD)" || true - fi - - if [[ -z "${PARENT_COMMIT:-}" ]]; then - PARENT_COMMIT=$(git rev-parse HEAD~1 2>/dev/null || echo "") - if [[ -n "${PARENT_COMMIT}" ]]; then - export PARENT_COMMIT - echo "Computed parent commit for cache fallback: ${PARENT_COMMIT}" - else - echo "Could not determine parent commit" - fi - else - echo "Using provided PARENT_COMMIT: ${PARENT_COMMIT}" - fi - - if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then - cache_branch_name=$(select_cache_branch_name) - if [[ -z "${cache_branch_name}" && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then - cache_branch_name="pr-${BUILDKITE_PULL_REQUEST}" - echo "Using pull request number for ROCm branch cache tag: ${cache_branch_name}" - fi - fi - - if [[ -z "${ROCM_CACHE_BRANCH_TAG:-}" && -n "${cache_branch_name}" ]]; then - ROCM_CACHE_BRANCH_TAG=$( - compose_cache_branch_tag "$(get_buildkite_repo_slug)" "${cache_branch_name}" - ) - export ROCM_CACHE_BRANCH_TAG - echo "Computed ROCm branch cache tag: ${ROCM_CACHE_BRANCH_TAG} (from ${cache_branch_name})" - elif [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then - echo "Using provided ROCM_CACHE_BRANCH_TAG: ${ROCM_CACHE_BRANCH_TAG}" - elif [[ -n "${BUILDKITE_BRANCH:-}" ]]; then - echo "Skipping ROCm branch cache tag: no usable branch name found" - echo " BUILDKITE_BRANCH=${BUILDKITE_BRANCH}" - fi - - if [[ -z "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" \ - && -n "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-}" \ - && "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then - target_repo_slug=$(get_buildkite_target_repo_slug) - ROCM_CACHE_UPSTREAM_BRANCH_TAG=$( - compose_cache_branch_tag "${target_repo_slug}" "${BUILDKITE_PULL_REQUEST_BASE_BRANCH}" - ) - export ROCM_CACHE_UPSTREAM_BRANCH_TAG - echo "Computed ROCm upstream branch cache tag: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" - elif [[ -n "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" ]]; then - echo "Using provided ROCM_CACHE_UPSTREAM_BRANCH_TAG: ${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" - fi - - validate_cache_branch_tag \ - ROCM_CACHE_BRANCH_TAG "${ROCM_CACHE_BRANCH_TAG:-}" - validate_cache_branch_tag \ - ROCM_CACHE_UPSTREAM_BRANCH_TAG "${ROCM_CACHE_UPSTREAM_BRANCH_TAG:-}" - - if [[ -z "${VLLM_MERGE_BASE_COMMIT:-}" ]]; then - target_repo_url=$(get_buildkite_target_repo_url) - merge_base_ref="refs/remotes/vllm-cache-upstream/${cache_base_branch}" - git_fetch_for_cache --no-tags --depth=200 "${target_repo_url}" \ - "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true - VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "") - if [[ -z "${VLLM_MERGE_BASE_COMMIT}" ]]; then - git_fetch_for_cache --no-tags --deepen=1000 "${target_repo_url}" \ - "+refs/heads/${cache_base_branch}:${merge_base_ref}" 2>/dev/null || true - VLLM_MERGE_BASE_COMMIT=$(git merge-base HEAD "${merge_base_ref}" 2>/dev/null || echo "") - fi - if [[ -n "${VLLM_MERGE_BASE_COMMIT}" ]]; then - export VLLM_MERGE_BASE_COMMIT - echo "Computed merge base commit for cache fallback: ${VLLM_MERGE_BASE_COMMIT}" - else - echo "Could not determine merge base with ${cache_base_branch}" - fi - else - echo "Using provided VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}" - fi -} - -ci_base_metadata_pairs() { - local dockerfile="${CI_BASE_DOCKERFILE:-${DEFAULT_CI_BASE_DOCKERFILE}}" - local stages="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" - local content_files="${CI_BASE_CONTENT_FILES:-${DEFAULT_CI_BASE_CONTENT_FILES}}" - local content_files_hash="" - local base_image="" - local base_image_digest="" - local -a content_paths=() - local -a content_args=() - - read -r -a content_paths <<< "${content_files}" - if [[ ${#content_paths[@]} -gt 0 ]]; then - if ! content_files_hash=$(compute_content_hash "${content_paths[@]}"); then - echo "Failed to hash ci_base metadata content files" >&2 - return 1 - fi - fi - mapfile -t content_args < <( - get_content_arg_names "${dockerfile}" "${stages}" "${CI_BASE_CONTENT_ARGS:-}" - ) - - base_image=$(resolve_dockerfile_arg_value "${dockerfile}" "BASE_IMAGE") - if [[ -n "${base_image}" ]]; then - if ! base_image_digest=$(resolve_image_digest "${base_image}"); then - echo "Failed to resolve ci_base metadata digest for ${base_image}" >&2 - return 1 - fi - fi - - metadata_pair "vllm.ci_base.metadata_version" "${CI_BASE_METADATA_VERSION:-${DEFAULT_CI_BASE_METADATA_VERSION}}" - metadata_pair "vllm.ci_base.content_hash" "${CI_BASE_CONTENT_HASH:-}" - metadata_pair "vllm.ci_base.content_files_hash" "${content_files_hash}" - metadata_pair "vllm.ci_base.content_files" "${content_files}" - metadata_pair "vllm.ci_base.content_args" "$(join_words "${content_args[@]}")" - metadata_pair "vllm.ci_base.dockerfile" "${dockerfile}" - metadata_pair "vllm.ci_base.dockerfile_stages" "${stages}" - - # The parent identity is digest-only: mutable aliases that resolve to the - # same image must produce byte-identical canonical metadata. - metadata_pair "vllm.rocm.base_image" "${base_image_digest}" - metadata_pair "vllm.rocm.base_image_digest" "${base_image_digest}" - metadata_pair "vllm.rocm.pytorch_rocm_arch" "${PYTORCH_ROCM_ARCH:-}" - metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")" - metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")" - metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")" - metadata_pair "vllm.rocm.max_jobs" "$(resolve_dockerfile_arg_value "${dockerfile}" "max_jobs")" - metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")" - metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}" - metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")" - metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}" - metadata_pair "vllm.rocm.deepep_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_REPO")" - metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}" - metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")" - metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")" - metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}" - metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}" -} - -write_ci_base_metadata_annotations() { - local metadata="$1" - local key="" - local value="" - - [[ -n "${metadata}" ]] || return 0 - while IFS=$'\t' read -r key value; do - [[ -n "${key}" && -n "${value}" ]] || continue - printf ' "%s",\n' \ - "$(hcl_escape_string "manifest:${key}=${value}")" - done <<< "${metadata}" -} - -write_ci_base_metadata_labels() { - local metadata="$1" - local key="" - local value="" - - [[ -n "${metadata}" ]] || return 0 - while IFS=$'\t' read -r key value; do - [[ -n "${key}" && -n "${value}" ]] || continue - printf ' "%s" = "%s"\n' \ - "$(hcl_escape_string "${key}")" \ - "$(hcl_escape_string "${value}")" - done <<< "${metadata}" -} - -write_ci_base_label_override() { - local target_name="" - local metadata="" - local -a ci_base_targets=() - - if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then - return 0 - fi - - mapfile -t ci_base_targets < <( - { - printf '%s\n' "ci-base-rocm" - sed -n -E 's/^target "(ci-base-rocm[^"]+)".*/\1/p' "${CI_HCL_PATH}" 2>/dev/null || true - } | awk '!seen[$0]++' - ) - - if [[ ${#ci_base_targets[@]} -eq 0 ]]; then - return 0 - fi - - metadata=$(ci_base_metadata_pairs) - - : > "${CI_BASE_LABEL_OVERRIDE_PATH}" - for target_name in "${ci_base_targets[@]}"; do - cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <> "${CI_BASE_LABEL_OVERRIDE_PATH}" - cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" <> "${CI_BASE_LABEL_OVERRIDE_PATH}" - cat >> "${CI_BASE_LABEL_OVERRIDE_PATH}" < "${ROCM_ARG_OVERRIDE_PATH}" - - BAKE_FILES+=(-f "${ROCM_ARG_OVERRIDE_PATH}") - echo "Appended resolved ROCm Docker ARG override" -} - -write_hcl_string_list_attr() { - local indent="$1" - local attr="$2" - shift 2 - - printf '%s%s = [\n' "${indent}" "${attr}" - write_hcl_string_list_entries "${indent} " "$@" - printf '%s]\n' "${indent}" -} - -validate_cache_export_mode() { - local mode="$1" - local env_name="$2" - - case "${mode}" in - min|max) - ;; - *) - echo "Error: ${env_name} must be one of: min, max" - exit 1 - ;; - esac -} - -validate_content_cache_export_mode() { - local mode="$1" - local env_name="$2" - - case "${mode}" in - missing|always|never) - ;; - *) - echo "Error: ${env_name} must be one of: missing, always, never" - exit 1 - ;; - esac -} - -should_export_content_cache_ref() { - local cache_ref="$1" - local cache_name="$2" - local trusted_ref="${3:-${cache_ref}}" - local mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}" - - case "${mode}" in - always) - echo "${cache_name} content cache export mode is always" - return 0 - ;; - never) - echo "${cache_name} content cache export mode is never" - return 1 - ;; - missing|"") - if registry_ref_exists_with_retry imagetools "${trusted_ref}"; then - echo "${cache_name} trusted content cache is visible: ${trusted_ref}" - return 1 - fi - if [[ "${cache_ref}" != "${trusted_ref}" ]] \ - && registry_ref_exists_with_retry imagetools "${cache_ref}"; then - echo "${cache_name} scoped content cache is visible: ${cache_ref}" - return 1 - fi - echo "${cache_name} content cache is missing: ${cache_ref}" - return 0 - ;; - *) - echo "Error: ROCM_CONTENT_CACHE_EXPORT_MODE must be one of: missing, always, never" - exit 1 - ;; - esac -} - -write_rocm_cache_override() { - local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}" - local content_cache_export_mode="${ROCM_CONTENT_CACHE_EXPORT_MODE:-missing}" - local csrc_cache_to_mode="${ROCM_CSRC_CACHE_TO_MODE:-max}" - local rust_cache_to_mode="${ROCM_RUST_CACHE_TO_MODE:-max}" - local rocm_cache_to_mode="${ROCM_FINAL_CACHE_TO_MODE:-min}" - local -a csrc_content_cache_from=() - local -a rust_content_cache_from=() - local -a combined_content_cache_from=() - local -a csrc_cache_to=() - local -a rust_cache_to=() - local -a rocm_cache_to=() - local -a export_wheel_cache_to=() - local export_csrc_cache=1 - local export_rust_cache=1 - - if ! uses_rocm_csrc_cache && ! uses_rocm_rust_cache; then - return 0 - fi - - validate_content_cache_export_mode \ - "${content_cache_export_mode}" \ - "ROCM_CONTENT_CACHE_EXPORT_MODE" - validate_cache_export_mode "${csrc_cache_to_mode}" "ROCM_CSRC_CACHE_TO_MODE" - validate_cache_export_mode "${rust_cache_to_mode}" "ROCM_RUST_CACHE_TO_MODE" - validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE" - echo "ROCm content cache export mode: ${content_cache_export_mode}" - echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}" - echo "ROCm Rust fallback cache export mode: ${rust_cache_to_mode}" - echo "ROCm final image cache export mode: ${rocm_cache_to_mode}" - - if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then - csrc_content_cache_from+=( - "type=registry,ref=${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}" - ) - if [[ "${ROCM_CSRC_CONTENT_CACHE_REF}" != \ - "${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}" ]]; then - csrc_content_cache_from+=( - "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF}" - ) - fi - if should_export_content_cache_ref \ - "${ROCM_CSRC_CONTENT_CACHE_REF}" "ROCm csrc" \ - "${ROCM_CSRC_TRUSTED_CONTENT_CACHE_REF}"; then - csrc_cache_to+=( - "type=registry,ref=${ROCM_CSRC_CONTENT_CACHE_REF},mode=${csrc_cache_to_mode},ignore-error=true" - ) - else - export_csrc_cache=0 - fi - fi - - if [[ -n "${ROCM_RUST_CONTENT_CACHE_REF:-}" ]]; then - rust_content_cache_from+=( - "type=registry,ref=${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}" - ) - if [[ "${ROCM_RUST_CONTENT_CACHE_REF}" != \ - "${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}" ]]; then - rust_content_cache_from+=( - "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF}" - ) - fi - # Legacy commit/branch exports are only needed while the exact-input - # content ref is absent. The exact ref itself is refreshed below. - if ! should_export_content_cache_ref \ - "${ROCM_RUST_CONTENT_CACHE_REF}" "ROCm Rust" \ - "${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}"; then - export_rust_cache=0 - fi - if [[ "${content_cache_export_mode}" != "never" ]]; then - # Refresh the exact-input ref in the original solve regardless of - # which local or remote cache supplied the Rust result. - rust_cache_to+=( - "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=min,ignore-error=true" - ) - echo "ROCm Rust exact-input cache will be refreshed (mode=min): ${ROCM_RUST_CONTENT_CACHE_REF}" - fi - fi - - combined_content_cache_from=("${csrc_content_cache_from[@]}" "${rust_content_cache_from[@]}") - - # Docker Hub cache exports are best-effort. A cache-only target failure can - # otherwise cancel the sibling image target before its manifest is pushed. - if [[ -n "${BUILDKITE_COMMIT:-}" ]]; then - if [[ ${export_csrc_cache} -eq 1 ]]; then - csrc_cache_to+=( - "type=registry,ref=${cache_repo}:csrc-rocm-${BUILDKITE_COMMIT},mode=${csrc_cache_to_mode},ignore-error=true" - ) - fi - if [[ ${export_rust_cache} -eq 1 ]]; then - rust_cache_to+=( - "type=registry,ref=${cache_repo}:rust-rocm-${BUILDKITE_COMMIT},mode=${rust_cache_to_mode},ignore-error=true" - ) - fi - rocm_cache_to+=( - "type=registry,ref=${cache_repo}:rocm-${BUILDKITE_COMMIT},mode=${rocm_cache_to_mode},ignore-error=true" - ) - fi - - if [[ -n "${ROCM_CACHE_BRANCH_TAG:-}" ]]; then - if [[ ${export_csrc_cache} -eq 1 ]]; then - csrc_cache_to+=( - "type=registry,ref=${cache_repo}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${csrc_cache_to_mode},ignore-error=true" - ) - fi - if [[ ${export_rust_cache} -eq 1 ]]; then - rust_cache_to+=( - "type=registry,ref=${cache_repo}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rust_cache_to_mode},ignore-error=true" - ) - fi - rocm_cache_to+=( - "type=registry,ref=${cache_repo}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${rocm_cache_to_mode},ignore-error=true" - ) - fi - - # Standalone image/wheel targets reach rust-build but not the cache-only - # exporter unless it is requested explicitly. - if ((${#rust_cache_to[@]} > 0)); then - case "${TARGET}" in - test-rocm-ci|export-wheel-rocm|smoke-test-rocm-ci) - BAKE_TARGETS=("rust-rocm-ci" "${BAKE_TARGETS[@]}") - ;; - esac - fi - - if [[ "${TARGET}" == "test-rocm-ci-with-wheel" ]]; then - export_wheel_cache_to=() - else - export_wheel_cache_to=("${rocm_cache_to[@]}") - fi - - { - cat < "${CSRC_CACHE_OVERRIDE_PATH}" - - BAKE_FILES+=(-f "${CSRC_CACHE_OVERRIDE_PATH}") - echo "Appended ROCm cache override with non-fatal registry exports" -} - -extract_dependency_pins() { - local bake_dir="" - local dockerfile_rocm="" - local physical_dockerfile="" - local var="" - local val="" - - bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" - physical_dockerfile="${dockerfile_rocm}" - if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile_rocm}" != /* ]]; then - physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile_rocm}" - fi - if [[ ! -f "${physical_dockerfile}" ]]; then - return 0 - fi - - for var in NIXL_BRANCH UCX_BRANCH DEEPEP_BRANCH; do - if [[ -n "${!var:-}" ]]; then - echo "Using provided ${var}: ${!var}" - continue - fi - - val=$( - sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ - "${physical_dockerfile}" | head -1 - ) - if [[ -n "${val}" ]]; then - export "${var}=${val}" - echo "Extracted ${var}=${val} from ${dockerfile_rocm}" - fi - done -} - -compute_dependency_cache_keys() { - local bake_dir="" - local dockerfile_rocm="" - local nixl_branch="" - local ucx_branch="" - local deepep_branch="" - local nixl_material="" - local deepep_material="" - - bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" - nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH") - ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH") - deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH") - - if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then - nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl") - NIXL_CACHE_KEY=$( - compose_dependency_cache_key \ - "${nixl_branch}-ucx-${ucx_branch}" \ - "${nixl_material}" - ) - export NIXL_CACHE_KEY - echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}" - fi - - if [[ -n "${deepep_branch}" ]]; then - deepep_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_deepep") - DEEPEP_CACHE_KEY=$( - compose_dependency_cache_key \ - "${deepep_branch}" \ - "${deepep_material}" - ) - export DEEPEP_CACHE_KEY - echo "DeepEP dependency cache key: ${DEEPEP_CACHE_KEY}" - fi -} - -compose_stage_cache_material() { - local dockerfile="$1" - local stages="$2" - local -a content_args=() - - mapfile -t content_args < <(get_content_arg_names "${dockerfile}" "${stages}" "") - { - printf 'dockerfile:%s\n' "${dockerfile}" - printf 'dockerfile-stages:%s\n' "${stages}" - hash_dockerfile_stages "${dockerfile}" "${stages}" - printf 'resolved-build-args:\n' - hash_dockerfile_arg_values "${dockerfile}" "${content_args[@]}" - } -} - -dependency_cache_ref_exists() { - local cache_ref="$1" - registry_ref_exists_with_retry imagetools "${cache_ref}" -} - -dependency_cache_ref_for_target() { - local target="$1" - local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}" - - case "${target}" in - nixl-rocm-ci) - if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then - printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}" - elif [[ -n "${NIXL_BRANCH:-}" ]]; then - printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}" - fi - ;; - deepep-rocm-ci) - if [[ -n "${DEEPEP_CACHE_KEY:-}" ]]; then - printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_CACHE_KEY}" - elif [[ -n "${DEEPEP_BRANCH:-}" ]]; then - printf '%s\n' "${cache_repo}:deepep-rocm-${DEEPEP_BRANCH}" - fi - ;; - esac -} - -add_dependency_cache_target() { - local target="$1" - - if printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" | grep -qx "${target}"; then - return 0 - fi - DEPENDENCY_CACHE_TARGETS+=("${target}") -} - -resolve_ci_base_dependency_targets() { - local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}" - local nixl_ref="" - local deepep_ref="" - - [[ "${TARGET}" == "ci-base-rocm-ci-with-deps" ]] || return 0 - - case "${mode}" in - always) - echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially" - for target in nixl-rocm-ci deepep-rocm-ci; do - if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then - add_dependency_cache_target "${target}" - fi - done - ;; - never) - BAKE_TARGETS=("ci-base-rocm-ci") - DEPENDENCY_CACHE_TARGETS=() - echo "ROCM_DEP_CACHE_EXPORT_MODE=never; building ci_base without dependency cache exports" - return 0 - ;; - missing|"") - ;; - *) - echo "Error: ROCM_DEP_CACHE_EXPORT_MODE must be one of: missing, always, never" - exit 1 - ;; - esac - - if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then - nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci") - if dependency_cache_ref_exists "${nixl_ref}"; then - echo "NIXL dependency cache exists: ${nixl_ref}" - else - echo "NIXL dependency cache missing; will seed: ${nixl_ref}" - add_dependency_cache_target "nixl-rocm-ci" - fi - fi - - if [[ "${mode}" != "always" && -n "${DEEPEP_CACHE_KEY:-}" ]]; then - deepep_ref=$(dependency_cache_ref_for_target "deepep-rocm-ci") - if dependency_cache_ref_exists "${deepep_ref}"; then - echo "DeepEP dependency cache exists: ${deepep_ref}" - else - echo "DeepEP dependency cache missing; will seed: ${deepep_ref}" - add_dependency_cache_target "deepep-rocm-ci" - fi - fi - - BAKE_TARGETS=("ci-base-rocm-ci") - if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then - echo "All dependency caches exist; building ci_base without dependency cache exports" - else - echo "Resolved dependency cache seed targets: ${DEPENDENCY_CACHE_TARGETS[*]}" - echo "Resolved ci_base bake targets: ${BAKE_TARGETS[*]}" - fi -} - -bake_config_targets() { - printf '%s\n' "${DEPENDENCY_CACHE_TARGETS[@]}" "${BAKE_TARGETS[@]}" \ - | awk 'NF && !seen[$0]++' -} - -print_bake_config() { - local -a print_targets=() - - echo "--- :page_facing_up: Resolved bake configuration" - mapfile -t print_targets < <(bake_config_targets) - docker buildx bake "${BAKE_ALLOW_ARGS[@]}" \ - "${BAKE_FILES[@]}" --print "${print_targets[@]}" | tee "${BAKE_CONFIG_FILE}" - - if command -v buildkite-agent >/dev/null 2>&1 && [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then - buildkite-agent artifact upload "${BAKE_CONFIG_FILE}" || true - echo "Uploaded ${BAKE_CONFIG_FILE} as Buildkite artifact" - else - echo "Saved bake config to ${BAKE_CONFIG_FILE} (not in Buildkite, skipping upload)" - fi -} - -confirm_remote_image_push() { - local image_ref="$1" - - if [[ -z "${CI_BASE_CONTENT_HASH:-}" ]]; then - remote_image_exists "${image_ref}" - return - fi - - if remote_ci_base_identity_is_current_with_retry "${image_ref}"; then - return 0 - fi - - echo "Remote image does not have the expected complete ci_base identity." - return 1 -} - -verify_dependency_cache_ref() { - local cache_ref="$1" - local attempts="${ROCM_DEP_CACHE_VERIFY_ATTEMPTS:-6}" - local delay_secs="${ROCM_DEP_CACHE_VERIFY_DELAY:-5}" - local attempt - - for ((attempt = 1; attempt <= attempts; attempt++)); do - if dependency_cache_ref_exists "${cache_ref}"; then - echo "Dependency cache confirmed: ${cache_ref}" - return 0 - fi - if [[ ${attempt} -lt ${attempts} ]]; then - echo "Dependency cache not visible yet (${attempt}/${attempts}): ${cache_ref}" - sleep "${delay_secs}" - fi - done - - echo "ERROR: dependency cache was not confirmed after upload: ${cache_ref}" - return 1 -} - -seed_dependency_caches_if_needed() { - local target="" - local cache_ref="" - - if [[ "${TARGET}" != "ci-base-rocm-ci-with-deps" ]]; then - return 0 - fi - if [[ ${#DEPENDENCY_CACHE_TARGETS[@]} -eq 0 ]]; then - return 0 - fi - - echo "--- :docker: Seeding ROCm dependency caches" - echo "Dependency cache uploads are required for this build." - echo "Seeding serially to avoid concurrent Docker Hub cache exporters." - - for target in "${DEPENDENCY_CACHE_TARGETS[@]}"; do - cache_ref=$(dependency_cache_ref_for_target "${target}") - if [[ -z "${cache_ref}" ]]; then - echo "ERROR: could not resolve dependency cache ref for ${target}" - return 1 - fi - - echo "--- :docker: Seeding ${target}" - echo "Expected cache ref: ${cache_ref}" - docker buildx bake \ - "${BAKE_ALLOW_ARGS[@]}" \ - "${BAKE_FILES[@]}" \ - --progress "${BUILDKIT_PROGRESS:-plain}" \ - "${target}" - verify_dependency_cache_ref "${cache_ref}" - done -} - -run_bake() { - local confirmation_ref="${IMAGE_TAG:-}" - - if is_ci_base_target && [[ -n "${CI_BASE_IMAGE_TAG_BUILD_REF:-}" ]]; then - confirmation_ref="${CI_BASE_IMAGE_TAG_BUILD_REF}" - fi - - echo "--- :docker: Building ${TARGET}" - docker buildx bake \ - "${BAKE_ALLOW_ARGS[@]}" \ - "${BAKE_FILES[@]}" \ - --progress "${BUILDKIT_PROGRESS:-plain}" \ - "${BAKE_TARGETS[@]}" - - if is_ci_base_target; then - if ! confirm_remote_image_push "${confirmation_ref}"; then - echo "Fresh ci_base image failed identity validation: ${confirmation_ref}" >&2 - return 1 - fi - refresh_ci_base_tags_from_ref "${confirmation_ref}" - fi - echo "--- :white_check_mark: Build complete" -} - -upload_wheel_artifacts_if_present() { - local wheel_dir="./wheel-export" - local artifact_dir="artifacts/vllm-rocm-install" - local archive_name="vllm-rocm-install.tar.gz" - local metadata_dir="${wheel_dir}/.vllm-ci-artifact" - local build_base_digest="" - local expected_native_base_image="" - local native_base_image="" - local native_base_digest="" - local whl="" - local whl_name="" - local -a wheels=() - - if ! should_upload_wheel_artifacts; then - return 0 - fi - - if [[ -d "${wheel_dir}" ]]; then - mapfile -t wheels < <(find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print) - fi - if [[ ${#wheels[@]} -ne 1 ]]; then - echo "Expected exactly one ROCm wheel in ${wheel_dir}; found ${#wheels[@]}" >&2 - return 1 - fi - whl="${wheels[0]}" - whl_name=$(basename "${whl}") - native_base_image="${CI_BASE_IMAGE_TAG_BUILD_REF:-${CI_BASE_IMAGE:-}}" - if [[ -z "${native_base_image}" ]]; then - echo "Native ROCm artifact requires a ci_base image reference" >&2 - return 1 - fi - if [[ "${BUILDKITE:-false}" == "true" ]]; then - expected_native_base_image=$(ci_base_tag_with_suffix \ - "${CI_BASE_IMAGE_TAG:-rocm/vllm-dev:ci_base}" \ - "build-${BUILDKITE_BUILD_ID:-}") - if [[ "${native_base_image}" != "${expected_native_base_image}" ]]; then - echo "Native ROCm artifact requires the exact ci_base build handoff: ${native_base_image}" >&2 - return 1 - fi - if [[ ! "${CI_BASE_IMAGE:-}" =~ @sha256:[0-9a-f]{64}$ ]]; then - echo "ROCm artifact build base must be digest-pinned: ${CI_BASE_IMAGE:-}" >&2 - return 1 - fi - build_base_digest="${CI_BASE_IMAGE##*@}" - if ! native_base_digest=$(resolve_image_digest "${native_base_image}"); then - echo "Could not resolve native ci_base handoff: ${native_base_image}" >&2 - return 1 - fi - if [[ "${native_base_digest}" != "${build_base_digest}" ]]; then - echo "Native ci_base handoff does not match the artifact build base" >&2 - echo " native: ${native_base_image}@${native_base_digest}" >&2 - echo " build: ${CI_BASE_IMAGE}" >&2 - return 1 - fi - fi - - echo "--- :package: Uploading ROCm vLLM install artifact" - rm -rf "${artifact_dir}" "${metadata_dir}" - mkdir -p "${artifact_dir}" "${metadata_dir}" - - printf '%s\n' "${BUILDKITE_COMMIT:-local}" > "${metadata_dir}/commit.txt" - printf '%s\n' "${native_base_image}" > "${metadata_dir}/native-base-image.txt" - printf '%s\n' "${CI_BASE_IMAGE:-}" > "${metadata_dir}/ci-base-image.txt" - printf '%s\n' "${IMAGE_TAG:-}" > "${metadata_dir}/fallback-image.txt" - printf '%s\n' "${whl_name}" > "${metadata_dir}/wheel-filename.txt" - - tar -C "${wheel_dir}" -czf "${artifact_dir}/${archive_name}" . - ( - cd "${artifact_dir}" - sha256sum "${archive_name}" > "${archive_name}.sha256" - ) - echo "Created ${archive_name}: $(du -sh "${artifact_dir}/${archive_name}" | cut -f1)" - cp "${metadata_dir}"/*.txt "${artifact_dir}/" - cp "${whl}" "${artifact_dir}/${whl_name}" - echo "Copied ${whl_name}: $(du -sh "${artifact_dir}/${whl_name}" | cut -f1)" - - if command -v buildkite-agent >/dev/null 2>&1; then - buildkite-agent artifact upload "${artifact_dir}/*" || return 1 - echo "ROCm vLLM install artifacts uploaded to ${artifact_dir}/" - elif [[ "${BUILDKITE:-false}" == "true" ]]; then - echo "buildkite-agent not found; cannot upload required ROCm artifacts" >&2 - return 1 - else - echo "Not in Buildkite, skipping artifact upload" - fi - - rm -rf "${wheel_dir}" -} - -main() { - init_config "$@" - configure_ci_base_write_scope - print_header - validate_inputs - load_ci_hcl - init_bake_files - if is_ci_base_target; then - prepare_ci_build_context - configure_custom_rocm_stages - fi - compute_ci_base_hash_if_needed - configure_ci_base_image_refs - maybe_skip_existing_image - setup_builder - prepare_git_cache_metadata - # Non-ci_base builds may deepen a shallow checkout above. Derive archival - # version metadata only after that lookup sees the available tag history. - if ! is_ci_base_target; then - prepare_ci_build_context - configure_custom_rocm_stages - fi - extract_dependency_pins - write_rocm_build_arg_override - compute_dependency_cache_keys - write_ci_base_label_override - compute_rocm_csrc_content_hash_if_needed - compute_rocm_rust_content_hash_if_needed - write_rocm_cache_override - # Keep the context override last so every bake target uses the owned tree. - write_build_context_override - resolve_ci_base_dependency_targets - print_bake_config - if [[ "${BAKE_PRINT_ONLY:-0}" == "1" ]]; then - echo "BAKE_PRINT_ONLY=1 set; skipping build" - return 0 - fi - if should_upload_wheel_artifacts; then - # wheel-export is an output directory, not a BuildKit cache. Starting - # clean prevents a failed/retried export from packaging a stale wheel. - rm -rf ./wheel-export - fi - if should_export_rocm_smoke; then - # The marker is a build output, not a cache. Never accept stale output - # from an earlier build or retry. - rm -rf ./build/rocm-smoke-export - fi - seed_dependency_caches_if_needed - run_bake - verify_rocm_smoke_export - promote_stable_ci_base_tag - publish_ci_base_handoff_ref - upload_wheel_artifacts_if_present -} - -if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then - main "$@" -fi diff --git a/.buildkite/scripts/hardware_ci/amd_ci_fast_exit.py b/.buildkite/scripts/hardware_ci/amd_ci_fast_exit.py new file mode 100644 index 000000000000..c77a19fbc966 --- /dev/null +++ b/.buildkite/scripts/hardware_ci/amd_ci_fast_exit.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""pytest plugin for AMD CI jobs: leave without native teardown once pytest is done. + +A pytest process can crash in native library destructors after its summary: +vllm/ci #96786 printed "236 passed" and then died with SIGSEGV (exit 139) during +interpreter shutdown, past the point where faulthandler reports. Once the outermost +session has finished, this runs the Python exit handlers (engine shutdown, Ray, +multiprocessing cleanup), flushes, and exits with pytest's own status, skipping +module teardown and native destructors. pytest's result is never changed, and a +command chain still stops on a real failure. + +Off when a kernel recorder is loaded (it writes its trace during native teardown) +or with VLLM_CI_FAST_EXIT=0. +""" + +import atexit +import os +import sys + +import pytest + +_outer = None +_session = None + + +def _recorder_loaded() -> bool: + return bool( + os.environ.get("ROCP_TOOL_LIBRARIES") + or os.environ.get("CUDA_INJECTION64_PATH") + or "rocprofiler" in os.environ.get("LD_PRELOAD", "") + ) + + +def pytest_configure(config): + global _outer + if _outer is None: # pytester runs nested sessions in-process; act on the outer one + _outer = config + + +def pytest_sessionfinish(session, exitstatus): + global _session + if session.config is _outer: + _session = session + + +@pytest.hookimpl(trylast=True) +def pytest_unconfigure(config): + if config is not _outer or _session is None: + return + if os.environ.get("VLLM_CI_FAST_EXIT") == "0" or _recorder_loaded(): + return + status = int(_session.exitstatus) + sys.stdout.flush() + sys.stderr.flush() + try: + atexit._run_exitfuncs() + except BaseException as exc: # an exit handler failing must not mask the result + print(f"[amd-ci] exit handler failed: {exc!r}", file=sys.stderr) + print( + f"[amd-ci] pytest finished with exit {status}; leaving without native teardown", + file=sys.stderr, + flush=True, + ) + sys.stdout.flush() + os._exit(status) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 59240b3d0a36..b649cfd7c847 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -1,1813 +1,212 @@ #!/bin/bash - -# shellcheck disable=SC2329 # Every function in this file is only ever reached -# transitively through handle_amd_runner_exit, the EXIT trap handler - shellcheck's -# unused-function check doesn't do full reachability analysis from an indirectly -# invoked entry point, so it flags the whole diagnostic-collection call chain as -# dead code even though it runs live on every nonzero exit. - -# This script runs ROCm tests either directly in a native CI pod or inside the -# corresponding Docker container. Multi-node tests continue to use Docker. -# -# Multi-node detection: Instead of matching on fragile group names, we detect -# multi-node jobs structurally by looking for the bracket command syntax -# "[node0_cmds] && [node1_cmds]" or via the NUM_NODES environment variable. -# -############################################################################### -# QUOTING / COMMAND PASSING -# -# Passing commands as positional arguments ($*) is fragile when the command -# string itself contains double quotes, e.g.: -# -# bash run-amd-test.sh "export FLAGS="value" && pytest -m "not slow"" -# -# The outer shell resolves the nested quotes *before* this script runs, so -# the script receives mangled input it cannot fully recover. -# -# Preferred: pass commands via the VLLM_TEST_COMMANDS environment variable: +# Run one AMD CI job in its Kubernetes pod. # -# export VLLM_TEST_COMMANDS='export FLAGS="value" && pytest -m "not slow"' -# bash run-amd-test.sh -# -# Single-quoted assignment preserves all inner double quotes verbatim. -# The $* path is kept for backward compatibility but callers should migrate. -############################################################################### +# The container image ($ROCM_CI_REPO:build-) already contains vLLM, +# all test dependencies, and /vllm-workspace. This script: +# 1. Sets up per-job local caches and validates the HF mount. +# 2. Fetches this build's wheel artifact for python_only_compile.sh. +# 3. Acquires the DPX GPU partition lock when running on a DPX agent. +# 4. Runs a preflight check (imports + GPU count). +# 5. Executes $VLLM_TEST_COMMANDS under a wall-clock timeout. +# 6. Writes a diagnostics bundle on failure. set -o pipefail -# shellcheck source=.buildkite/scripts/rocm/build-config.sh -source "$(dirname "${BASH_SOURCE[0]}")/../rocm/build-config.sh" || exit $? - -: "${BUILDKIT_PROGRESS:=plain}" -: "${TERM:=xterm-256color}" -: "${FORCE_COLOR:=1}" -: "${CLICOLOR_FORCE:=1}" -: "${PY_COLORS:=1}" -: "${ROCM_DOCKER_TTY:=1}" -: "${PYTHONFAULTHANDLER:=1}" -: "${PYTEST_TIMEOUT:=2400}" - -# Preserve runner-only values before clear_ci_orchestration_env removes them -# from the environment inherited by the test process. -amd_diagnostics_expected_gpu_count="${VLLM_CI_EXPECTED_GPU_COUNT:-1}" -amd_diagnostics_dir="${VLLM_CI_DIAGNOSTICS_DIR:-artifacts/amd-gpu-diagnostics}" -amd_diagnostics_checkout_root="${BUILDKITE_BUILD_CHECKOUT_PATH:-$(pwd -P)}" -amd_diagnostics_execution_mode="${VLLM_CI_EXECUTION_MODE:-single-node}" -amd_diagnostics_test_group="${VLLM_TEST_GROUP_NAME:-${BUILDKITE_STEP_KEY:-unknown}}" -amd_diagnostics_workspace="${VLLM_CI_WORKSPACE:-/vllm-workspace}" -amd_diagnostics_pod_name="${VLLM_CI_K8S_POD_NAME:-${POD_NAME:-${HOSTNAME:-}}}" -amd_diagnostics_k8s_namespace="${VLLM_CI_K8S_NAMESPACE:-${POD_NAMESPACE:-unknown}}" -amd_diagnostics_k8s_node_name="${VLLM_CI_K8S_NODE_NAME:-${NODE_NAME:-}}" -amd_diagnostics_collected=0 -amd_diagnostics_memory_events_path="" -amd_diagnostics_probe_budget_seconds=25 -amd_diagnostics_command_timeout_seconds=5 -amd_diagnostics_upload_timeout_seconds=20 -if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes" -fi -if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations="* ]]; then - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations=25" -fi -if [[ " ${PYTEST_ADDOPTS:-} " != *" --durations-min="* ]]; then - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--durations-min=1.0" -fi -# Dump stacks after 25 minutes, then stop an individual test after 40 minutes. -if [[ " ${PYTEST_ADDOPTS:-} " != *" faulthandler_timeout="* ]]; then - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }-o faulthandler_timeout=1500" -fi -if [[ " ${PYTEST_ADDOPTS:-} " != *" --timeout-method="* && - " ${PYTEST_ADDOPTS:-} " != *" --timeout-method "* ]]; then - PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--timeout-method=thread" -fi -export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS PYTEST_TIMEOUT ROCM_DOCKER_TTY -export PYTHONFAULTHANDLER - -# Export Python path for commands that run directly on the host. Containerized -# tests set this to /vllm-workspace below so spawned Python processes do not -# depend on their current working directory. -export PYTHONPATH="${PYTHONPATH:-..}" - -ci_started_at=$SECONDS - -############################################################################### -# Helper Functions -############################################################################### - -report_docker_usage() { - echo "--- Docker usage" - docker system df || true -} - -clear_ci_orchestration_env() { - unset -v \ - VLLM_TEST_GROUP_NAME \ - VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE \ - VLLM_CI_ARTIFACT_STEP \ - VLLM_TEST_CACHE \ - VLLM_CI_EXECUTION_MODE \ - VLLM_CI_WORKSPACE \ - VLLM_CI_REQUIRE_WORKSPACE_MOUNT \ - VLLM_TEST_COMMANDS \ - VLLM_CI_BRANCH \ - CI_ROCM_DOCKERFILE_BASE \ - CI_ROCM_DOCKERFILE \ - ROCM_BASE_DOCKERFILE \ - CI_BASE_DOCKERFILE \ - VLLM_CI_BASE_IMAGE \ - VLLM_CI_FALLBACK_IMAGE \ - VLLM_CI_DOCKER_DISABLED \ - VLLM_CI_DIAGNOSTICS_DIR \ - VLLM_CI_ARTIFACT_GLOB \ - VLLM_CI_ARTIFACT_CHECKSUM_GLOB \ - VLLM_CI_EXPECTED_GPU_COUNT \ - VLLM_CI_K8S_POD_NAME \ - VLLM_CI_K8S_NAMESPACE \ - VLLM_CI_K8S_NODE_NAME \ - VLLM_CI_USE_ARTIFACTS \ - VLLM_CI_RESULTS_ROOT \ - VLLM_ALLOW_DEPRECATED_BEAM_SEARCH -} - -amd_ci_teardown_log() { - local event=$1 - shift - - printf '[amd-ci-teardown] event=%s' "${event}" >&2 - if (($#)); then - printf ' %s' "$@" >&2 - fi - printf '\n' >&2 -} - -run_docker_with_ci_timeout() { - local raw_timeout=${CONTAINER_TIMEOUT_S:-0} - local configured_timeout=0 - local elapsed=0 - local remaining_timeout=0 - local status=0 - - if [[ ! "${raw_timeout}" =~ ^(0|[1-9][0-9]{0,5})$ ]]; then - amd_ci_teardown_log configuration_error \ - "reason=invalid_timeout" "expected=integer_0_to_604800" - return 2 - fi - configured_timeout=$((10#${raw_timeout})) - if ((configured_timeout > 604800)); then - amd_ci_teardown_log configuration_error \ - "reason=invalid_timeout" "expected=integer_0_to_604800" - return 2 - fi - - remaining_timeout=${configured_timeout} - if ((configured_timeout > 0)); then - if ! command -v timeout >/dev/null 2>&1; then - amd_ci_teardown_log configuration_error \ - "reason=timeout_command_not_found" - return 127 - fi - elapsed=$((SECONDS - ci_started_at)) - if ((elapsed >= configured_timeout)); then - amd_ci_teardown_log deadline_exhausted \ - "mode=docker" "configured_s=${configured_timeout}" "elapsed_s=${elapsed}" - return 124 - fi - remaining_timeout=$((configured_timeout - elapsed)) - fi - - amd_ci_teardown_log workload_started \ - "mode=docker" "timeout_s=${remaining_timeout}" - if ((remaining_timeout == 0)); then - "$@" || status=$? - else - # Docker runs interactively, so keep it in the foreground process group. - # --verbose records both TERM and any KILL escalation caused by timeout. - timeout --verbose --foreground --signal=TERM --kill-after=10s \ - "${remaining_timeout}s" "$@" || status=$? - fi - amd_ci_teardown_log workload_finished "mode=docker" "status=${status}" - return "${status}" -} - -prepare_artifact_image() { - if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then - return 1 - fi - if ! command -v buildkite-agent >/dev/null 2>&1; then - echo "buildkite-agent not found; cannot download ROCm wheel artifact" - return 1 - fi - - local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}" - local archive="" - local metadata_file="" - local base_image="${VLLM_CI_BASE_IMAGE:-rocm/vllm-dev:ci_base}" - local artifact_image="" - local artifact_key="" - local base_digest="" - local wheel_dir="" - local context_dir="" - local workspace_dir="" - - artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) - wheel_dir="${artifact_work_dir}/wheels" - context_dir="${artifact_work_dir}/context" - workspace_dir="${context_dir}/workspace" - mkdir -p "${wheel_dir}" "${context_dir}/wheels" "${workspace_dir}" - - echo "--- Downloading ROCm wheel artifact" - if ! buildkite-agent artifact download "${artifact_glob}" "${artifact_work_dir}"; then - echo "Failed to download ${artifact_glob}" - return 1 - fi - buildkite-agent artifact download \ - "artifacts/vllm-rocm-install/ci-base-image.txt" \ - "${artifact_work_dir}" >/dev/null 2>&1 || true - - archive=$(find "${artifact_work_dir}" -name "vllm-rocm-install.tar.gz" -type f | head -1) - if [[ -z "${archive}" || ! -f "${archive}" ]]; then - echo "ROCm wheel artifact archive was not found" - return 1 - fi - - metadata_file=$(find "${artifact_work_dir}" -name "ci-base-image.txt" -type f | head -1) - if [[ -n "${metadata_file}" && -s "${metadata_file}" ]]; then - base_image=$(tr -d '[:space:]' < "${metadata_file}") - elif using_custom_rocm_dockerfiles; then - echo "Custom ROCm ci_base metadata is missing; using the full CI image" - return 1 - fi - - echo "--- Preparing local ROCm test image" - echo "Base image: ${base_image}" - docker pull "${base_image}" || return 1 - base_digest=$( - docker image inspect \ - --format='{{if .RepoDigests}}{{index .RepoDigests 0}}{{else}}{{.Id}}{{end}}' \ - "${base_image}" 2>/dev/null || printf '%s' "${base_image}" - ) - - artifact_key=$( - { - printf 'base-image:%s\n' "${base_digest}" - sha256sum "${archive}" - } | sha256sum | cut -c1-24 - ) - artifact_image="rocm/vllm-ci-artifact:${artifact_key}" - - if docker image inspect "${artifact_image}" >/dev/null 2>&1; then - echo "Using existing local ROCm artifact image: ${artifact_image}" - image_name="${artifact_image}" - return 0 - fi - - tar -xzf "${archive}" -C "${wheel_dir}" || return 1 - if ! ls "${wheel_dir}"/*.whl >/dev/null 2>&1; then - echo "ROCm wheel artifact did not contain a wheel" - return 1 - fi - if [[ ! -d "${wheel_dir}/tests" ]]; then - echo "ROCm wheel artifact did not contain the test workspace" - return 1 - fi - - cp "${wheel_dir}"/*.whl "${context_dir}/wheels/" || return 1 - tar -C "${wheel_dir}" \ - --exclude='*.whl' \ - --exclude='.vllm-ci-artifact' \ - --exclude='./.vllm-ci-artifact' \ - -cf - . \ - | tar -C "${workspace_dir}" -xf - || return 1 - cat > "${context_dir}/Dockerfile" <<'EOF' -ARG BASE_IMAGE -FROM ${BASE_IMAGE} -COPY wheels/ /tmp/vllm-wheels/ -COPY workspace/ /vllm-workspace/ -RUN python3 -m pip install --no-deps --force-reinstall /tmp/vllm-wheels/*.whl \ - && rm -rf /tmp/vllm-wheels -WORKDIR /vllm-workspace -EOF - - echo "--- Building local ROCm test image" - docker build \ - --pull=false \ - --progress "${BUILDKIT_PROGRESS}" \ - --build-arg "BASE_IMAGE=${base_image}" \ - -t "${artifact_image}" \ - "${context_dir}" || return 1 - image_name="${artifact_image}" - return 0 -} - -is_native_runtime() { - [[ "${AMD_CI_RUNTIME:-}" == "native" || "${NATIVE_CI:-}" == "true" ]] -} - -validate_native_workspace() { - local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}" - local workspace_real="" - local checkout_real="" - local workspace_mount="" - - mkdir -p "${workspace_dir}" || return 1 - workspace_real=$(readlink -m "${workspace_dir}") || return 1 - if [[ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ]]; then - checkout_real=$(readlink -m "${BUILDKITE_BUILD_CHECKOUT_PATH}") || return 1 - if [[ "${checkout_real}" == "${workspace_real}" \ - || "${checkout_real}" == "${workspace_real}/"* \ - || "${workspace_real}" == "${checkout_real}/"* ]]; then - echo "Refusing to replace ${workspace_real}; it overlaps the Buildkite checkout ${checkout_real}" >&2 - return 1 - fi - fi - if [[ "${VLLM_CI_REQUIRE_WORKSPACE_MOUNT:-1}" == "1" ]]; then - if ! command -v findmnt >/dev/null 2>&1; then - echo "findmnt is required to verify the native workspace mount" >&2 - return 1 - fi - workspace_mount=$(findmnt -n -T "${workspace_real}" -o TARGET 2>/dev/null || true) - if [[ "$(readlink -m "${workspace_mount:-/}")" != "${workspace_real}" ]]; then - echo "Native CI requires a dedicated volume mounted at ${workspace_real}" >&2 - return 1 - fi - fi -} - -prepare_native_workspace() { - local test_commands="${1:-}" - - if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then - echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1" - return 1 - fi - if ! command -v buildkite-agent >/dev/null 2>&1; then - echo "buildkite-agent not found; cannot download ROCm wheel artifact" - return 1 - fi - validate_native_workspace || return 1 - - local artifact_glob="${VLLM_CI_ARTIFACT_GLOB:-artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz}" - local artifact_checksum_glob="${VLLM_CI_ARTIFACT_CHECKSUM_GLOB:-${artifact_glob}.sha256}" - local artifact_step="${VLLM_CI_ARTIFACT_STEP:-image-build-amd}" - local archive="" - local checksum="" - local download_dir="" - local metadata_dir="" - local recorded_base="" - local recorded_commit="" - local recorded_wheel="" - local checkout="" - local checkout_commit="" - local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}" - local wheel_dir="" - local attempt=0 - local attempt_dir="" - local -a archives=() - local -a checksums=() - local -a wheels=() - - artifact_work_dir=$(mktemp -d -t vllm-rocm-artifact.XXXXXX) || return 1 - wheel_dir="${artifact_work_dir}/wheels" - mkdir -p "${wheel_dir}" || return 1 - - echo "--- Downloading ROCm wheel artifact from ${artifact_step} (native in-pod)" - for attempt in 1 2 3; do - attempt_dir="${artifact_work_dir}/download-${attempt}" - rm -rf "${attempt_dir}" || return 1 - mkdir -p "${attempt_dir}" || return 1 - if buildkite-agent artifact download \ - "${artifact_glob}" "${attempt_dir}" --step "${artifact_step}" \ - && buildkite-agent artifact download \ - "${artifact_checksum_glob}" "${attempt_dir}" --step "${artifact_step}"; then - download_dir="${attempt_dir}" - break - fi - echo "Artifact download attempt ${attempt}/3 failed" - if [[ "${attempt}" -lt 3 ]]; then - sleep $((attempt * 2)) - fi - done - if [[ -z "${download_dir}" ]]; then - echo "Failed to download ${artifact_glob} and ${artifact_checksum_glob} from ${artifact_step}" - return 1 - fi - - mapfile -t archives < <( - find "${download_dir}" -name "vllm-rocm-install.tar.gz" -type f -print - ) - mapfile -t checksums < <( - find "${download_dir}" -name "vllm-rocm-install.tar.gz.sha256" -type f -print - ) - if [[ ${#archives[@]} -ne 1 || ${#checksums[@]} -ne 1 ]]; then - echo "Expected exactly one ROCm archive and checksum; found ${#archives[@]} archive(s) and ${#checksums[@]} checksum(s)" >&2 - return 1 - fi - archive="${archives[0]}" - checksum="${checksums[0]}" - if [[ "$(dirname "${archive}")" != "$(dirname "${checksum}")" ]]; then - echo "ROCm archive and checksum were downloaded to different directories" >&2 - return 1 - fi - ( - cd "$(dirname "${archive}")" - sha256sum -c "$(basename "${checksum}")" - ) || return 1 - - tar --no-same-owner -xzf "${archive}" -C "${wheel_dir}" || return 1 - mapfile -t wheels < <( - find "${wheel_dir}" -maxdepth 1 -type f -name '*.whl' -print - ) - if [[ ${#wheels[@]} -ne 1 ]]; then - echo "ROCm artifact must contain exactly one top-level wheel; found ${#wheels[@]}" >&2 - return 1 - fi - metadata_dir="${wheel_dir}/.vllm-ci-artifact" - for metadata_file in commit.txt native-base-image.txt wheel-filename.txt; do - if [[ ! -s "${metadata_dir}/${metadata_file}" ]]; then - echo "ROCm artifact metadata is missing ${metadata_file}" >&2 - return 1 - fi - done - for metadata_file in ci-base-image.txt fallback-image.txt; do - if [[ ! -f "${metadata_dir}/${metadata_file}" ]]; then - echo "ROCm artifact metadata is missing ${metadata_file}" >&2 - return 1 - fi - done - - recorded_commit=$(tr -d '\r\n' < "${metadata_dir}/commit.txt") - recorded_base=$(tr -d '\r\n' < "${metadata_dir}/native-base-image.txt") - recorded_wheel=$(tr -d '\r\n' < "${metadata_dir}/wheel-filename.txt") - if [[ -z "${BUILDKITE_COMMIT:-}" || "${recorded_commit}" != "${BUILDKITE_COMMIT}" ]]; then - echo "ROCm artifact commit ${recorded_commit} does not match ${BUILDKITE_COMMIT:-unset}" >&2 - return 1 - fi - if [[ -z "${BUILDKITE_BUILD_ID:-}" \ - || "${recorded_base}" != *":ci_base-build-${BUILDKITE_BUILD_ID}" ]]; then - echo "ROCm artifact base is not scoped to Buildkite build ${BUILDKITE_BUILD_ID:-unset}: ${recorded_base}" >&2 - return 1 - fi - if [[ -z "${VLLM_CI_BASE_IMAGE:-}" || "${recorded_base}" != "${VLLM_CI_BASE_IMAGE}" ]]; then - echo "ROCm artifact base ${recorded_base} does not match ${VLLM_CI_BASE_IMAGE:-unset}" >&2 - return 1 - fi - if [[ "${recorded_wheel}" != "$(basename "${wheels[0]}")" ]]; then - echo "ROCm artifact wheel manifest ${recorded_wheel} does not match $(basename "${wheels[0]}")" >&2 - return 1 - fi - for required_dir in tests .buildkite requirements; do - if [[ ! -d "${wheel_dir}/${required_dir}" ]]; then - echo "ROCm wheel artifact did not contain ${required_dir}/" >&2 - return 1 - fi +# --------------------------------------------------------------------------- +# Inputs — set by the pipeline / ci-infra template +# --------------------------------------------------------------------------- +commands=${VLLM_TEST_COMMANDS:?VLLM_TEST_COMMANDS must be set by the pipeline} +gpus=${VLLM_CI_EXPECTED_GPU_COUNT:-1} +[[ $gpus =~ ^(0|[1-9][0-9]*)$ ]] || { echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=$gpus" >&2; exit 1; } +job=${BUILDKITE_JOB_ID:-local} +job=${job//[^A-Za-z0-9_.-]/_} +# Short: Ray puts AF_UNIX sockets under $TMPDIR, whose paths may not exceed 107 bytes. +root=/tmp/vllm-${job##*-} +root=${root:0:22} + +# --------------------------------------------------------------------------- +# Cache directories — all under $root so they are job-local and isolated. +# HF hub is the one exception: it lives on the node's persistent NFS mount. +# --------------------------------------------------------------------------- +export HF_HOME=${HF_HOME:-/home/buildkite-agent/huggingface} +export TIKTOKEN_RS_CACHE_DIR=$HF_HOME/tiktoken-rs-cache +export PYTHONPATH=/vllm-workspace +export PYTORCH_ROCM_ARCH="" # let vLLM detect arches at runtime +export TMPDIR=$root/tmp +export VLLM_RPC_BASE_PATH=/tmp +export TORCHINDUCTOR_CACHE_DIR=$root/inductor +export TRITON_CACHE_DIR=$root/triton +export VLLM_CACHE_ROOT=$root/vllm +export XDG_CACHE_HOME=$root/xdg +export HF_DATASETS_CACHE=$root/datasets +export HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT:-300} +export HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT:-60} + +mkdir -p \ + "$TMPDIR" "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR" \ + "$VLLM_CACHE_ROOT" "$XDG_CACHE_HOME" "$HF_DATASETS_CACHE" \ + "$TIKTOKEN_RS_CACHE_DIR" || exit 1 + +# HF_HOME must be a persistent mount, not rootfs (catches misconfigured nodes). +[[ $(findmnt -n -T "$HF_HOME" -o TARGET) != / ]] \ + || { echo "HF_HOME is on rootfs — it must be a persistent node mount" >&2; exit 1; } + +# hf-xet state corrupts on NFS; redirect it to a local tmpdir when needed. +if [[ $(findmnt -n -T "$HF_HOME" -o FSTYPE) == nfs* ]]; then + export HF_XET_CACHE=$root/hf-xet + export HF_XET_HIGH_PERFORMANCE=0 + export HF_XET_RECONSTRUCTION_USE_VECTORED_WRITE=0 +fi + +# --------------------------------------------------------------------------- +# Terminal and pytest settings +# --------------------------------------------------------------------------- +# Keep CI log output colorized. +: "${TERM:=xterm-256color}" "${FORCE_COLOR:=1}" "${CLICOLOR_FORCE:=1}" "${PY_COLORS:=1}" +export TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS + +# Fault-handler dumps stacks after 25 min; pytest-timeout kills a single test +# after 40 min (PYTEST_TIMEOUT=2400 s). Both are overridable by the pipeline. +export PYTHONFAULTHANDLER=1 +export PYTEST_TIMEOUT=${PYTEST_TIMEOUT:-2400} +export PYTEST_ADDOPTS="${PYTEST_ADDOPTS:-} \ + --color=yes \ + --durations=25 \ + --durations-min=1.0 \ + -o faulthandler_timeout=1500 \ + --timeout-method=thread" + +# After its summary, a pytest process leaves without native library teardown, where it +# has crashed with SIGSEGV after every test passed (amd_ci_fast_exit.py). +# Installed into site-packages, not put on PYTHONPATH: commands that set PYTHONPATH themselves +# still inherit PYTEST_ADDOPTS and must find the plugin. +ci_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +cp "$ci_dir/amd_ci_fast_exit.py" \ + "$(python3 -P -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')/" || exit 1 +export PYTEST_ADDOPTS="$PYTEST_ADDOPTS -p amd_ci_fast_exit" + +[[ $gpus == 0 ]] && export VLLM_TARGET_DEVICE=cpu + +# --------------------------------------------------------------------------- +# Wheel reinstall — python_only_compile.sh path +# --------------------------------------------------------------------------- +# python_only_compile.sh reinstalls vLLM around this build's wheel, which build.sh uploads as an +# artifact of the image build step (kept out of the image: only this test needs it). +if [[ $commands == *python_only_compile.sh* ]]; then + wheel_dir=$(mktemp -d) + for try in 1 2 3; do + buildkite-agent artifact download ".rocm-wheel/vllm-*.whl" "$wheel_dir" --step image-build-amd \ + && break + ((try == 3)) && exit 1 + sleep $((try * 5)) done - - echo "--- Installing ROCm wheel into pod environment" - python3 -m pip install --no-deps --force-reinstall "${wheels[0]}" || return 1 - - echo "--- Preparing ${workspace_dir} from artifact" - find "${workspace_dir}" -mindepth 1 -maxdepth 1 -exec rm -rf -- {} + || return 1 - tar -C "${wheel_dir}" \ - --exclude='*.whl' \ - --exclude='.vllm-ci-artifact' \ - --exclude='./.vllm-ci-artifact' \ - -cf - . | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1 - if [[ ! -d "${workspace_dir}/tests" ]]; then - echo "Failed to stage the native test workspace" >&2 - return 1 - fi - - # The ROCm artifact intentionally contains only the installed wheel and the - # test workspace. The Python-only compilation job also needs setup.py and the - # vllm source tree, so overlay the matching Buildkite checkout for that job. - if [[ "${test_commands}" == *python_only_compile.sh* ]]; then - checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}" - if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then - echo "Python-only native CI requires BUILDKITE_BUILD_CHECKOUT_PATH" >&2 - return 1 - fi - if ! git -c "safe.directory=${checkout}" -C "${checkout}" \ - rev-parse --is-inside-work-tree >/dev/null 2>&1; then - echo "Buildkite checkout is not a Git worktree: ${checkout}" >&2 - return 1 - fi - checkout_commit=$( - git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse HEAD - ) || return 1 - if [[ "${checkout_commit}" != "${recorded_commit}" ]]; then - echo "Buildkite checkout ${checkout_commit} does not match ROCm artifact ${recorded_commit}" >&2 - return 1 - fi - - # setup.py normally derives this from .git via setuptools-scm. The native - # source overlay deliberately excludes Git metadata, so preserve the exact - # version from the already installed, artifact-matched wheel. - VLLM_VERSION_OVERRIDE=$( - python3 -c 'import importlib.metadata as m; print(m.version("vllm"))' - ) || return 1 - export VLLM_VERSION_OVERRIDE - VLLM_PRECOMPILED_WHEEL_LOCATION="${wheels[0]}" - export VLLM_PRECOMPILED_WHEEL_LOCATION - echo "INFO: native Python-only wheel=${VLLM_PRECOMPILED_WHEEL_LOCATION}" - - echo "--- Overlaying full source checkout for Python-only compilation" - # Archive the verified commit instead of copying the worktree so dirty or - # untracked agent files cannot contaminate the artifact-matched workspace. - git -c "safe.directory=${checkout}" -C "${checkout}" \ - archive --format=tar "${recorded_commit}" \ - | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1 - for required_source in setup.py pyproject.toml vllm; do - if [[ ! -e "${workspace_dir}/${required_source}" ]]; then - echo "Full source checkout is missing ${required_source}" >&2 - return 1 - fi - done - fi - - return 0 -} - -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 - echo "Native ROCm CI currently requires the ci_base container to run as root" >&2 - return 1 - fi - - job_id="${job_id//[^A-Za-z0-9_.-]/_}" - job_id_suffix="${job_id##*-}" - job_id_suffix="${job_id_suffix:0:12}" - 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" - : "${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" - # openai-harmony downloads its tiktoken vocab on first use and caches it under - # $TMPDIR by default, which is job-local; keep it with the persistent Hub cache. - TIKTOKEN_RS_CACHE_DIR="${HF_HOME}/tiktoken-rs-cache" - : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" - : "${HF_HUB_ETAG_TIMEOUT:=60}" - if [[ "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" == "0" ]]; then - # CPU-only native jobs intentionally reuse the ROCm wheel. Make that target - # explicit so platform selection does not depend on wheel metadata. - VLLM_TARGET_DEVICE=cpu - export VLLM_TARGET_DEVICE - fi - export TMPDIR VLLM_RPC_BASE_PATH - export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME - export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT - export TIKTOKEN_RS_CACHE_DIR - export PYTORCH_ROCM_ARCH="" - - mkdir -p "${TMPDIR}" \ - "${TORCHINDUCTOR_CACHE_DIR}" \ - "${TRITON_CACHE_DIR}" \ - "${VLLM_CACHE_ROOT}" \ - "${XDG_CACHE_HOME}" \ - "${HF_HOME}" \ - "${TIKTOKEN_RS_CACHE_DIR}" \ - "${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 - echo "findmnt is required to verify the native Hugging Face cache mount" >&2 - return 1 - fi - hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true) - if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then - echo "Native CI requires a persistent volume mounted at or above ${HF_HOME}" >&2 - 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 -} - -check_dpx_gpu_exclusivity() { - local devices=(/dev/dri/renderD*) - local device_id lock_status - - if [[ "${#devices[@]}" -ne 1 || ! -c "${devices[0]}" ]]; then - echo "DPX guard requires exactly one render device; refusing to start tests." >&2 - return 1 - fi - # This queue mounts the same host-local HF cache into every pod on the node. - if ! mountpoint -q "${HF_HOME}"; then - echo "DPX guard requires the shared HF cache mount; refusing to start tests." >&2 - return 1 - fi - device_id=$(stat -Lc '%t-%T' "${devices[0]}") || return 1 - mkdir -p "${HF_HOME}/.vllm-dpx-locks" || return 1 - # Hold the descriptor through workload and teardown; never delete the file. - exec {dpx_gpu_lock_fd}>>"${HF_HOME}/.vllm-dpx-locks/${device_id}.lock" || return 1 - flock -n -E 75 "${dpx_gpu_lock_fd}" && return 0 - lock_status=$? - if [[ "${lock_status}" -eq 75 ]]; then - echo "DPX GPU collision: ${devices[0]} (${device_id}) is already locked by another CI job; refusing to start tests." >&2 - else - echo "DPX GPU lock failed (status ${lock_status}); refusing to start tests." >&2 - fi - return 1 + wheels=("$wheel_dir"/.rocm-wheel/vllm-*.whl) + [[ ${#wheels[@]} == 1 && -f ${wheels[0]} ]] \ + || { echo "Expected exactly one wheel artifact, found: ${wheels[*]}" >&2; exit 1; } + VLLM_PRECOMPILED_WHEEL_LOCATION=${wheels[0]} + VLLM_VERSION_OVERRIDE=$(python3 -P -c \ + 'import importlib.metadata as m; print(m.version("vllm"))') || exit 1 + export VLLM_TARGET_DEVICE=rocm VLLM_VERSION_OVERRIDE VLLM_PRECOMPILED_WHEEL_LOCATION +fi + +# --------------------------------------------------------------------------- +# DPX GPU partition lock +# --------------------------------------------------------------------------- +# DPX nodes expose exactly one GPU partition per pod. A host-wide lock file +# prevents a second job from landing on the same device simultaneously. +_acquire_dpx_lock() { + local dev=(/dev/dri/renderD*) + [[ ${#dev[@]} == 1 ]] \ + || { echo "DPX: expected exactly one render device, found ${#dev[@]}" >&2; return 1; } + mkdir -p "$HF_HOME/.vllm-dpx-locks" + exec 9>>"$HF_HOME/.vllm-dpx-locks/$(stat -Lc %t-%T "${dev[0]}").lock" + flock -n 9 || { echo "DPX GPU collision on ${dev[0]}" >&2; return 1; } +} + +if [[ ${BUILDKITE_AGENT_META_DATA_QUEUE:-} == *dpx* && $gpus != 0 ]]; then + _acquire_dpx_lock || exit 1 +fi + +# --------------------------------------------------------------------------- +# Lifecycle logging +# --------------------------------------------------------------------------- +# Structured events for infra dashboards keyed on [amd-ci-teardown]. +# Usage: _teardown_log [key=val ...] +_teardown_log() { + local event=$1; shift + printf '[amd-ci-teardown] event=%s' "$event" >&2 + (($#)) && printf ' %s' "$@" >&2 + printf '\n' >&2 } -run_native_preflight() { - local expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}" - - if [[ ! "${expected_gpus}" =~ ^[0-9]+$ ]]; then - echo "Invalid VLLM_CI_EXPECTED_GPU_COUNT=${expected_gpus}" >&2 - return 1 - fi - - python3 -c "import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing" || return 1 - - if [[ "${expected_gpus}" == "0" ]]; then - echo "Native CPU-only AMD job: skipping ROCm device validation" - return 0 - fi - +# --------------------------------------------------------------------------- +# Diagnostics bundle (always collected on failure) +# --------------------------------------------------------------------------- +# Writes to artifacts/amd-gpu-diagnostics//diagnostics.log, which the +# pipeline's artifact_paths uploads even when the step fails. +diagnostics() { + local d=artifacts/amd-gpu-diagnostics/$job + mkdir -p "$d" + { + date + echo "node=${VLLM_CI_K8S_NODE_NAME:-?} pod=${VLLM_CI_K8S_POD_NAME:-?}" + echo "--- amd-smi ---" + timeout 10 amd-smi list + echo "--- rocm-smi pids ---" + timeout 10 rocm-smi --showpids + echo "--- top processes by RSS ---" + ps -eo pid,etime,rss,args --sort=-rss | head -30 + echo "--- disk / memory ---" + df -h /tmp /dev/shm "$HF_HOME" + free -g + echo "--- devices ---" + ls -l /dev/kfd /dev/dri + } >"$d/diagnostics.log" 2>&1 +} + +# --------------------------------------------------------------------------- +# Preflight: verify imports and GPU visibility +# --------------------------------------------------------------------------- +# -P: the checkout (cwd) has a vllm/ source tree that would shadow the installed package. +python3 -P -c "import vllm, torch, transformers, ray" || exit 1 + +if [[ $gpus != 0 ]]; then echo "--- ROCm info" - rocminfo || return 1 - VLLM_CI_EXPECTED_GPU_COUNT="${expected_gpus}" python3 - <<'PY' -import os - + rocminfo || { diagnostics; exit 1; } + python3 - <> "${log_file}" - printf '%s\n' "${title}" >> "${log_file}" - printf '===============================================================================\n' \ - >> "${log_file}" -} - -append_failure_diagnostic_note() { - local log_file=$1 - shift - printf '%s\n' "$@" >> "${log_file}" -} - -run_failure_diagnostic() { - local log_file=$1 - local probe_deadline=$2 - shift 2 - local command_status=0 - local command_timeout_seconds=${amd_diagnostics_command_timeout_seconds} - local remaining_seconds=$((probe_deadline - SECONDS)) - - { - printf '\n$' - printf ' %q' "$@" - printf '\n' - } >> "${log_file}" - - if ! command -v timeout >/dev/null 2>&1; then - echo "[timeout unavailable; bounded diagnostic command skipped]" \ - >> "${log_file}" - return 0 - fi - if [[ "${remaining_seconds}" -le 0 ]]; then - echo "[diagnostic probe budget exhausted; command skipped]" \ - >> "${log_file}" - return 0 - fi - if [[ "${remaining_seconds}" -lt "${command_timeout_seconds}" ]]; then - command_timeout_seconds=${remaining_seconds} - fi - - timeout --kill-after=1s "${command_timeout_seconds}s" "$@" \ - >> "${log_file}" 2>&1 - command_status=$? - if [[ "${command_status}" -ne 0 ]]; then - echo "[diagnostic command exited with status ${command_status}]" \ - >> "${log_file}" - fi - return 0 -} - -upload_failure_diagnostics_artifact() { - local upload_root=$1 upload_path=$2 agent_bin="" - local agent_upload_help="" - local -a upload_options=() - - if [[ -n "${BUILDKITE_BIN_PATH:-}" \ - && -x "${BUILDKITE_BIN_PATH}/buildkite-agent" ]]; then - agent_bin="${BUILDKITE_BIN_PATH}/buildkite-agent" - else - agent_bin=$(command -v buildkite-agent 2>/dev/null || true) - fi - if [[ -z "${agent_bin}" && -x /workspace/buildkite-agent ]]; then - agent_bin=/workspace/buildkite-agent - fi - [[ -n "${agent_bin}" && -n "${BUILDKITE_JOB_ID:-}" ]] || return 1 - command -v timeout >/dev/null 2>&1 || return 1 - - # Older AMD runners predate the literal upload options. A restricted path - # has no glob or delimiter characters, so it is safe with either agent. - if [[ ! "${upload_path}" =~ ^[A-Za-z0-9_][A-Za-z0-9_.-]*(/[A-Za-z0-9_][A-Za-z0-9_.-]*)*$ ]]; then - agent_upload_help=$("${agent_bin}" artifact upload --help 2>&1 || true) - if [[ "${agent_upload_help}" != *"--literal"* \ - || "${agent_upload_help}" != *"--delimiter"* ]]; then - return 1 - fi - upload_options=(--literal --delimiter "") - fi - - ( - cd "${upload_root}" || exit 1 - BUILDKITE_AGENT_DEBUG=false \ - BUILDKITE_AGENT_DEBUG_HTTP=false \ - BUILDKITE_AGENT_TRACE_HTTP=false \ - BUILDKITE_AGENT_LOG_LEVEL=error \ - timeout --kill-after=2s "${amd_diagnostics_upload_timeout_seconds}s" \ - "${agent_bin}" artifact upload \ - "${upload_options[@]}" "./${upload_path}" - ) >/dev/null 2>&1 -} - -append_failure_diagnostic_file() { - local log_file=$1 - local label=$2 - local path=$3 - local command_status=0 - - printf '\n%s:\n' "${label}" >> "${log_file}" - printf '$ cat %q\n' "${path}" >> "${log_file}" - if [[ ! -r "${path}" ]]; then - echo "[file unavailable]" >> "${log_file}" - return 0 - fi - - cat "${path}" >> "${log_file}" 2>&1 - command_status=$? - if [[ "${command_status}" -ne 0 ]]; then - echo "[diagnostic file read exited with status ${command_status}]" \ - >> "${log_file}" - fi - return 0 -} - -decode_mountinfo_path() { - local value=$1 - - value=${value//\\040/ } - value=${value//\\011/$'\t'} - value=${value//\\012/$'\n'} - value=${value//\\134/\\} - printf '%s\n' "${value}" -} - -resolve_current_cgroup_v2_dir() { - local cgroup_path="" - local mount_root="" - local mount_point="" - local relative_path="" - local candidate="" - - cgroup_path=$(awk -F: '$2 == "" { print $3; exit }' \ - /proc/self/cgroup 2>/dev/null) - [[ -n "${cgroup_path}" ]] || return 1 - - while IFS=$'\t' read -r mount_root mount_point; do - mount_root=$(decode_mountinfo_path "${mount_root}") - mount_point=$(decode_mountinfo_path "${mount_point}") - if [[ "${cgroup_path}" == "/" ]]; then - relative_path="" - elif [[ "${mount_root}" == "/" ]]; then - relative_path="${cgroup_path}" - elif [[ "${cgroup_path}" == "${mount_root}" ]]; then - relative_path="" - elif [[ "${cgroup_path}" == "${mount_root}/"* ]]; then - relative_path="/${cgroup_path#"${mount_root}/"}" - else - continue - fi - - candidate="${mount_point%/}${relative_path}" - if [[ -d "${candidate}" ]]; then - printf '%s\n' "${candidate}" - return 0 - fi - done < <( - awk ' - { - for (i = 6; i <= NF; i++) { - if ($i == "-" && $(i + 1) == "cgroup2") { - print $4 "\t" $5 - break - } - } - } - ' /proc/self/mountinfo 2>/dev/null - ) - return 1 -} - -collect_cgroup_files() { - local log_file=$1 - local title=$2 - local cgroup_dir=$3 - shift 3 - local cgroup_file="" - local files_collected=0 - - append_failure_diagnostic_section "${log_file}" "${title}" - if [[ -z "${cgroup_dir}" || ! -d "${cgroup_dir}" ]]; then - append_failure_diagnostic_note "${log_file}" \ - "resolved_cgroup_path=unavailable" - return 0 - fi - - append_failure_diagnostic_note "${log_file}" \ - "resolved_cgroup_path=${cgroup_dir}" - for cgroup_file in "$@"; do - if [[ -r "${cgroup_dir}/${cgroup_file}" ]]; then - append_failure_diagnostic_file "${log_file}" "${cgroup_file}" \ - "${cgroup_dir}/${cgroup_file}" - files_collected=$((files_collected + 1)) - fi - done - if [[ "${files_collected}" -eq 0 ]]; then - append_failure_diagnostic_note "${log_file}" \ - "[no requested cgroup files were readable]" - fi -} - -collect_cgroup_diagnostics() { - local log_file=$1 - local cgroup_dir="" - - append_failure_diagnostic_section "${log_file}" "Cgroup membership" - append_failure_diagnostic_file "${log_file}" "current_process" \ - /proc/self/cgroup - append_failure_diagnostic_file "${log_file}" "pid_namespace_init_process" \ - /proc/1/cgroup - - cgroup_dir=$(resolve_current_cgroup_v2_dir 2>/dev/null || true) - collect_cgroup_files "${log_file}" \ - "Cgroup v2 current-process resource state" "${cgroup_dir}" \ - cgroup.type cgroup.events \ - memory.current memory.peak memory.max memory.high memory.low memory.min \ - memory.oom.group memory.swap.current memory.swap.max \ - memory.events memory.events.local memory.swap.events memory.stat \ - memory.pressure \ - cpu.max cpu.max.burst cpu.weight cpu.stat cpu.pressure \ - cpuset.cpus cpuset.cpus.effective cpuset.mems cpuset.mems.effective \ - pids.current pids.max pids.events pids.events.local \ - io.stat io.max io.weight io.pressure - if [[ -n "${cgroup_dir}" && -r "${cgroup_dir}/memory.events" ]]; then - amd_diagnostics_memory_events_path="${cgroup_dir}/memory.events" - fi -} - -collect_process_diagnostics() { - local log_file=$1 - local probe_deadline=$2 - local shell_proc_dir="/proc/$$" - - append_failure_diagnostic_section "${log_file}" \ - "Runner process constraints and PID-namespace-visible processes" - append_failure_diagnostic_file "${log_file}" "diagnostic_shell_limits" \ - "${shell_proc_dir}/limits" - if command -v awk >/dev/null 2>&1; then - # shellcheck disable=SC2016 # The expression is evaluated by awk. - run_failure_diagnostic "${log_file}" "${probe_deadline}" awk ' - $1 ~ /^(Pid:|PPid:|Threads:|VmPeak:|VmSize:|VmHWM:|VmRSS:|RssAnon:|RssFile:|RssShmem:|Cpus_allowed_list:|Mems_allowed_list:|voluntary_ctxt_switches:|nonvoluntary_ctxt_switches:)$/ { - print - } - ' "${shell_proc_dir}/status" - fi - if command -v nproc >/dev/null 2>&1; then - run_failure_diagnostic "${log_file}" "${probe_deadline}" nproc - fi - run_failure_diagnostic "${log_file}" "${probe_deadline}" \ - /bin/bash -c 'ulimit -a' - if command -v ps >/dev/null 2>&1 && command -v head >/dev/null 2>&1; then - append_failure_diagnostic_note "${log_file}" \ - "process_snapshot=top_100_by_rss" - run_failure_diagnostic "${log_file}" "${probe_deadline}" \ - /bin/bash -c \ - 'ps -eo pid,ppid,stat,etimes,nlwp,pcpu,rss,vsz,comm --sort=-rss | head -n 101' - fi - return 0 -} - -collect_mount_diagnostic() { - local log_file=$1 - local probe_deadline=$2 - local label=$3 - local path=$4 - - printf '\n%s (%s):\n' "${label}" "${path}" >> "${log_file}" - if [[ ! -e "${path}" ]]; then - echo "[path unavailable]" >> "${log_file}" - return 0 - fi - if command -v findmnt >/dev/null 2>&1; then - run_failure_diagnostic "${log_file}" "${probe_deadline}" findmnt \ - -n -T "${path}" -o TARGET,FSTYPE,VFS-OPTIONS - fi - if command -v df >/dev/null 2>&1; then - run_failure_diagnostic "${log_file}" "${probe_deadline}" df -h \ - --output=fstype,size,used,avail,pcent,itotal,iused,iavail,ipcent,target \ - -- "${path}" - fi - return 0 -} - -collect_execution_resource_diagnostics() { - local log_file=$1 - local probe_deadline=$2 - - # free and /proc/meminfo generally describe the node from a pod. Cgroups are - # the live source for this container's resource limits, usage, and events. - collect_cgroup_diagnostics "${log_file}" - collect_process_diagnostics "${log_file}" "${probe_deadline}" - - append_failure_diagnostic_section "${log_file}" \ - "Mount-visible filesystem capacity" - append_failure_diagnostic_note "${log_file}" \ - "These values are filesystem views, not Kubernetes ephemeral-storage quotas." - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "checkout" "${amd_diagnostics_checkout_root}" - if is_native_runtime; then - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "native workspace emptyDir" "${amd_diagnostics_workspace}" - else - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "test workspace (if mounted)" "${amd_diagnostics_workspace}" - fi - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "temporary directory" "${TMPDIR:-/tmp}" - if [[ "${TMPDIR:-/tmp}" != "/tmp" ]]; then - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "system temporary directory" /tmp - fi - if is_native_runtime; then - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "shared-memory emptyDir" /dev/shm - else - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "shared-memory mount" /dev/shm - fi - if [[ -n "${HF_HOME:-}" ]]; then - collect_mount_diagnostic "${log_file}" "${probe_deadline}" \ - "Hugging Face cache" "${HF_HOME}" - fi - return 0 -} - -collect_amd_device_nodes() { - local log_file=$1 - local probe_deadline=$2 - local device="" - local device_count=0 - - for device in /dev/kfd /dev/dri/renderD*; do - if [[ ! -e "${device}" ]]; then - continue - fi - device_count=$((device_count + 1)) - if command -v stat >/dev/null 2>&1; then - run_failure_diagnostic "${log_file}" "${probe_deadline}" stat -Lc \ - '%n type=%F mode=%a owner=%u:%g device=%t:%T' -- "${device}" - else - append_failure_diagnostic_note "${log_file}" "device=${device}" - fi - done - if [[ "${device_count}" -eq 0 ]]; then - append_failure_diagnostic_note "${log_file}" \ - "No /dev/kfd or /dev/dri/renderD* device nodes are visible." - fi - return 0 -} - -collect_rocm_failure_diagnostics() { - local exit_code=$1 - local job_id="${BUILDKITE_JOB_ID:-local}" - local retry_count="${BUILDKITE_RETRY_COUNT:-0}" - local parallel_job="${BUILDKITE_PARALLEL_JOB:-0}" - local diagnostics_relative_path="" - local diagnostics_path="" - local diagnostics_parent="" - local diagnostics_fallback_root="" - local diagnostics_storage="checkout" - local diagnostics_artifact_summary="" - local checkout_real="" - local diagnostics_parent_real="" - local exit_signal="" - local probe_deadline=0 - local runtime="single-node-docker" - local diagnostics_scope="outer-runner-after-test-container-exit" - local identity_label="runner" - local identity_value="${HOSTNAME:-unknown}" - local k8s_pod="unknown" - local k8s_namespace="${amd_diagnostics_k8s_namespace}" - local oom_kill_count="" - local summary_identity_label="Runner" - local -a summary_rows=() - - if [[ "${amd_diagnostics_collected}" == "1" ]]; then - return 0 - fi - amd_diagnostics_collected=1 - - if is_native_runtime; then - runtime="native-kubernetes" - diagnostics_scope="current-container-cgroup-and-namespaces" - if [[ -n "${amd_diagnostics_pod_name}" ]]; then - identity_label="pod" - identity_value="${amd_diagnostics_pod_name}" - k8s_pod="${amd_diagnostics_pod_name}" - else - identity_label="container_hostname" - fi - elif [[ "${amd_diagnostics_execution_mode}" == "multi-node" ]]; then - runtime="multi-node-docker" - fi - if [[ "${k8s_namespace}" == "unknown" \ - && -r /var/run/secrets/kubernetes.io/serviceaccount/namespace ]]; then - k8s_namespace=$( - tr -d '\r\n' \ - < /var/run/secrets/kubernetes.io/serviceaccount/namespace 2>/dev/null - ) - fi - - job_id="${job_id//[^A-Za-z0-9_.-]/_}" - retry_count="${retry_count//[^A-Za-z0-9_.-]/_}" - parallel_job="${parallel_job//[^A-Za-z0-9_.-]/_}" - - if [[ "${exit_code}" -gt 128 && "${exit_code}" -le 192 ]]; then - exit_signal=$(kill -l "$((exit_code - 128))" 2>/dev/null || true) - fi - - # Buildkite keeps the supplied artifact path. Restrict the configurable - # directory to a checkout-relative path so the UI shows a clean artifact key. - if [[ -z "${amd_diagnostics_dir}" \ - || "${amd_diagnostics_dir}" == /* \ - || "${amd_diagnostics_dir}" == *".."* ]]; then - echo "WARNING: ignoring unsafe VLLM_CI_DIAGNOSTICS_DIR" - amd_diagnostics_dir="artifacts/amd-gpu-diagnostics" - fi - diagnostics_relative_path="${amd_diagnostics_dir}/${job_id}/diagnostics.log" - diagnostics_path="${amd_diagnostics_checkout_root}/${diagnostics_relative_path}" - diagnostics_parent=$(dirname "${diagnostics_path}") - checkout_real=$(readlink -m "${amd_diagnostics_checkout_root}" 2>/dev/null || true) - diagnostics_parent_real=$(readlink -m "${diagnostics_parent}" 2>/dev/null || true) - if [[ -z "${checkout_real}" || -z "${diagnostics_parent_real}" \ - || ("${diagnostics_parent_real}" != "${checkout_real}" \ - && "${diagnostics_parent_real}" != "${checkout_real}/"*) ]] \ - || ! mkdir -p "${diagnostics_parent}" \ - || [[ -L "${diagnostics_path}" ]] \ - || ! (set -o noclobber; : > "${diagnostics_path}") 2>/dev/null; then - echo "WARNING: unable to create AMD CI diagnostics in the checkout; using temporary storage." - diagnostics_path="" - if diagnostics_fallback_root=$(mktemp -d -t \ - vllm-amd-diagnostics.XXXXXX 2>/dev/null); then - diagnostics_path="${diagnostics_fallback_root}/${diagnostics_relative_path}" - diagnostics_parent=$(dirname "${diagnostics_path}") - if (umask 077; mkdir -p "${diagnostics_parent}" \ - && (set -o noclobber; : > "${diagnostics_path}")) 2>/dev/null; then - diagnostics_storage="temporary-fallback" - else - rm -rf -- "${diagnostics_fallback_root}" || true - diagnostics_fallback_root="" - diagnostics_path="" - fi - fi - fi - if [[ -z "${diagnostics_path}" ]]; then - echo "WARNING: unable to create AMD CI diagnostics at ${diagnostics_relative_path}." - printf '\nAMD CI failure summary\n' - printf '%-22s | %s\n' \ - "Field" "Value" \ - "----------------------" "-----" \ - "Exit code" "${exit_code}" \ - "Signal" "${exit_signal:-none}" \ - "Runtime" "${runtime}" \ - "Pod/container" "${identity_value}" \ - "Diagnostics artifact" "unavailable: ${diagnostics_relative_path}" - return 0 - fi - - { - echo "AMD CI failure diagnostics" - echo "timestamp_utc=$(date -u +%Y-%m-%dT%H:%M:%SZ)" - echo "exit_code=${exit_code}" - echo "exit_signal=${exit_signal:-none}" - echo "build_url=${BUILDKITE_BUILD_URL:-unknown}" - echo "job_id=${BUILDKITE_JOB_ID:-unknown}" - echo "step_key=${BUILDKITE_STEP_KEY:-unknown}" - echo "step_label=${BUILDKITE_LABEL:-unknown}" - echo "test_group=${amd_diagnostics_test_group}" - echo "retry_count=${retry_count}" - echo "parallel_job=${parallel_job}" - echo "execution_mode=${amd_diagnostics_execution_mode}" - echo "runtime=${runtime}" - echo "diagnostics_scope=${diagnostics_scope}" - echo "expected_gpu_count=${amd_diagnostics_expected_gpu_count}" - echo "probe_budget_seconds=${amd_diagnostics_probe_budget_seconds}" - echo "command_timeout_seconds=${amd_diagnostics_command_timeout_seconds}" - echo "diagnostics_storage=${diagnostics_storage}" - echo "agent_name=${BUILDKITE_AGENT_NAME:-unknown}" - echo "container_hostname=${HOSTNAME:-unknown}" - echo "k8s_pod=${k8s_pod}" - echo "k8s_namespace=${k8s_namespace:-unknown}" - echo "k8s_node=${amd_diagnostics_k8s_node_name:-unknown}" - echo "kernel=$(uname -srmo 2>/dev/null || echo unknown)" - } > "${diagnostics_path}" - - probe_deadline=$((SECONDS + amd_diagnostics_probe_budget_seconds)) - - # Cgroup, process, and mount state is fast and most useful for explaining OOM, - # throttling, PID exhaustion, and emptyDir pressure. Capture it before external - # probes can consume the shared deadline. - collect_execution_resource_diagnostics \ - "${diagnostics_path}" "${probe_deadline}" - - append_failure_diagnostic_section "${diagnostics_path}" \ - "AMD GPU diagnostics (tool-visible devices)" - if [[ "${amd_diagnostics_expected_gpu_count}" == "0" ]]; then - append_failure_diagnostic_note "${diagnostics_path}" \ - "CPU-only job; AMD GPU probes skipped." - else - if [[ "${runtime}" == "native-kubernetes" ]]; then - append_failure_diagnostic_note "${diagnostics_path}" \ - "SMI visibility may be broader than the Kubernetes device allocation." - else - append_failure_diagnostic_note "${diagnostics_path}" \ - "Outer-runner SMI visibility may be broader than the test container." - fi - collect_amd_device_nodes "${diagnostics_path}" "${probe_deadline}" - fi - if [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \ - && command -v amd-smi >/dev/null 2>&1; then - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ - amd-smi version - # Bus data identifies a card within the public node without publishing its - # persistent UUID, serial number, or process list. - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ - amd-smi static --bus --gpu all - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ - amd-smi metric --ecc --ecc-blocks --pcie --gpu all - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ - amd-smi static --ras --gpu all - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ - amd-smi metric --power --temperature --usage --mem-usage --gpu all - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" \ - amd-smi xgmi --link-status --gpu all - elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]] \ - && command -v rocm-smi >/dev/null 2>&1; then - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" rocm-smi \ - --showbus --showreplaycount --showrasinfo --showpagesinfo - elif [[ "${amd_diagnostics_expected_gpu_count}" != "0" ]]; then - append_failure_diagnostic_note "${diagnostics_path}" \ - "Neither amd-smi nor rocm-smi is available." - fi - - if [[ "${runtime}" == "native-kubernetes" ]]; then - append_failure_diagnostic_section "${diagnostics_path}" \ - "Unauthenticated pod network reachability" - else - append_failure_diagnostic_section "${diagnostics_path}" \ - "Unauthenticated outer-runner network reachability" - fi - if command -v curl >/dev/null 2>&1; then - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \ - --silent --location --proto '=https' --proto-redir '=https' \ - --max-redirs 3 --output /dev/null --connect-timeout 2 --max-time 4 \ - --write-out 'target=huggingface http_code=%{http_code} dns_done_s=%{time_namelookup} connect_done_s=%{time_connect} tls_done_s=%{time_appconnect} first_byte_s=%{time_starttransfer} total_s=%{time_total}\n' \ - https://huggingface.co/api/models/gpt2 - run_failure_diagnostic "${diagnostics_path}" "${probe_deadline}" curl -q \ - --silent --location --proto '=https' --proto-redir '=https' \ - --max-redirs 3 --output /dev/null --connect-timeout 2 --max-time 4 \ - --write-out 'target=github_git_smart_http http_code=%{http_code} dns_done_s=%{time_namelookup} connect_done_s=%{time_connect} tls_done_s=%{time_appconnect} first_byte_s=%{time_starttransfer} total_s=%{time_total}\n' \ - 'https://github.com/vllm-project/ci-infra.git/info/refs?service=git-upload-pack' - else - append_failure_diagnostic_note "${diagnostics_path}" \ - "curl unavailable; reachability probes skipped." - fi - - if [[ "${runtime}" == "native-kubernetes" \ - && -n "${amd_diagnostics_memory_events_path}" \ - && -r "${amd_diagnostics_memory_events_path}" ]]; then - oom_kill_count=$( - awk '$1 == "oom_kill" { print $2; exit }' \ - "${amd_diagnostics_memory_events_path}" 2>/dev/null - ) - if [[ ! "${oom_kill_count}" =~ ^[0-9]+$ ]]; then - oom_kill_count="" - fi - fi - - case "${identity_label}" in - pod) - summary_identity_label="Kubernetes pod" - ;; - container_hostname) - summary_identity_label="Container hostname" - ;; - esac - - summary_rows=( - "Field" "Value" - "----------------------" "-----" - "Exit code" "${exit_code}" - "Signal" "${exit_signal:-none}" - "Runtime" "${runtime}" - "${summary_identity_label}" "${identity_value}" - ) - if [[ -n "${amd_diagnostics_k8s_node_name}" ]]; then - summary_rows+=("Kubernetes node" "${amd_diagnostics_k8s_node_name}") - fi - if [[ -n "${oom_kill_count}" ]]; then - summary_rows+=("Cgroup OOM kills" "${oom_kill_count}") - fi - diagnostics_artifact_summary="${diagnostics_relative_path}" - if [[ -n "${diagnostics_fallback_root}" ]]; then - if upload_failure_diagnostics_artifact \ - "${diagnostics_fallback_root}" "${diagnostics_relative_path}"; then - echo "Uploaded AMD CI diagnostics artifact from temporary storage: ${diagnostics_relative_path}" - diagnostics_artifact_summary="${diagnostics_relative_path} (direct upload)" - else - echo "WARNING: failed to upload temporary AMD CI diagnostics artifact: ${diagnostics_relative_path}" - echo "--- AMD CI diagnostics (artifact upload failed)" - cat "${diagnostics_path}" || true - diagnostics_artifact_summary="job log only: direct upload failed" - fi - rm -rf -- "${diagnostics_fallback_root}" || \ - echo "WARNING: unable to remove temporary AMD CI diagnostics storage." - fi - summary_rows+=("Diagnostics artifact" "${diagnostics_artifact_summary}") - - printf '\nAMD CI failure summary\n' - printf '%-22s | %s\n' "${summary_rows[@]}" - - return 0 -} - -handle_pytest_exit() { - local exit_code=$1 - if [ "$exit_code" -eq 5 ]; then - echo "Pytest exit code 5 (no tests collected) - treating as success." - exit 0 - fi - exit "$exit_code" -} - -############################################################################### -# Pytest marker/keyword re-quoting -# -# When commands are passed through Buildkite -> shell -> $* -> bash -c, -# quotes around multi-word pytest -m/-k expressions get stripped: -# pytest -v -s -m 'not cpu_test' v1/core -# becomes: -# pytest -v -s -m not cpu_test v1/core -# -# pytest then interprets "cpu_test" as a file path, not part of the marker. -# -# This function detects unquoted expressions after -m/-k and re-quotes them -# by collecting tokens until a recognizable boundary is reached: -# - test path (contains '/') -# - test file (ends with '.py') -# - another pytest flag (--xxx or -x single-char flags) -# - command separator (&& || ; |) -# - environment variable assignment (FOO=bar) -# -# Single-word markers (e.g. -m cpu_test, -m hybrid_model) pass through -# unquoted since they have no spaces and work fine. -# -# Already-quoted expressions (containing literal single quotes) are passed -# through untouched to avoid double-quoting well-formed shell fragments. -# -# NOTE: This ONLY fixes -m/-k flags. It cannot recover arbitrary inner -# double-quotes stripped by the calling shell (see header comment). -# Use VLLM_TEST_COMMANDS to avoid the problem entirely. -############################################################################### -re_quote_pytest_markers() { - local input="$1" - local output="" - local collecting=false - local marker_buf="" - - # Strip backslash-newline continuations, then flatten remaining newlines - local flat="${input//$'\\\n'/ }" - flat="${flat//$'\n'/ }" - - # Disable globbing to prevent *.py etc. from expanding during read -ra - local restore_glob - restore_glob="$(shopt -p -o noglob 2>/dev/null || true)" - set -o noglob - local -a words - read -ra words <<< "$flat" - eval "$restore_glob" - - for word in "${words[@]}"; do - if $collecting; then - # If the token we're about to collect already contains a literal - # single quote, the expression was already quoted upstream. - # Flush and stop collecting. - if [[ "$word" == *"'"* ]]; then - if [[ -n "$marker_buf" ]]; then - # Should not normally happen (partial buf + quote), flush raw - output+="${marker_buf} " - marker_buf="" - fi - output+="${word} " - collecting=false - continue - fi - - local is_boundary=false - case "$word" in - # Line-continuation artifact - "\\") - is_boundary=true ;; - # Command separators - "&&"|"||"|";"|"|") - is_boundary=true ;; - # Long flags (--ignore, --shard-id, etc.) - --*) - is_boundary=true ;; - # Short flags (-v, -s, -x, etc.) but NOT negative marker tokens - # like "not" which don't start with "-". Also skip -k/-m which - # would start a new marker (handled below). - -[a-zA-Z]) - is_boundary=true ;; - # Test path (contains /) - */*) - is_boundary=true ;; - # Test file (ends with .py, possibly with ::method) - *.py|*.py::*) - is_boundary=true ;; - # Environment variable assignment preceding a command (FOO=bar) - *=*) - # Only treat as boundary if it looks like VAR=value, not - # pytest filter expressions like num_gpus=2 inside markers - if [[ "$word" =~ ^[A-Z_][A-Z0-9_]*= ]]; then - is_boundary=true - fi - ;; - esac - - if $is_boundary; then - # Strip surrounding double quotes if present (from upstream - # single-to-double conversion); without this, wrapping below - # would produce '"expr"' with literal double-quote characters. - if [[ "$marker_buf" == '"'*'"' ]]; then - marker_buf="${marker_buf#\"}" - marker_buf="${marker_buf%\"}" - fi - # Flush the collected marker expression - if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then - output+="'${marker_buf}' " - else - output+="${marker_buf} " - fi - collecting=false - marker_buf="" - # Check if this boundary word itself starts a new -m/-k - if [[ "$word" == "-m" || "$word" == "-k" ]]; then - output+="${word} " - collecting=true - # Drop stray backslash tokens silently - elif [[ "$word" == "\\" ]]; then - : - else - output+="${word} " - fi - else - # Accumulate into marker buffer - if [[ -n "$marker_buf" ]]; then - marker_buf+=" ${word}" - else - marker_buf="${word}" - fi - fi - elif [[ "$word" == "-m" || "$word" == "-k" ]]; then - output+="${word} " - collecting=true - marker_buf="" - else - output+="${word} " - fi - done - - # Flush any trailing marker expression (marker at end of command) - if $collecting && [[ -n "$marker_buf" ]]; then - # Strip surrounding double quotes (see mid-stream flush comment) - if [[ "$marker_buf" == '"'*'"' ]]; then - marker_buf="${marker_buf#\"}" - marker_buf="${marker_buf%\"}" - fi - if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then - output+="'${marker_buf}'" - else - output+="${marker_buf}" - fi - fi - - echo "${output% }" -} - -# shellcheck disable=SC2317 # Called indirectly by the EXIT trap. -handle_amd_runner_exit() { - local exit_code=${1:-$?} - if [[ "${exit_code}" -ne 0 ]]; then - collect_rocm_failure_diagnostics "${exit_code}" - fi - return 0 -} - -# Catch both test failures and wrapper/setup failures. Runtime-specific cleanup -# traps below replace this trap and call the same handler after cleanup. -trap handle_amd_runner_exit EXIT - -############################################################################### -# Main -############################################################################### - -if is_native_runtime; then - echo "--- Native in-pod ROCm CI (AMD_CI_RUNTIME=${AMD_CI_RUNTIME:-unset}, NATIVE_CI=${NATIVE_CI:-unset})" - artifact_work_dir="" - - # shellcheck disable=SC2317 # Called indirectly by the EXIT trap. - cleanup_native_workspace() { - local exit_code=$? - if [[ -n "${artifact_work_dir}" ]]; then - rm -rf "${artifact_work_dir}" - fi - handle_amd_runner_exit "${exit_code}" - } - trap cleanup_native_workspace EXIT - - if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then - commands="${VLLM_TEST_COMMANDS}" - commands_source="env" - else - commands="$*" - commands_source="argv" - if [[ -z "$commands" ]]; then - echo "Error: No test commands provided for native CI." >&2 - exit 1 - fi - fi - - if [[ "$commands_source" == "argv" ]]; then - commands=$(re_quote_pytest_markers "$commands") - fi - - if is_multi_node "$commands"; then - echo "Native CI does not support multi-node jobs yet." - exit 1 - fi - - if ! initialize_native_environment; then - echo "Failed to initialize the native test environment" - exit 1 - fi - if [[ "${BUILDKITE_AGENT_META_DATA_QUEUE:-}" == *dpx* \ - && "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" != "0" ]]; then - check_dpx_gpu_exclusivity || exit 1 - fi - if [[ "${commands}" == *python_only_compile.sh* ]]; then - # This no-GPU job validates the ROCm precompiled/editable install path, - # rather than CPU runtime platform selection. - VLLM_TARGET_DEVICE=rocm - export VLLM_TARGET_DEVICE - fi - if ! prepare_native_workspace "${commands}"; then - echo "Failed to prepare native test workspace" - exit 1 - fi - - export PYTHONPATH="${VLLM_CI_WORKSPACE:-/vllm-workspace}" - - echo "Native test commands: $commands" - run_native_preflight || exit 1 - echo "--- Test log" - # Keep AMD CI orchestration variables out of vLLM's runtime environment. - clear_ci_orchestration_env - /bin/bash -o pipefail -c "${commands}" - handle_pytest_exit "$?" -fi - -# --- GPU initialization for legacy Docker execution --- -echo "--- ROCm info" -rocminfo - -# --- Docker status --- -report_docker_usage - -# --- Pull test image --- -echo "--- Pulling container" -image_name="${VLLM_CI_FALLBACK_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:-local}}" -artifact_work_dir="" -container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)" - -# shellcheck disable=SC2317 # Called indirectly by the EXIT trap. -remove_docker_container() { - local exit_code=$? - if docker container inspect "${container_name}" >/dev/null 2>&1; then - docker rm -f "${container_name}" || true - fi - if [[ "${VLLM_CI_REMOVE_TEST_IMAGE:-0}" == "1" ]]; then - docker image rm -f "${image_name}" || true - else - # Keep images by default so later jobs on the same AMD node can reuse layers. - echo "Keeping ROCm test image locally: ${image_name}" - fi - if [[ -n "${artifact_work_dir}" ]]; then - rm -rf "${artifact_work_dir}" - fi - handle_amd_runner_exit "${exit_code}" -} -trap remove_docker_container EXIT - -# python_only_compile.sh runs `python setup.py develop` and needs the full repo tree -# under /vllm-workspace (Dockerfile.rocm test stage: mkdir src && mv vllm). -# The ROCm wheel artifact tarball only ships a thin tree (tests, etc.), so -# artifact images cannot satisfy that test — use the full rocm/vllm-ci image. -_cmd_probe="${VLLM_TEST_COMMANDS:-}" -if [[ -z "${_cmd_probe}" ]]; then - _cmd_probe="$*" -fi -if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" == "1" && "${_cmd_probe}" == *python_only_compile.sh* ]]; then - echo "INFO: disabling VLLM_CI_USE_ARTIFACTS for python_only_compile (requires full /vllm-workspace tree)" - export VLLM_CI_USE_ARTIFACTS=0 -fi -unset -v _cmd_probe - -if ! prepare_artifact_image; then - echo "Using full ROCm CI image: ${image_name}" - docker pull "${image_name}" || exit 1 -fi - -# --- Prepare commands --- -echo "--- Running container" - -HF_CACHE="$(realpath ~)/huggingface" -mkdir -p "${HF_CACHE}" -HF_MOUNT="/root/.cache/huggingface" - -# Hugging Face Hub defaults to 10s request/download timeouts, while the ROCm -# CI image currently raises downloads to 60s. AMD model-test jobs routinely -# start from a cold or partially-populated shared cache, and the 60s read cap -# has still timed out before pytest reached the vLLM behavior under test. -# Keep the CI default explicit and overridable from the Buildkite environment. -: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" -: "${HF_HUB_ETAG_TIMEOUT:=60}" - -# ---- Command source selection ---- -# Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact). -# Fall back to $* for backward compatibility, but warn that inner -# double-quotes will have been stripped by the calling shell. -if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then - commands="${VLLM_TEST_COMMANDS}" - commands_source="env" - echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)" -else - commands="$*" - commands_source="argv" - if [[ -z "$commands" ]]; then - echo "Error: No test commands provided." >&2 - echo "Usage:" >&2 - echo " Preferred: VLLM_TEST_COMMANDS='...' bash $0" >&2 - echo " Legacy: bash $0 \"commands here\"" >&2 - exit 1 - fi - echo "Commands sourced from positional args (legacy mode)" - echo "WARNING: Inner double-quotes in the command string may have been" - echo " stripped by the calling shell. If you see syntax errors, switch to:" - echo " export VLLM_TEST_COMMANDS='your commands here'" - echo " bash $0" -fi - -echo "Raw commands: $commands" - -# Only try to repair stripped pytest -m/-k quoting in legacy argv mode. -# VLLM_TEST_COMMANDS preserves inner quoting already, and re-quoting that path -# can corrupt embedded echo strings or otherwise well-formed shell fragments. -if [[ "$commands_source" == "argv" ]]; then - commands=$(re_quote_pytest_markers "$commands") - echo "After re-quoting: $commands" -else - echo "Skipping re-quoting for VLLM_TEST_COMMANDS input" -fi - -echo "Final commands: $commands" - -# Match native CPU jobs even when the container can see AMD devices. -cpu_platform_env=() -if [[ "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" == "0" \ - && "$commands" != *python_only_compile.sh* ]]; then - cpu_platform_env=(-e "VLLM_TARGET_DEVICE=cpu") -fi - -standalone_merge_base_env=() -if [[ "$commands" == *python_only_compile.sh* ]]; then - # The ROCm test image often ships /vllm-workspace without .git. Resolve the - # wheels.vllm.ai commit from the agent checkout for this test only. - vllm_standalone_merge_base="" - checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}" - if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then - checkout="." - fi - # Pass safe.directory per-command because Buildkite uses mixed user IDs. - if git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse --is-inside-work-tree >/dev/null 2>&1; then - vllm_standalone_merge_base="$( - git -c "safe.directory=${checkout}" -C "${checkout}" merge-base HEAD origin/main 2>/dev/null || true - )" - fi - if [[ -z "${vllm_standalone_merge_base}" ]]; then - vllm_standalone_merge_base="${BUILDKITE_COMMIT:-}" - fi - echo "INFO: passing CI_STANDALONE_MERGE_BASE into container: ${vllm_standalone_merge_base}" - standalone_merge_base_env=(-e "CI_STANDALONE_MERGE_BASE=${vllm_standalone_merge_base}") -fi - -MYPYTHONPATH="/vllm-workspace" - -container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}" -container_job_id="${container_job_id//[^A-Za-z0-9_.-]/_}" -container_job_id_short="${container_job_id:0:8}" -CONTAINER_TMPDIR="/tmp/vllm-${container_job_id_short}" -CONTAINER_CACHE_ROOT="/tmp/vllm-buildkite-${container_job_id}/cache" -CONTAINER_PREFLIGHT="mkdir -p \"\$TMPDIR\" \"\$TIKTOKEN_RS_CACHE_DIR\" \"\$TORCHINDUCTOR_CACHE_DIR\" \"\$TRITON_CACHE_DIR\" \"\$VLLM_CACHE_ROOT\" \"\$XDG_CACHE_HOME\" && python -c \"import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing\"" - -# Verify GPU access -render_gid=$(getent group render | cut -d: -f3) -if [[ -z "$render_gid" ]]; then - echo "Error: 'render' group not found. This is required for GPU access." >&2 - exit 1 +n = torch.cuda.device_count() +assert torch.version.hip and torch.cuda.is_available() and n == $gpus, \ + f"preflight: expected $gpus ROCm GPU(s), found {n}" +EOF fi -# --- RDMA device passthrough (conditional) --- -# If the host has RDMA devices, pass them through so tests like -# test_moriio_connector can access ibverbs. On hosts without RDMA -# hardware the tests will gracefully skip via _rdma_available(). -RDMA_FLAGS="" -if [ -d /dev/infiniband ]; then - echo "RDMA devices detected on host, enabling passthrough" - RDMA_FLAGS="--device /dev/infiniband --cap-add=IPC_LOCK" +# --------------------------------------------------------------------------- +# Run +# --------------------------------------------------------------------------- +# Scrub orchestration vars before handing control to test commands. +# (VLLM_CI_* that are test knobs are intentionally left — see tests/ci_envs.py) +unset VLLM_TEST_COMMANDS VLLM_CI_EXPECTED_GPU_COUNT VLLM_CI_DIAGNOSTICS_DIR + +echo "Test commands: $commands" +echo "--- Test log" +_teardown_log workload_started "timeout_s=${CONTAINER_TIMEOUT_S:-0}" + +rc=0 +if [[ ${CONTAINER_TIMEOUT_S:-0} -gt 0 ]]; then + # timeout fires TERM, then KILL after 10 s — mirrors the old docker run path. + # Not --foreground: the whole process group (pytest, engine workers) must stop, not just bash. + # That group is in the background of the job's PTY, where reading the terminal stops it + # (SIGTTIN: apt, transformers' trust_remote_code prompt), so stdin is /dev/null. + timeout --verbose --signal=TERM --kill-after=10s \ + "${CONTAINER_TIMEOUT_S}s" bash -o pipefail -c "$commands" /dev/null 2>&1; then + present+=("${arch}") + else + missing+=("${arch}") + fi + done + + if [ "${#present[@]}" -eq 0 ]; then + echo "WARNING: no ${suffix} images found in ECR for ${COMMIT}, skipping the ${suffix} publish" + elif [ "${#missing[@]}" -eq 0 ]; then + publish_image_family "${suffix}" + else + echo "ERROR: partial ${suffix} build (present: ${present[*]}; missing: ${missing[*]})." >&2 + echo " Refusing to publish to avoid per-arch and multi-arch tags disagreeing." >&2 + exit 1 + fi +} + RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//') if [ -z "${RELEASE_VERSION}" ]; then echo "ERROR: release-version metadata not set" @@ -41,25 +98,7 @@ aws ecr-public get-login-password --region us-east-1 | \ # ---- CUDA (default: 13.0) ---- if target_enabled cuda-13-0; then - docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" - docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" - - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" vllm/vllm-openai:latest-x86_64 - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" - docker push vllm/vllm-openai:latest-x86_64 - docker push "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" - - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" vllm/vllm-openai:latest-aarch64 - docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" - docker push vllm/vllm-openai:latest-aarch64 - docker push "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" - - docker manifest rm vllm/vllm-openai:latest || true - docker manifest rm "vllm/vllm-openai:v${RELEASE_VERSION}" || true - docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 - docker manifest create "vllm/vllm-openai:v${RELEASE_VERSION}" "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64" "vllm/vllm-openai:v${RELEASE_VERSION}-aarch64" - docker manifest push vllm/vllm-openai:latest - docker manifest push "vllm/vllm-openai:v${RELEASE_VERSION}" + publish_image_family "" ZSTD_DIGEST=$(docker buildx imagetools inspect \ "public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64" \ @@ -69,36 +108,46 @@ if target_enabled cuda-13-0; then "vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-zstd" fi +# ---- AWS EFA variant (docker/Dockerfile.efa) ---- +# No cu134 entry: Rubin CUDA 13.4 is nightly-only and is not promoted to +# DockerHub for the base variant either. + +if target_enabled cuda-13-0-efa; then + publish_optional_family "-efa" +fi + # ---- ROCm ---- -# Each ROCm stack is identified by its base Dockerfile (see rocm/stack.sh). The -# default stack also owns the plain :latest / :v tags. +# Copied from ECR (publish.sh). The default stack (rocm) also owns the plain :latest / :v tags. + +# pin_rocm — @, so what was checked is what is published +pin_rocm() { + local d + d=$(docker buildx imagetools inspect "$1" | awk '$1 == "Digest:" { print $2; exit }') + [[ $d =~ ^sha256:[0-9a-f]{64}$ ]] || { echo "cannot resolve the digest of $1" >&2; return 1; } + echo "$1@$d" +} publish_rocm_stack() { - local base_dockerfile="$1" default="$2" suffix - ( - # shellcheck source=.buildkite/scripts/rocm/stack.sh - source .buildkite/scripts/rocm/stack.sh "$base_dockerfile" - local suffixes=("-${ROCM_STACK_VARIANT}") - [[ "$default" == "1" ]] && suffixes=("" "-${ROCM_STACK_VARIANT}") - docker pull "$ROCM_STACK_ECR_IMAGE" - docker pull "$ROCM_STACK_ECR_BASE" - for suffix in "${suffixes[@]}"; do - for tag in "latest${suffix}" "v${RELEASE_VERSION}${suffix}"; do - docker tag "$ROCM_STACK_ECR_IMAGE" "vllm/vllm-openai-rocm:${tag}" - docker tag "$ROCM_STACK_ECR_BASE" "vllm/vllm-openai-rocm:${tag}-base" - docker push "vllm/vllm-openai-rocm:${tag}" - docker push "vllm/vllm-openai-rocm:${tag}-base" - done - done - ) + local v s sfx img base ver ecr=public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT} hub=vllm/vllm-openai-rocm + v=$(buildkite-agent meta-data get "rocm-$1-variant") + img=$(pin_rocm "$ecr-$v") + base=$(pin_rocm "$ecr-$v-base") + # Another build of this commit (e.g. a branch push) may have retagged it with its own version. + ver=$(docker buildx imagetools inspect "$img" \ + --format '{{index .Image.Config.Labels "org.opencontainers.image.version"}}') + if [[ ${ver%%+*} != "$RELEASE_VERSION" ]]; then + echo "ERROR: $img is version '$ver', not $RELEASE_VERSION" >&2 + exit 1 + fi + sfx=("-$v"); if [[ $1 == rocm ]]; then sfx+=(""); fi + for s in "${sfx[@]}"; do + docker buildx imagetools create -t "$hub:latest$s" -t "$hub:v${RELEASE_VERSION}$s" "$img" + docker buildx imagetools create -t "$hub:latest$s-base" -t "$hub:v${RELEASE_VERSION}$s-base" "$base" + done } if target_enabled rocm; then - publish_rocm_stack docker/Dockerfile.rocm_base 1 -fi - -if target_enabled rocm72; then - publish_rocm_stack docker/Dockerfile.rocm_72_base 0 + publish_rocm_stack rocm fi # ---- XPU ---- diff --git a/.buildkite/scripts/push-nightly-builds-rocm.sh b/.buildkite/scripts/push-nightly-builds-rocm.sh deleted file mode 100644 index 8b120883eff4..000000000000 --- a/.buildkite/scripts/push-nightly-builds-rocm.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# -# Push a ROCm stack's nightly base image and nightly image from ECR to Docker Hub -# under vllm/vllm-openai-rocm, as base-nightly-, nightly- and their -# - forms, then prune old nightly tags. With --default, also push the plain -# base-nightly and nightly tags for the same images. -# Run when NIGHTLY=1 after the stack's release image has been pushed to ECR. -# -# Usage: push-nightly-builds-rocm.sh [--default] -# e.g. "docker/Dockerfile.rocm_base --default" (ROCm 10.0, also :nightly) -# and "docker/Dockerfile.rocm_72_base" (ROCm 7.2, :nightly-rocm72). -# -# Local testing (no push to Docker Hub): -# BUILDKITE_COMMIT= DRY_RUN=1 \ -# bash .buildkite/scripts/push-nightly-builds-rocm.sh docker/Dockerfile.rocm_base --default -# Requires: AWS CLI configured (for ECR public login), Docker. For full run: Docker Hub login. - -set -euxo pipefail - -usage() { - echo "Usage: $0 [--default]" >&2 -} - -PUSH_DEFAULT=0 -case "$#" in - 1) ;; - 2) - if [[ "$2" != "--default" ]]; then - usage - exit 2 - fi - PUSH_DEFAULT=1 - ;; - *) - usage - exit 2 - ;; -esac - -# Use BUILDKITE_COMMIT from env (required; set to a commit that has the image in ECR for local test) -BUILDKITE_COMMIT="${BUILDKITE_COMMIT:?Set BUILDKITE_COMMIT to the commit SHA that has the image in ECR (e.g. from a previous release pipeline run)}" -DRY_RUN="${DRY_RUN:-0}" - -# shellcheck source=.buildkite/scripts/rocm/stack.sh -source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "$1" -DOCKERHUB_REPO="vllm/vllm-openai-rocm" -VARIANT="$ROCM_STACK_VARIANT" - -BASE_TAGS=("base-nightly-${VARIANT}" "base-nightly-${VARIANT}-${BUILDKITE_COMMIT}") -TAGS=("nightly-${VARIANT}" "nightly-${VARIANT}-${BUILDKITE_COMMIT}") -if [[ "$PUSH_DEFAULT" == "1" ]]; then - BASE_TAGS+=("base-nightly" "base-nightly-${BUILDKITE_COMMIT}") - TAGS+=("nightly" "nightly-${BUILDKITE_COMMIT}") -fi - -echo "Pushing base image $ROCM_STACK_ECR_BASE as: ${BASE_TAGS[*]}" -echo "Pushing release image $ROCM_STACK_ECR_IMAGE as: ${TAGS[*]}" - -# Login to ECR and pull the stack's images -aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7 -docker pull "$ROCM_STACK_ECR_BASE" -docker pull "$ROCM_STACK_ECR_IMAGE" - -for tag in "${BASE_TAGS[@]}"; do docker tag "$ROCM_STACK_ECR_BASE" "$DOCKERHUB_REPO:$tag"; done -for tag in "${TAGS[@]}"; do docker tag "$ROCM_STACK_ECR_IMAGE" "$DOCKERHUB_REPO:$tag"; done - -if [[ "$DRY_RUN" == "1" ]]; then - echo "[DRY_RUN] Local tags created. Exiting without push." - exit 0 -fi - -# Push to Docker Hub (docker-login plugin runs before this step in CI) -for tag in "${BASE_TAGS[@]}" "${TAGS[@]}"; do docker push "$DOCKERHUB_REPO:$tag"; done -echo "Pushed $DOCKERHUB_REPO: ${BASE_TAGS[*]} ${TAGS[*]}" - -# Keep only the last 14 builds of each tag family. The plain nightly- families -# exclude the variant ones sharing their prefix; commit SHAs never start with "rocm". -CLEANUP=.buildkite/scripts/cleanup-nightly-builds.sh -bash "$CLEANUP" "nightly-${VARIANT}-" "$DOCKERHUB_REPO" -bash "$CLEANUP" "base-nightly-${VARIANT}-" "$DOCKERHUB_REPO" -if [[ "$PUSH_DEFAULT" == "1" ]]; then - bash "$CLEANUP" "nightly-" "$DOCKERHUB_REPO" "nightly-rocm" - bash "$CLEANUP" "base-nightly-" "$DOCKERHUB_REPO" "base-nightly-rocm" -fi diff --git a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh b/.buildkite/scripts/rocm/aiter-nightly-overlay.sh deleted file mode 100755 index bed4a3e09aec..000000000000 --- a/.buildkite/scripts/rocm/aiter-nightly-overlay.sh +++ /dev/null @@ -1,116 +0,0 @@ -#!/bin/bash -# Install the AITER nightly wheel over this build's ci_base. -# -# Runs in CI when AITER_NIGHTLY=1, once the stock ci_base it builds on is -# published. In order: -# 1. Select the newest nightly wheel that fits the stock ci_base -# (select_aiter_nightly_wheel.py, beside this script, run inside that image) -# 2. Build a new image: pip install that wheel over the stock ci_base -# 3. Push the new image as this build's ci_base-build tag, replacing the stock -# one, so later steps and native jobs use it; record its digest -# -# Buildkite meta-data, a key-value store shared by every job in this build: -# reads rocm-ci-base-image the stock ci_base -# r/w aiter-nightly-wheel-url the wheel step 1 chose; a retry reads it -# back and installs that same wheel -# writes aiter-nightly-ci-base the new image as tag@digest; the vLLM -# image build uses it as its ci_base -# -# Exit status: -# 10 no wheel fits this ci_base -# 11 the wheel does not install or its modules do not load -# Anything else is unexpected and passed through as is. -# Buildkite retries this step once on exit status 1 or 2, which it treats as -# transient Docker/registry failures (amd_infra_retry in the template). -# 10 and 11 sit outside that range: they are real AITER failures, and a retry -# would only fail the same way. - -set -euo pipefail - -SELECT_WHEEL="$(dirname "${BASH_SOURCE[0]}")/select_aiter_nightly_wheel.py" -BUILD_TAG="rocm/vllm-dev:ci_base-build-${BUILDKITE_BUILD_ID:?BUILDKITE_BUILD_ID is required}" - -# Scratch space for the build context and captured errors. The trap deletes it -# however the script exits: success, fail, or set -e. -scratch="$(mktemp -d)" -trap 'rm -rf "${scratch}"' EXIT -errors="${scratch}/errors" - -fail() { - local status="$1" title="$2" detail="$3" - echo "+++ :x: AITER nightly: ${title}" - echo "${detail}" >&2 - # shellcheck disable=SC2016 # the backticks are a literal markdown fence - printf ':x: **AITER nightly: %s**\n\n```\n%s\n```\n' "${title}" "${detail}" \ - | buildkite-agent annotate --style error --context aiter-nightly - exit "${status}" -} - -# Set by ensure-ci-base-amd; this step's depends_on guarantees it is there. -ci_base="$(buildkite-agent meta-data get rocm-ci-base-image)" -# Set below on the first attempt. Meta-data is per build, so a retry reuses -# that wheel, not a newer nightly. -pinned="$(buildkite-agent meta-data get aiter-nightly-wheel-url --default "")" - -echo "--- :docker: Pulling ${ci_base}" -docker pull "${ci_base}" - -echo "--- :mag: Selecting the AITER nightly for ${ci_base}" -set +e -out="$(docker run --rm -i --entrypoint python3 "${ci_base}" \ - -P - ${pinned:+--wheel-url "${pinned}"} \ - < "${SELECT_WHEEL}" 2> "${errors}")" -status=$? -set -e -cat "${errors}" >&2 -case "${status}" in - 0) ;; - 10) fail 10 "no wheel for this image" "${out}" ;; - *) fail "${status}" "could not select a wheel" "$(tail -n 20 "${errors}")" ;; -esac -# The first line is the wheel's URL; the rest, the annotation naming it. -wheel_url="$(head -n 1 <<< "${out}")" -note="$(tail -n +2 <<< "${out}")" -buildkite-agent meta-data set aiter-nightly-wheel-url "${wheel_url}" - -cat > "${scratch}/Dockerfile" <<'EOF' -ARG CI_BASE_IMAGE -FROM ${CI_BASE_IMAGE} -ARG AITER_WHEEL_URL -RUN python3 -m pip install --no-cache-dir "$AITER_WHEEL_URL" -# A wheel built against a newer glibc or libstdc++ than ci_base installs fine -# but fails every test that loads it; fail here instead. -RUN aiter="$(python3 -c 'import importlib.util as u; print(u.find_spec("aiter").origin)')" \ - && ! find "${aiter%/*}" -name '*.so' -exec ldd {} + 2>&1 | grep "version .* not found" -EOF - -echo "--- :docker: Installing ${wheel_url##*/}" -if ! docker buildx build \ - --builder default \ - --provenance=false \ - --progress plain \ - --build-arg "CI_BASE_IMAGE=${ci_base}" \ - --build-arg "AITER_WHEEL_URL=${wheel_url}" \ - -t "${BUILD_TAG}" \ - "${scratch}"; then - fail 11 "the wheel does not install" \ - "Installing or loading ${wheel_url##*/} over ${ci_base} failed; see this step's log." -fi - -docker push "${BUILD_TAG}" -# Handed off as tag@digest: build-test-image.sh rejects a ci_base that is not -# digest-pinned. -digest="$(docker buildx imagetools inspect "${BUILD_TAG}" | awk '$1 == "Digest:" { print $2; exit }')" -if [[ ! "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then - echo "Could not resolve the pushed digest of ${BUILD_TAG}" >&2 - exit 1 -fi -# The vLLM image build uses this as its ci_base. -buildkite-agent meta-data set aiter-nightly-ci-base "${BUILD_TAG}@${digest}" -# The annotation starts with :warning: when the wheel is stale. -style=info -if [[ "${note}" == :warning:* ]]; then - style=warning -fi -echo "${note}" | buildkite-agent annotate --style "${style}" --context aiter-nightly -echo "--- :white_check_mark: ${BUILD_TAG}@${digest}" diff --git a/.buildkite/scripts/rocm/build-ci-base.sh b/.buildkite/scripts/rocm/build-ci-base.sh deleted file mode 100755 index 7d4eb381bc27..000000000000 --- a/.buildkite/scripts/rocm/build-ci-base.sh +++ /dev/null @@ -1,48 +0,0 @@ -#!/usr/bin/env bash -# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base. - -set -euo pipefail - -metadata_get() { - local key="$1" - if command -v buildkite-agent >/dev/null 2>&1; then - buildkite-agent meta-data get "${key}" 2>/dev/null || true - fi -} - -main() { - local base_refreshed="" - local selected_base="" - - # The content identity below describes this checkout. Do not let legacy - # external pipeline settings switch the build to a remote source tree. - export REMOTE_VLLM=0 - unset VLLM_BRANCH - - base_refreshed="$(metadata_get rocm-base-refresh)" - selected_base="$(metadata_get rocm-base-image)" - if [[ -n "${selected_base}" ]]; then - if [[ ! "${selected_base}" =~ @sha256:[0-9a-f]{64}$ ]]; then - echo "Selected ROCm base handoff is not digest-pinned: ${selected_base}" >&2 - return 1 - fi - export BASE_IMAGE="${selected_base}" - echo "Using selected ROCm base image for ci_base: ${BASE_IMAGE}" - elif [[ "${BUILDKITE:-false}" == "true" ]]; then - echo "Required ROCm base handoff metadata is missing" >&2 - return 1 - fi - - if [[ "${base_refreshed}" == "1" ]]; then - export CI_BASE_PUSH_STABLE_TAG - - CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)" - CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}" - - echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}" - fi - - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps -} - -main "$@" diff --git a/.buildkite/scripts/rocm/build-config.sh b/.buildkite/scripts/rocm/build-config.sh deleted file mode 100644 index f6c6ff56ae38..000000000000 --- a/.buildkite/scripts/rocm/build-config.sh +++ /dev/null @@ -1,85 +0,0 @@ -#!/usr/bin/env bash -# Shared Dockerfile selection for the normal AMD CI build chain. -# CI_ROCM_DOCKERFILE_BASE and CI_ROCM_DOCKERFILE accept repository-relative paths. -# For the legacy ROCm 7.2 stack, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rocm_72_base -# and CI_ROCM_DOCKERFILE=docker/Dockerfile.rocm_72. -# Custom recipes must preserve the CI targets and artifact layout. Additional -# COPY inputs also need the corresponding *_CONTENT_FILES overrides. - -rocm_base_dockerfile() { - local default="docker/Dockerfile.rocm_base" - printf '%s\n' "${CI_ROCM_DOCKERFILE_BASE:-${ROCM_BASE_DOCKERFILE:-${default}}}" -} - -rocm_ci_dockerfile() { - local default="docker/Dockerfile.rocm" - printf '%s\n' "${CI_ROCM_DOCKERFILE:-${CI_BASE_DOCKERFILE:-${default}}}" -} - -using_custom_rocm_dockerfiles() { - [[ "$(rocm_base_dockerfile)" != "docker/Dockerfile.rocm_base" \ - || "$(rocm_ci_dockerfile)" != "docker/Dockerfile.rocm" ]] -} - -rocm_dockerfile_stages() { - awk ' - toupper($1) == "FROM" { - for (idx = 2; idx < NF; idx++) { - if (tolower($idx) == "as") print tolower($(idx + 1)) - } - } - ' "$1" -} - -validate_rocm_dockerfile() { - local dockerfile="$1" - local stage="" - local stages="" - shift - - case "${dockerfile}" in - /*|../*|*/../*|*/..|*[[:space:]]*) - echo "ROCm Dockerfile must be a repository-relative path without '..' or whitespace: ${dockerfile}" >&2 - return 2 - ;; - esac - if [[ ! -f "${dockerfile}" ]]; then - echo "ROCm Dockerfile not found: ${dockerfile}" >&2 - return 2 - fi - if (($#)) && ! awk ' - toupper($1) == "FROM" { - named = 0 - for (idx = 2; idx < NF; idx++) { - if (tolower($idx) == "as") named = 1 - } - if (!named) exit 1 - } - ' "${dockerfile}"; then - echo "ROCm Dockerfile must name every build stage: ${dockerfile}" >&2 - return 2 - fi - stages=$(rocm_dockerfile_stages "${dockerfile}") - for stage in "$@"; do - if ! grep -Fxq -- "${stage}" <<< "${stages}"; then - echo "ROCm Dockerfile ${dockerfile} is missing required stage: ${stage}" >&2 - return 2 - fi - done -} - -configure_rocm_build() { - ROCM_BASE_DOCKERFILE=$(rocm_base_dockerfile) - CI_BASE_DOCKERFILE=$(rocm_ci_dockerfile) - export ROCM_BASE_DOCKERFILE CI_BASE_DOCKERFILE - using_custom_rocm_dockerfiles || return 0 - - # Custom stacks may supply dependencies without separate build stages. - export ROCM_DEP_CACHE_EXPORT_MODE=never - - # Preserve the standard runtime tags, which are unique to each build, - # while preventing experimental builds from promoting stable images. - export ROCM_BASE_PUSH_STABLE_TAG=0 - export CI_BASE_PUSH_STABLE_TAG=0 - echo "Custom ROCm Dockerfiles: ${ROCM_BASE_DOCKERFILE} and ${CI_BASE_DOCKERFILE}" -} diff --git a/.buildkite/scripts/rocm/build-test-image.sh b/.buildkite/scripts/rocm/build-test-image.sh deleted file mode 100755 index 69c4f77d20ca..000000000000 --- a/.buildkite/scripts/rocm/build-test-image.sh +++ /dev/null @@ -1,77 +0,0 @@ -#!/usr/bin/env bash -# Build the ROCm CI test image or wheel artifact. -# -# When base selection changes, build the full image so downstream ROCm tests -# validate the selected base -> ci_base -> CI image chain. - -set -euo pipefail - -metadata_get() { - local key="$1" - if command -v buildkite-agent >/dev/null 2>&1; then - buildkite-agent meta-data get "${key}" 2>/dev/null || true - fi -} - -load_digest_handoff() { - local metadata_key="$1" - local env_name="$2" - local description="$3" - local image_ref="" - - image_ref="$(metadata_get "${metadata_key}")" - if [[ -z "${image_ref}" ]]; then - return 1 - fi - if [[ ! "${image_ref}" =~ @sha256:[0-9a-f]{64}$ ]]; then - echo "${description} is not digest-pinned: ${image_ref}" >&2 - return 1 - fi - - printf -v "${env_name}" '%s' "${image_ref}" - export "${env_name?}" - echo "Using ${description}: ${image_ref}" -} - -main() { - local base_refreshed=0 - - # This job always builds the checked-out commit. Some externally generated - # pipeline templates still inject remote-fetch settings; do not let those - # settings make the source identity commit-specific or bypass local edits. - export REMOTE_VLLM=0 - unset VLLM_BRANCH - - if ! load_digest_handoff \ - rocm-ci-base-image CI_BASE_IMAGE "ROCm ci_base handoff"; then - if [[ "${BUILDKITE:-false}" == "true" ]]; then - echo "Required ROCm ci_base handoff metadata is missing or invalid" >&2 - return 1 - fi - echo "No ROCm ci_base handoff metadata found; using the local default" - fi - - if ! load_digest_handoff \ - rocm-base-image BASE_IMAGE "ROCm base handoff"; then - if [[ "${BUILDKITE:-false}" == "true" ]]; then - echo "Required ROCm base handoff metadata is missing or invalid" >&2 - return 1 - fi - echo "No ROCm base handoff metadata found; using the local default" - fi - - if [[ "$(metadata_get rocm-base-refresh)" == "1" ]]; then - echo "The selected ROCm base differs from the current stable base" - base_refreshed=1 - fi - - if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then - echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only" - IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts - return - fi - - bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel -} - -main "$@" diff --git a/.buildkite/scripts/rocm/build.sh b/.buildkite/scripts/rocm/build.sh new file mode 100755 index 000000000000..b61dcb913f87 --- /dev/null +++ b/.buildkite/scripts/rocm/build.sh @@ -0,0 +1,330 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Build and promote ROCm images via docker/docker-bake-rocm.hcl. +# Pushes only to $ROCM_CI_REPO and the BuildKit cache that goes with it. +# +# Usage: build.sh +# +# Modes: +# keys Print current Ks/Kr keys and whether their images exist (on-call debug). +# base Ensure stack, runtime and test-deps images exist; write their refs to Buildkite metadata. +# ci Build and push this build's per-commit test image. +# release Release image + wheels from source into the release repo (release pipeline, AWS agents). +# release-base Only the release stack and runtime (a serialized step, so one run builds a new key). +# +# Stack: rocm (TheRock ROCm 10). Force a stack rebuild: bump ARG REFRESH +# in docker/Dockerfile.rocm_base +# +# Trust: ROCM_CI_REPO is set by ci-infra (CI) or release-pipeline.yaml (release). +# Trusted main builds: rocm/vllm-ci +# PR / fork builds: rocm/vllm-dev (pr- tags and caches; trusted builds never read it) +# Other untrusted: rocm/vllm-dev (br--- tags and caches: nightlies, branch builds) +# Release pipeline: its ECR repo (never reads rocm/vllm-ci or its cache: it builds the +# stack and runtime from source, with only its own ECR cache; +# csrc objects come from the shared S3 sccache, as in CUDA) +# ROCM_CI_TRUSTED_REPO / ROCM_CI_CACHE_REPO only point a local run at a sandbox registry. +set -euo pipefail + +# --------------------------------------------------------------------------- +# Arguments and trust +# --------------------------------------------------------------------------- +STACK=$1 +MODE=$2 +REPO=${ROCM_CI_REPO:?ROCM_CI_REPO must be set (by ci-infra or release-pipeline.yaml)} +TRUSTED=${ROCM_CI_TRUSTED_REPO:-rocm/vllm-ci} + +export STACK PUSH=true +export CACHE=${ROCM_CI_CACHE_REPO:-rocm/vllm-ci-cache} # every CI build reads the trusted cache +export CACHE_TO=$CACHE: # ref prefix trusted builds write (HCL appends -) +export BUILDKIT_PROGRESS=${BUILDKIT_PROGRESS:-plain} # TTY progress redraws bloat CI logs + +# PR builds write to a per-PR cache and prefix their image tags. The release pipeline's agents +# cannot write rocm/vllm-ci or its cache: release writes its images and its own cache to $REPO, +# so the next release reuses this one's all-arch compile. +PFX="" +if [[ $MODE == release* ]]; then + CACHE="" CACHE_TO=$REPO:cache- # and reads only that cache: CI agents also run PR code +elif [[ $REPO != "$TRUSTED" ]]; then + if [[ ${BUILDKITE_PULL_REQUEST:-false} != false ]]; then + PFX="pr$BUILDKITE_PULL_REQUEST-" + else # not a PR (AITER nightly, branch or manual builds): one namespace per branch + br=${BUILDKITE_BRANCH:-local} # the hash keeps e.g. feat/x and feat-x apart + PFX="br-$(printf %s "$br" | tr -c 'A-Za-z0-9.-' - | cut -c1-40)-$(printf %s "$br" | sha256sum | cut -c1-8)-" + fi + CACHE_TO=$REPO:cache-$PFX + # Only the compile stages' cache: everything else is cheap or in the trusted cache. + export CACHE_CI=false +fi + +# Stack Dockerfile — literal path used by ci-infra ci_selector (images.py). +case $STACK in + rocm) SDF=docker/Dockerfile.rocm_base ;; + *) echo "Unknown stack: $STACK" >&2; exit 2 ;; +esac + +# --------------------------------------------------------------------------- +# BuildKit builder +# --------------------------------------------------------------------------- +# Prefer the host buildkitd (it has its own GC policy and cache). Fall back +# to a docker-container builder if no socket is present. Release RUN steps must reach EC2 IMDS +# for sccache's S3 credentials: their fallback builder uses the host network, under its own name. +B=vllm-rocm net=() +[[ $MODE == release* ]] && B=vllm-rocm-release net=(--driver-opt network=host) +if ! docker buildx inspect "$B" >/dev/null 2>&1; then + if [[ -S /run/buildkit/buildkitd.sock ]]; then + drv=(--driver remote unix:///run/buildkit/buildkitd.sock) + else + drv=(--driver docker-container "${net[@]}" \ + ${BUILDKITD_CONFIG:+--buildkitd-config "$BUILDKITD_CONFIG"}) + fi + # || inspect: tolerate a race where another job created it between the check and create. + docker buildx create --name "$B" "${drv[@]}" >/dev/null \ + || docker buildx inspect "$B" >/dev/null +fi +export BUILDX_BUILDER=$B +docker buildx version + +# --------------------------------------------------------------------------- +# Build context — an export of HEAD +# --------------------------------------------------------------------------- +# Bake builds .rocm-src, not the checkout: the content keys below hash HEAD, and BuildKit keys +# COPY on file modes, which drift on shared agents. umask 022 gives every agent git's modes. +[[ $(git rev-parse HEAD) == $(git rev-parse "${BUILDKITE_COMMIT:-HEAD}^{commit}") ]] \ + || { echo "HEAD is not BUILDKITE_COMMIT ($BUILDKITE_COMMIT)" >&2; exit 1; } +umask 022 +rm -rf .rocm-src && mkdir .rocm-src +git archive HEAD | tar -x --no-same-permissions -C .rocm-src -f - +export CONTEXT=.rocm-src + +# --------------------------------------------------------------------------- +# Git tags — required by the version stage +# --------------------------------------------------------------------------- +# setuptools_scm needs a reachable v* tag. Agents clone without tags, partially +# and with their objects in a host mirror (--reference) that the build context +# cannot see; --force updates tags that moved (e.g. a re-tag after a release +# candidate). The context has no .git: the version stage reads .git_archival.txt. +git fetch -q --force --tags --filter=tree:0 origin +git describe --tags --match 'v[0-9]*' >/dev/null 2>&1 \ + || git fetch -q --unshallow --filter=tree:0 origin +printf 'node: %s\nnode-date: %s\ndescribe-name: %s\n' "$(git rev-parse HEAD)" \ + "$(git log -1 --format=%cI)" "$(git describe --tags --long --abbrev=40 --match 'v[0-9]*')" \ + > .rocm-src/.git_archival.txt + +# --------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------- + +# digest +# Print the sha256 digest of . Returns 1 if the tag does not exist. +# The "Digest:" line of the plain output is parsed and validated (as the old CI +# scripts did): --format output and stray stderr are not safe to splice into a ref. +# Any other registry error, after retries, exits 75 (the steps retry it); in a $() +# subshell that becomes the substitution's status, which ensure() and set -e pass on. +digest() { + local out rc d try + for try in 1 2 3; do + rc=0 + out=$(docker buildx imagetools inspect "$1" 2>&1) || rc=$? + d=$(awk '$1 == "Digest:" { print $2; exit }' <<<"$out") + if ((rc == 0)) && [[ $d =~ ^sha256:[0-9a-f]{64}$ ]]; then echo "$d"; return 0; fi + if ((rc != 0)) && grep -qiE 'manifest unknown|no such manifest|name unknown|not found|does not exist' <<<"$out"; then + return 1 + fi + ((try == 3)) || sleep 5 + done + echo "registry error for $1 (exit $rc): $out" >&2 + exit 75 +} + +# label — read one image label +label() { docker buildx imagetools inspect "$1" \ + --format "{{index .Image.Config.Labels \"$2\"}}"; } + +# bake — invoke docker buildx bake with the project HCL +bake() { docker buildx bake -f docker/docker-bake-rocm.hcl "$@"; } + +# meta [value] — Buildkite build metadata +meta() { buildkite-agent meta-data "$@"; } + +# pushed — the digest that bake pushed for +pushed() { + local d + d=$(awk -F'"' -v t="$2" '/^ "/ { k = $2 } k == t && $2 == "containerimage.digest" { print $4 }' "$1") + [[ $d =~ ^sha256:[0-9a-f]{64}$ ]] && echo "$d" +} + +# --------------------------------------------------------------------------- +# Content keys +# --------------------------------------------------------------------------- +# Ks (stack key): hash of the stack Dockerfile + therock_wheels.py. +# A Ks cache miss triggers a full stack rebuild (hours). +# Kr (runtime key): Ks + the portion of Dockerfile.rocm above the per-commit +# marker + all pip requirement files. +# A Kr cache miss triggers a runtime rebuild (minutes). + +_trw=$( [[ $STACK == rocm ]] && echo tools/vllm-rocm/therock_wheels.py || true ) +# The layer format (compression, and mtimes as built: see the bake file) is part of every image's +# identity: an image left in an older format would hand later builds a second copy of the same layers. +Ks=$( { echo "layers: zstd-3, mtimes kept"; git ls-tree HEAD -- "$SDF" ${_trw:+"$_trw"}; } | sha256sum | cut -c1-16) +Kr=$( + { + echo "$Ks" + sed '/^# ==== per-commit stages/q' .rocm-src/docker/Dockerfile.rocm + git ls-tree HEAD -- \ + requirements/common.txt \ + requirements/rocm.txt \ + requirements/test/rocm.txt \ + rust-toolchain.toml \ + tests/vllm_test_utils \ + tools/install_torchcodec_rocm.sh + } | sha256sum | cut -c1-16 +) +echo "build.sh $STACK $MODE: repo=$REPO trusted=$TRUSTED Ks=$Ks Kr=$Kr" + +# --------------------------------------------------------------------------- +# ensure [] +# --------------------------------------------------------------------------- +# Resolve a content-keyed image without building if it already exists. +# Resolution order: +# 1. Trusted registry (rocm/vllm-ci:) — safe for any CI build to reuse. Not for release: +# CI agents also run PR code, so release never ships an image it did not build itself. +# 2. This repo's own tag ($REPO:$PFX) — PR builds and the release pipeline write here. +# 3. Neither found: build via bake and push. +# A means the image only counts as present when the sibling +# tag also exists (used by stack: must have both the image and -wheels); then +# both pinned refs are printed, space-separated. +ensure() { + local target=$1 tag=$2 sibling=${3:-} ref d ds refs rc mf + refs=("$REPO:$PFX$tag") + [[ $MODE == release* || $REPO == "$TRUSTED" ]] || refs=("$TRUSTED:$tag" "${refs[@]}") + + for ref in "${refs[@]}"; do + if d=$(digest "$ref") && { [[ -z $sibling ]] || ds=$(digest "$ref$sibling"); }; then + echo "Reusing $ref@$d" >&2 + echo "$ref@$d${sibling:+ $ref$sibling@$ds}" + return 0 + else + rc=$?; ((rc == 1)) || return $rc # a registry error is not a miss + fi + done + + # Neither copy found — build and push. + ref=$REPO:$PFX$tag + echo "--- :docker: Building $target as $ref" >&2 + # Explicit checks: $(ensure ...) runs in a subshell, where set -e does not apply. + mf=$(mktemp) + TAG=$ref bake --metadata-file "$mf" "$target" >&2 || { rm -f "$mf"; echo "building $ref failed" >&2; return 1; } + # The digests this bake pushed (base: stack-, stack-wheels), not the tags', which by now + # may hold a concurrent build's push of the same key. + if ! d=$(pushed "$mf" "${target/#base/stack-$STACK}") \ + || { [[ -n $sibling ]] && ! ds=$(pushed "$mf" stack-wheels); }; then + rm -f "$mf"; echo "bake recorded no digest for $ref" >&2; return 1 + fi + rm -f "$mf" + echo "$ref@$d${sibling:+ $ref$sibling@$ds}" +} + +# load — read stack/runtime refs from Buildkite metadata (written by base mode). +load() { + STACK_IMAGE=$(meta get "rocm-$STACK-stack") + STACK_WHEELS=$(meta get "rocm-$STACK-stack-wheels") + RUNTIME=$(meta get "rocm-$STACK-runtime") + TEST_DEPS=$(meta get "rocm-$STACK-test-deps") + export STACK_IMAGE STACK_WHEELS RUNTIME TEST_DEPS + for ref in "$STACK_IMAGE" "$STACK_WHEELS" "$RUNTIME" "$TEST_DEPS"; do + [[ $ref =~ @sha256:[0-9a-f]{64}$ ]] || { echo "not a pinned image ref in the metadata: '$ref'" >&2; exit 1; } + done +} + +# --------------------------------------------------------------------------- +# Mode dispatch +# --------------------------------------------------------------------------- +case $MODE in + + keys) + # Print current content keys and registry status (useful for on-call triage). + for t in "stack-$STACK-$Ks" "runtime-$STACK-$Kr" "test-deps-$STACK-$Kr"; do + d=$(digest "$TRUSTED:$t") || { (($? == 1)) || exit 75; d=MISSING; } + echo "$t $d" + done + ;; + + base) + # Ensure the stack and runtime images exist; record their digests in metadata + # so the ci mode can pull them without re-resolving. + # The "base" bake group builds the stack and its -wheels sibling together. + STACK_IMAGE=$(ensure base "stack-$STACK-$Ks" -wheels) + STACK_WHEELS=${STACK_IMAGE#* } STACK_IMAGE=${STACK_IMAGE%% *} + export STACK_IMAGE + RUNTIME=$(ensure runtime "runtime-$STACK-$Kr") + export RUNTIME + TEST_DEPS=$(ensure test-deps "test-deps-$STACK-$Kr") + + meta set "rocm-$STACK-stack" "$STACK_IMAGE" + meta set "rocm-$STACK-runtime" "$RUNTIME" + meta set "rocm-$STACK-test-deps" "$TEST_DEPS" + meta set "rocm-$STACK-stack-wheels" "$STACK_WHEELS" + echo "stack=$STACK_IMAGE runtime=$RUNTIME test-deps=$TEST_DEPS stack-wheels=$STACK_WHEELS" + ;; + + ci) + # Build and push the per-commit test image for GPU test steps. + load + export TAG=$REPO:build-$BUILDKITE_BUILD_ID + # AITER nightly (AITER_NIGHTLY=1 from the schedule): the test stage installs the newest + # AITER nightly that fits the image; an AITER_SPEC in the environment pins one instead. + rm -rf .rocm-wheel .rocm-kernel-map + bake ci + # For the python-only install test (run-amd-test.sh downloads it); not in the image. + buildkite-agent artifact upload ".rocm-wheel/vllm-*.whl" + # CI test selection (VLLM_KERNEL_SYMBOL_MAP=1, set by ci-infra): the kernel symbol map of this + # build's csrc, from the build cache. Optional: a missing map never fails the image build. + if [[ ${VLLM_KERNEL_SYMBOL_MAP:-0} == 1 ]]; then + if bake kernel-symbol-map && [[ -s .rocm-kernel-map/kernel_symbol_map.rocm.json.gz ]] \ + && (cd .rocm-kernel-map && buildkite-agent artifact upload kernel_symbol_map.rocm.json.gz); then + : + else + echo "ROCm kernel symbol map export or upload failed; continuing without it" >&2 + fi + fi + ;; + + release|release-base) + # Build and push the release image + wheels into $REPO. The stack and runtime come from this + # repo only (built here once per key, from the same Dockerfiles and keys CI tests), never from + # rocm/vllm-ci or its cache. + STACK_IMAGE=$(ensure base "stack-$STACK-$Ks" -wheels) + STACK_WHEELS=${STACK_IMAGE#* } STACK_IMAGE=${STACK_IMAGE%% *} + export STACK_IMAGE # before ensure runtime: a runtime built here must use this stack + RUNTIME=$(ensure runtime "runtime-$STACK-$Kr") + [[ $MODE == release ]] || exit 0 + export STACK_WHEELS RUNTIME + export TAG=$REPO:$STACK-$BUILDKITE_COMMIT + + # Resolve the vLLM version from the commit (setuptools_scm). + PUSH=false bake version + VLLM_VERSION=$(cat .rocm-version/vllm-version.txt) + export VLLM_VERSION + + # Skip if this exact version was already released (handles rc → final re-tags). + if digest "$TAG" >/dev/null \ + && [[ $(label "$TAG" org.opencontainers.image.version) == "$VLLM_VERSION" ]]; then + echo "Release image $TAG already exists with version $VLLM_VERSION — skipping." >&2 + exit 0 + fi + + # Wheels first: their all-arch csrc compile is the long part, and vllm-openai reuses it. + bake wheels openai-smoke + # Alias the runtime as the -base tag (the base-nightly consumer reference). + docker buildx imagetools create -t "$TAG-base" "$RUNTIME" + # Push the vllm-openai tag last — its presence marks the set as complete. + bake vllm-openai + ;; + + *) + echo "Unknown mode: $MODE" >&2 + exit 2 + ;; +esac diff --git a/.buildkite/scripts/rocm/publish.sh b/.buildkite/scripts/rocm/publish.sh new file mode 100755 index 000000000000..3593d00ada39 --- /dev/null +++ b/.buildkite/scripts/rocm/publish.sh @@ -0,0 +1,208 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Promote ROCm release artifacts without rebuilding anything. +# The release pipeline's build step (build.sh release) pushed the source images to +# ECR; this script only copies digests and uploads wheels. +# +# Usage: publish.sh +# +# Modes: +# wheels Extract the vLLM wheel from the -wheels image and upload to S3. +# images Tag the release image and its runtime base for their variant in ECR. +# With NIGHTLY=1, also push the nightly tags to Docker Hub. +# +# Required environment variables: +# BUILDKITE_COMMIT — the commit being released +# ECR — ECR repository prefix (release-pipeline.yaml's ROCM_CI_REPO) +# +# Optional: +# HUB — Docker Hub repo for nightly images (default: vllm/vllm-openai-rocm) +# NIGHTLY — set to 1 to push nightly Hub tags (default: 0) +# DOCKERHUB_USERNAME — Hub username for cleanup script (default: vllmbot) +set -euo pipefail + +# --------------------------------------------------------------------------- +# Arguments +# --------------------------------------------------------------------------- +STACK=$1 +MODE=$2 +COMMIT=${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set} + +ECR=${ECR:?ECR must be set (by release-pipeline.yaml)} +SRC=$ECR:${STACK}-${COMMIT} # build.sh release, in the same build +HUB=${HUB:-vllm/vllm-openai-rocm} + +# --------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------- + +# retry — registry lookups: these steps have no retry rules +retry() { "$@" || { sleep 5; "$@"; } || { sleep 5; "$@"; }; } + +# pin — resolve to a pinned ref@digest (resolves once, stable for the rest of the script). +# Parses the validated "Digest:" line, as build.sh does. +pin() { + local d + d=$(retry docker buildx imagetools inspect "$1" | awk '$1 == "Digest:" { print $2; exit }') + [[ $d =~ ^sha256:[0-9a-f]{64}$ ]] || { echo "cannot resolve the digest of $1" >&2; return 1; } + echo "$1@$d" +} + +# variant — read the ai.vllm.rocm.variant image label +variant() { + retry docker buildx imagetools inspect "$1" \ + --format '{{index .Image.Config.Labels "ai.vllm.rocm.variant"}}' +} + +# --------------------------------------------------------------------------- +# Resolve source image refs (done once; images step uses them too) +# --------------------------------------------------------------------------- +# Logged in, the ECR lookups below are not subject to the anonymous rate limits. +aws ecr-public get-login-password --region us-east-1 \ + | docker login -u AWS --password-stdin "${ECR%%/*}" +IMG=$(pin "$SRC") +V=$(variant "$IMG") # e.g. rocm100 — used in S3 paths, ECR tags and Hub tags +[[ $V =~ ^rocm[0-9]+$ ]] || { echo "$IMG has no valid ai.vllm.rocm.variant label: '$V'" >&2; exit 1; } +# For the steps after either mode: root-index generation and publish-release-images.sh. +buildkite-agent meta-data set "rocm-$STACK-variant" "$V" + +# --------------------------------------------------------------------------- +# Mode dispatch +# --------------------------------------------------------------------------- +case $MODE in + + # ------------------------------------------------------------------------- + wheels) + # ------------------------------------------------------------------------- + # Extract the vLLM wheel from this commit's -wheels image and upload it + # to S3 so users can pip-install it directly. + # ------------------------------------------------------------------------- + + WHEELS_IMG=$(pin "$SRC-wheels") + + # Export wheel files out of the scratch image into ./dist/ + rm -rf dist + echo "FROM $WHEELS_IMG" | docker buildx build -q -o dist - + + # Sanity: the wheel's variant suffix must match the stack label. + # (setup.py's rocm_version_to_variant produces e.g. rocm100 for 10.0.x) + actual_variant=$(basename dist/vllm-*.whl | grep -oE 'rocm[0-9]+' | tail -1) + [[ $actual_variant == "$V" ]] \ + || { echo "Wheel variant $actual_variant != image variant $V" >&2; exit 1; } + + bash .buildkite/scripts/upload-rocm-wheels.sh "$STACK" dist + + # Buildkite annotation: pip install + direct S3 download commands. + WHL_URL="https://wheels.vllm.ai/rocm/$COMMIT/$V" + buildkite-agent annotate --style success --context "rocm-$STACK-wheels" << EOF +## :package: ROCm \`$V\` wheels published + +| | | +|---|---| +| **Variant** | \`$V\` | +| **Commit** | \`$COMMIT\` | +| **Branch** | \`${BUILDKITE_BRANCH:-?}\` | + +### Install (this commit) +\`\`\`bash +pip install vllm --extra-index-url $WHL_URL +\`\`\` + +### Install (latest nightly) +\`\`\`bash +pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/nightly/$V +\`\`\` + +### Download wheels directly +\`\`\`bash +aws s3 ls s3://${S3_BUCKET:-vllm-wheels}/rocm/$COMMIT/$V-wheels/ +aws s3 cp --recursive --exclude "*" --include "*.whl" s3://${S3_BUCKET:-vllm-wheels}/rocm/$COMMIT/$V-wheels/ . +\`\`\` +EOF + ;; + + # ------------------------------------------------------------------------- + images) + # ------------------------------------------------------------------------- + # Tag the release image and its runtime base for their variant in ECR. + # With NIGHTLY=1, also tag them on Docker Hub and prune old nightly tags. + # ------------------------------------------------------------------------- + + RT=$(pin "$SRC-base") # runtime base alias (published as base-nightly*) + + # --- ECR ----------------------------------------------------------------- + docker buildx imagetools create -t "$ECR:$COMMIT-$V" "$IMG" + docker buildx imagetools create -t "$ECR:$COMMIT-$V-base" "$RT" + + ECR_IMG="$ECR:$COMMIT-$V" + ECR_BASE="$ECR:$COMMIT-$V-base" + + # --- Non-nightly annotation ---------------------------------------------- + if [[ ${NIGHTLY:-0} != 1 ]]; then + buildkite-agent annotate --style success --context "rocm-$STACK-images" << EOF +## :docker: ROCm \`$V\` images published to ECR + +| | | +|---|---| +| **Variant** | \`$V\` | +| **Commit** | \`$COMMIT\` | +| **ECR image** | \`$ECR_IMG\` | +| **ECR base** | \`$ECR_BASE\` | + +\`\`\`bash +docker pull $ECR_IMG +\`\`\` +EOF + exit 0 + fi + + # --- Nightly: Docker Hub ------------------------------------------------- + export DOCKERHUB_USERNAME=${DOCKERHUB_USERNAME:-vllmbot} + + # The default stack (rocm) owns the unqualified :nightly / :base-nightly tags. + # All stacks own the variant-qualified :nightly- tags. + suffixes=("-$V") + [[ $STACK == rocm ]] && suffixes+=("") + + hub_tags="" + for sfx in "${suffixes[@]}"; do + docker buildx imagetools create \ + -t "$HUB:nightly$sfx" -t "$HUB:nightly$sfx-$COMMIT" "$IMG" + docker buildx imagetools create \ + -t "$HUB:base-nightly$sfx" -t "$HUB:base-nightly$sfx-$COMMIT" "$RT" + hub_tags="$hub_tags \`$HUB:nightly$sfx\` \`$HUB:nightly$sfx-$COMMIT\`" + done + + # Prune old nightly tags (keep last 14 builds per tag family). + clean=.buildkite/scripts/cleanup-nightly-builds.sh + bash "$clean" "nightly-$V-" "$HUB" + bash "$clean" "base-nightly-$V-" "$HUB" + if [[ $STACK == rocm ]]; then + bash "$clean" "nightly-" "$HUB" "nightly-rocm" + bash "$clean" "base-nightly-" "$HUB" "base-nightly-rocm" + fi + + buildkite-agent annotate --style success --context "rocm-$STACK-images" << EOF +## :docker: ROCm \`$V\` nightly images published + +| | | +|---|---| +| **Variant** | \`$V\` | +| **Commit** | \`$COMMIT\` | +| **ECR image** | \`$ECR_IMG\` | +| **ECR base** | \`$ECR_BASE\` | +| **Docker Hub** |$hub_tags | + +\`\`\`bash +docker pull $HUB:nightly-$V +\`\`\` +EOF + ;; + + *) + echo "Unknown mode: $MODE" >&2 + exit 2 + ;; +esac diff --git a/.buildkite/scripts/rocm/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh deleted file mode 100755 index 10e594577a87..000000000000 --- a/.buildkite/scripts/rocm/refresh-base-image.sh +++ /dev/null @@ -1,613 +0,0 @@ -#!/usr/bin/env bash -# Select the content-addressed ROCm base for this checkout, building it only -# when the exact content ref is absent. - -set -euo pipefail - -# shellcheck source=.buildkite/scripts/rocm/build-config.sh -source "$(dirname "${BASH_SOURCE[0]}")/build-config.sh" -configure_rocm_build - -DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}" -BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}" -CACHE_REPO="${ROCM_BASE_CACHE_REPO:-${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}}" -BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}" -DEFAULT_ROCM_BASE_METADATA_VERSION="2" -DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}" - -ROCM_BASE_LAYER_CACHE_REF="" -BASE_CACHE_PREFIX="rocm-base" -if using_custom_rocm_dockerfiles; then - BASE_CACHE_PREFIX="rocm-base-$(printf '%s\n%s\n' \ - "${ROCM_BASE_DOCKERFILE}" "${CI_BASE_DOCKERFILE}" | sha256sum | cut -c1-12)" -fi -ROCM_BASE_TRUSTED_LAYER_CACHE_REF="${CACHE_REPO}:${BASE_CACHE_PREFIX}-main" -ROCM_BASE_STABLE_TAG_UPDATED=0 -declare -a ROCM_BASE_CACHE_ARGS=() - -metadata_set() { - local key="$1" - local value="$2" - - [[ -n "${value}" ]] || return 0 - if command -v buildkite-agent >/dev/null 2>&1; then - buildkite-agent meta-data set "${key}" "${value}" - elif [[ "${BUILDKITE:-false}" == "true" ]]; then - echo "buildkite-agent is required to publish ROCm base metadata" >&2 - return 1 - fi -} - -compute_content_hash() { - local path="" - local file="" - - for path in "$@"; do - if [[ -d "${path}" ]]; then - while IFS= read -r -d '' file; do - printf 'file:%s\n' "${file}" - sha256sum "${file}" - done < <(find "${path}" -type f -print0 | sort -z) - elif [[ -f "${path}" ]]; then - printf 'file:%s\n' "${path}" - sha256sum "${path}" - else - printf 'missing:%s\n' "${path}" - fi - done | sha256sum | cut -d' ' -f1 -} - -clean_docker_tag() { - local input="$1" - echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128 -} - -tag_component() { - local input="$1" - local max_chars="${2:-24}" - - clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}" -} - -normalize_repo_slug() { - local repo_slug="${1:-}" - - repo_slug="${repo_slug%/}" - repo_slug="${repo_slug%.git}" - repo_slug="${repo_slug#https://github.com/}" - repo_slug="${repo_slug#http://github.com/}" - repo_slug="${repo_slug#ssh://git@github.com/}" - repo_slug="${repo_slug#git@github.com:}" - repo_slug="${repo_slug#github.com/}" - printf '%s\n' "${repo_slug}" -} - -is_trusted_main_build() { - local actual_repo="" - local trusted_repo="" - - [[ "${BUILDKITE:-false}" == "true" ]] || return 1 - [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" ]] || return 1 - [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] || return 1 - actual_repo=$(normalize_repo_slug "${BUILDKITE_REPO:-}") - trusted_repo=$(normalize_repo_slug \ - "${ROCM_BASE_STABLE_REPO_SLUG:-vllm-project/vllm}") - [[ -n "${actual_repo}" && "${actual_repo}" == "${trusted_repo}" ]] -} - -rocm_base_layer_cache_scope() { - local pull_request="${BUILDKITE_PULL_REQUEST:-false}" - local branch="${BUILDKITE_PULL_REQUEST_HEAD_BRANCH:-${BUILDKITE_BRANCH:-local}}" - local identity="" - local repo_slug="" - - if is_trusted_main_build; then - printf 'main\n' - return 0 - fi - if [[ "${pull_request}" != "false" && -n "${pull_request}" ]]; then - repo_slug=$(normalize_repo_slug "${BUILDKITE_REPO:-local}") - identity=$(printf '%s\n' "${repo_slug:-local}" | sha256sum | cut -c1-12) - printf 'pr-%s-%s\n' \ - "$(tag_component "${pull_request}" 32)" "${identity}" - return 0 - fi - - identity=$(printf '%s\n%s\n' "${BUILDKITE_REPO:-local}" "${branch}" \ - | sha256sum | cut -c1-12) - printf 'preview-%s-%s\n' "$(tag_component "${branch}" 24)" "${identity}" -} - -configure_rocm_base_layer_cache() { - local scope="" - - ROCM_BASE_CACHE_ARGS=() - if [[ "${ROCM_BASE_NO_CACHE:-0}" == "1" ]]; then - ROCM_BASE_CACHE_ARGS+=(--no-cache) - ROCM_BASE_LAYER_CACHE_REF="disabled" - return 0 - fi - - scope=$(rocm_base_layer_cache_scope) - ROCM_BASE_LAYER_CACHE_REF="${CACHE_REPO}:${BASE_CACHE_PREFIX}-${scope}" - ROCM_BASE_CACHE_ARGS+=( - --cache-from "type=registry,ref=${ROCM_BASE_LAYER_CACHE_REF}" - ) - if [[ "${ROCM_BASE_LAYER_CACHE_REF}" != \ - "${ROCM_BASE_TRUSTED_LAYER_CACHE_REF}" ]]; then - ROCM_BASE_CACHE_ARGS+=( - --cache-from "type=registry,ref=${ROCM_BASE_TRUSTED_LAYER_CACHE_REF}" - ) - fi - ROCM_BASE_CACHE_ARGS+=( - --cache-to \ - "type=registry,ref=${ROCM_BASE_LAYER_CACHE_REF},mode=max,ignore-error=true" - ) -} - -extract_arg_default() { - local arg_name="$1" - - sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ - "${DOCKERFILE}" | head -1 -} - -resolve_image_digest() { - local image_ref="$1" - local attempts="${ROCM_IMAGE_DIGEST_ATTEMPTS:-4}" - local delay_secs="${ROCM_IMAGE_DIGEST_RETRY_DELAY:-2}" - local attempt=0 - local output="" - local digest="" - local status=0 - - if [[ "${image_ref}" =~ @(sha256:[0-9a-f]{64})$ ]]; then - printf '%s\n' "${BASH_REMATCH[1]}" - return 0 - fi - if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \ - || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then - echo "Invalid image digest retry configuration" >&2 - return 1 - fi - - for ((attempt = 1; attempt <= attempts; attempt++)); do - status=0 - output=$(docker buildx imagetools inspect "${image_ref}" 2>&1) || status=$? - digest=$(awk '$1 == "Digest:" { print $2; exit }' <<< "${output}") - if ((status == 0)) && [[ "${digest}" =~ ^sha256:[0-9a-f]{64}$ ]]; then - printf '%s\n' "${digest}" - return 0 - fi - if ((attempt < attempts)); then - printf \ - 'Image digest lookup failed for %s (%d/%d, status %d); retrying\n' \ - "${image_ref}" "${attempt}" "${attempts}" "${status}" >&2 - sleep "${delay_secs}" - fi - done - - printf 'Failed to resolve digest for %s (status %d)\n%s\n' \ - "${image_ref}" "${status}" "${output:-}" >&2 - return 1 -} - -trusted_base_content_ref() { - local base_hash="$1" - local metadata_version="$2" - - printf '%s:base-v%s-%s\n' "${BASE_REPO}" "${metadata_version}" "${base_hash}" -} - -scoped_base_content_ref() { - local base_hash="$1" - local metadata_version="$2" - - printf '%s:base-v%s-preview-%s\n' \ - "${BASE_REPO}" "${metadata_version}" "${base_hash}" -} - -find_matching_base_content_ref() { - local expected_hash="$1" - local expected_version="$2" - local image_ref="" - local digest="" - local immutable_ref="" - local label_values="" - local remote_hash="" - local remote_version="" - local attempts="${ROCM_IMAGE_DIGEST_ATTEMPTS:-4}" - local delay_secs="${ROCM_IMAGE_DIGEST_RETRY_DELAY:-2}" - local attempt=0 - local inspect_status=0 - shift 2 - - if [[ ! "${attempts}" =~ ^[1-9][0-9]*$ \ - || ! "${delay_secs}" =~ ^[0-9]+$ ]]; then - echo "Invalid image identity retry configuration" >&2 - return 1 - fi - - for image_ref in "$@"; do - if ! digest=$(resolve_image_digest "${image_ref}"); then - continue - fi - immutable_ref="${image_ref}@${digest}" - label_values="" - for ((attempt = 1; attempt <= attempts; attempt++)); do - inspect_status=0 - label_values=$(docker buildx imagetools inspect "${immutable_ref}" \ - --format '{{ index .Image.Config.Labels "vllm.rocm_base.content_hash" }} {{ index .Image.Config.Labels "vllm.rocm_base.metadata_version" }}' \ - 2>/dev/null) || inspect_status=$? - remote_hash="" - remote_version="" - if ((inspect_status == 0)) && [[ -n "${label_values}" ]]; then - read -r remote_hash remote_version _ <<< "${label_values}" - if [[ "${remote_hash}" == "${expected_hash}" \ - && "${remote_version}" == "${expected_version}" ]]; then - printf '%s\n' "${immutable_ref}" - return 0 - fi - fi - if ((attempt < attempts)); then - printf \ - 'ROCm base identity lookup incomplete or mismatched for %s (%d/%d); retrying\n' \ - "${immutable_ref}" "${attempt}" "${attempts}" >&2 - sleep "${delay_secs}" - fi - done - if ((inspect_status == 0)) \ - && [[ -n "${remote_hash}" && -n "${remote_version}" ]]; then - echo "ROCm base content ref has unexpected identity: ${immutable_ref}" >&2 - echo " expected hash/version: ${expected_hash} / ${expected_version}" >&2 - echo " found hash/version: ${remote_hash} / ${remote_version}" >&2 - # A content-addressed ref with the wrong identity is corrupt. The - # stable alias is allowed to describe older content and is only a - # migration/recovery candidate. - if [[ "${image_ref}" != "${BASE_REPO}:base" ]]; then - return 2 - fi - fi - echo "Failed to read complete ROCm base identity labels: ${immutable_ref}" >&2 - done - return 1 -} - -rocm_version_from_base_image() { - local base_image="$1" - local version="" - - version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")" - tag_component "${version:-${base_image}}" 16 -} - -should_push_stable_tag() { - if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then - return 1 - fi - - if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then - return 1 - fi - - is_trusted_main_build -} - -trusted_main_tip_matches_build() { - local branch="${ROCM_BASE_STABLE_BRANCH:-main}" - local build_commit="${BUILDKITE_COMMIT:-}" - local remote_tip="" - - is_trusted_main_build || return 1 - if [[ ! "${build_commit}" =~ ^[0-9a-fA-F]{40}$ ]]; then - echo "Skipping ROCm stable tag: Buildkite commit is missing or invalid" >&2 - return 1 - fi - remote_tip=$(git ls-remote --exit-code "${BUILDKITE_REPO}" \ - "refs/heads/${branch}" 2>/dev/null | awk 'NR == 1 { print $1 }') - if [[ ! "${remote_tip}" =~ ^[0-9a-fA-F]{40}$ ]]; then - echo "Skipping ROCm stable tag: could not resolve remote ${branch} tip" >&2 - return 1 - fi - if [[ "${remote_tip,,}" != "${build_commit,,}" ]]; then - echo "Skipping ROCm stable tag: ${branch} changed after ${build_commit} (${remote_tip})" >&2 - return 1 - fi - return 0 -} - -promote_stable_base_tag() { - local source_ref="$1" - local stable_tag="$2" - - ROCM_BASE_STABLE_TAG_UPDATED=0 - should_push_stable_tag && trusted_main_tip_matches_build || return 0 - - docker buildx imagetools create --prefer-index=false \ - -t "${stable_tag}" "${source_ref}" - ROCM_BASE_STABLE_TAG_UPDATED=1 -} - -setup_builder() { - echo "--- :buildkite: Setting up buildx builder for ROCm base" - if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then - docker buildx use "${BUILDER_NAME}" - else - docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use - fi - docker buildx inspect --bootstrap -} - -compute_base_content_hash() { - local use_sccache="$1" - local base_image_digest="$2" - local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}" - local sccache_arg="" - local sccache_value="" - local -a content_paths=() - - read -r -a content_paths <<< "${content_files}" - - { - printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")" - printf 'dockerfile:%s\n' "${DOCKERFILE}" - printf 'base-image-digest:%s\n' "${base_image_digest}" - printf 'use-sccache:%s\n' "${use_sccache}" - if [[ "${use_sccache}" == "1" ]]; then - # These values can change the installed binary or final image - # configuration. SCCACHE_ENDPOINT is transport-only and therefore - # deliberately excluded from image identity. - for sccache_arg in \ - SCCACHE_DOWNLOAD_URL \ - SCCACHE_BUCKET_NAME \ - SCCACHE_REGION_NAME \ - SCCACHE_S3_NO_CREDENTIALS; do - sccache_value="${!sccache_arg:-}" - if [[ -z "${sccache_value}" ]]; then - sccache_value=$(extract_arg_default "${sccache_arg}") - fi - printf 'arg:%s=%s\n' \ - "${sccache_arg}" "${sccache_value:-}" - done - fi - } | sha256sum | cut -d' ' -f1 -} - -build_base_image() { - local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}" - local base_hash="" - local base_image_arg="" - local base_image_digest="" - local pinned_base_image="" - local rocm_version="" - local triton_arg="" - local pytorch_arg="" - local pytorch_vision_arg="" - local pytorch_audio_arg="" - local fa_arg="" - local aiter_arg="" - local mori_arg="" - local python_version_arg="" - local pytorch_rocm_arch_arg="" - local dependency_summary="" - local stable_tag="${BASE_REPO}:base" - local trusted_content_tag="" - local scoped_content_tag="" - local preview_content_tag="" - local writable_content_tag="" - local build_ref="" - local immutable_ref="" - local selected_digest="" - local stable_digest="" - local reuse_status=0 - local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}" - local content_files_hash="" - local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}" - local -a sccache_args=() - local -a content_paths=() - - if [[ ! -f "${DOCKERFILE}" ]]; then - echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2 - exit 1 - fi - if [[ ! "${metadata_version}" =~ ^[A-Za-z0-9][A-Za-z0-9_.-]{0,15}$ ]]; then - echo "Invalid ROCm base metadata version: ${metadata_version}" >&2 - return 1 - fi - if [[ "${use_sccache}" != "0" && "${use_sccache}" != "1" ]]; then - echo "ROCm base USE_SCCACHE must be 0 or 1: ${use_sccache}" >&2 - return 1 - fi - - base_image_arg="$(extract_arg_default BASE_IMAGE)" - if ! base_image_digest=$(resolve_image_digest "${base_image_arg}"); then - echo "Failed to resolve ROCm base input image: ${base_image_arg}" >&2 - return 1 - fi - pinned_base_image="${base_image_arg%@*}@${base_image_digest}" - read -r -a content_paths <<< "${content_files}" - content_files_hash="$(compute_content_hash "${content_paths[@]}")" - base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}") - rocm_version="$(extract_arg_default ROCM_SDK_VERSION)" - rocm_version="${rocm_version:-$(rocm_version_from_base_image "${base_image_arg}")}" - triton_arg="$(extract_arg_default TRITON_BRANCH)" - pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)" - pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)" - pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)" - triton_arg="${triton_arg:-$(extract_arg_default TRITON_VERSION)}" - pytorch_arg="${pytorch_arg:-$(extract_arg_default TORCH_VERSION)}" - pytorch_vision_arg="${pytorch_vision_arg:-$(extract_arg_default TORCHVISION_VERSION)}" - pytorch_audio_arg="${pytorch_audio_arg:-$(extract_arg_default TORCHAUDIO_VERSION)}" - fa_arg="$(extract_arg_default FA_BRANCH)" - aiter_arg="$(extract_arg_default AITER_BRANCH)" - mori_arg="$(extract_arg_default MORI_BRANCH)" - python_version_arg="$(extract_arg_default PYTHON_VERSION)" - pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)" - dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}" - trusted_content_tag=$(trusted_base_content_ref "${base_hash}" "${metadata_version}") - scoped_content_tag=$(scoped_base_content_ref "${base_hash}" "${metadata_version}") - preview_content_tag="${scoped_content_tag}" - # Preview writes share an exact-content namespace so stacked PRs can reuse - # identical images. They may import the trusted ref, but never overwrite it. - if is_trusted_main_build; then - scoped_content_tag="${trusted_content_tag}" - writable_content_tag="${trusted_content_tag}" - else - writable_content_tag="${scoped_content_tag}" - fi - - configure_rocm_base_layer_cache - - if [[ "${use_sccache}" == "1" ]]; then - for env_name in \ - SCCACHE_DOWNLOAD_URL \ - SCCACHE_ENDPOINT \ - SCCACHE_BUCKET_NAME \ - SCCACHE_REGION_NAME \ - SCCACHE_S3_NO_CREDENTIALS; do - if [[ -n "${!env_name:-}" ]]; then - sccache_args+=(--build-arg "${env_name}=${!env_name}") - fi - done - fi - - echo "--- :docker: Preparing ROCm base image" - echo "Dockerfile: ${DOCKERFILE}" - echo "Trusted content tag: ${trusted_content_tag}" - echo "Writable content tag: ${writable_content_tag}" - echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))" - echo "Content hash: ${base_hash}" - echo "Dependency summary: ${dependency_summary}" - echo "USE_SCCACHE: ${use_sccache}" - echo "BuildKit layer cache: ${ROCM_BASE_LAYER_CACHE_REF}" - - if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" != "1" \ - && "${ROCM_BASE_NO_CACHE:-0}" != "1" ]]; then - reuse_status=0 - if [[ "${trusted_content_tag}" == "${scoped_content_tag}" ]]; then - immutable_ref=$(find_matching_base_content_ref \ - "${base_hash}" "${metadata_version}" \ - "${trusted_content_tag}" "${preview_content_tag}" \ - "${stable_tag}") || reuse_status=$? - else - immutable_ref=$(find_matching_base_content_ref \ - "${base_hash}" "${metadata_version}" \ - "${trusted_content_tag}" "${scoped_content_tag}" \ - "${stable_tag}") \ - || reuse_status=$? - fi - else - reuse_status=1 - echo "Forced/no-cache ROCm base build; bypassing content image reuse" - fi - - if ((reuse_status > 1)); then - return "${reuse_status}" - fi - - if [[ ${reuse_status} -eq 0 && -n "${immutable_ref}" ]]; then - echo "Reusing ROCm base image with matching content: ${immutable_ref}" - if [[ "${immutable_ref%@*}" != "${writable_content_tag}" \ - && ( "${immutable_ref%@*}" == "${stable_tag}" \ - || "${immutable_ref%@*}" == "${preview_content_tag}" ) ]]; then - echo "Backfilling ROCm base content tag: ${writable_content_tag}" - docker buildx imagetools create --prefer-index=false \ - -t "${writable_content_tag}" "${immutable_ref}" - if ! immutable_ref=$(find_matching_base_content_ref \ - "${base_hash}" "${metadata_version}" \ - "${writable_content_tag}"); then - echo "Backfilled ROCm base content tag failed identity validation" >&2 - return 1 - fi - fi - else - build_ref="${writable_content_tag}" - echo "No reusable ROCm base content image found; building ${build_ref}" - setup_builder - docker buildx build \ - "${ROCM_BASE_CACHE_ARGS[@]}" \ - --pull \ - --provenance=false \ - --progress "${BUILDKIT_PROGRESS:-plain}" \ - --file "${DOCKERFILE}" \ - --build-arg "BASE_IMAGE=${pinned_base_image}" \ - --build-arg "USE_SCCACHE=${use_sccache}" \ - "${sccache_args[@]}" \ - --label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \ - --label "org.opencontainers.image.vendor=vLLM" \ - --label "org.opencontainers.image.title=vLLM ROCm base" \ - --label "vllm.rocm_base.metadata_version=${metadata_version}" \ - --label "vllm.rocm_base.content_hash=${base_hash}" \ - --label "vllm.rocm_base.content_files_hash=${content_files_hash}" \ - --label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \ - --label "vllm.rocm_base.dependency_summary=${dependency_summary}" \ - --label "vllm.rocm_base.base_image=${pinned_base_image}" \ - --label "vllm.rocm_base.base_image_digest=${base_image_digest}" \ - --label "vllm.rocm_base.dependency.rocm=${rocm_version}" \ - --label "vllm.rocm_base.dependency.python=${python_version_arg}" \ - --label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \ - --label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \ - --label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \ - --label "vllm.rocm_base.dependency.triton=${triton_arg}" \ - --label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \ - --label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \ - --label "vllm.rocm_base.dependency.mori=${mori_arg}" \ - --label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \ - -t "${build_ref}" \ - --push \ - . - if ! immutable_ref=$(find_matching_base_content_ref \ - "${base_hash}" "${metadata_version}" "${build_ref}"); then - echo "Published ROCm base image failed identity validation: ${build_ref}" >&2 - return 1 - fi - fi - - selected_digest="${immutable_ref##*@}" - metadata_set "rocm-base-image" "${immutable_ref}" - metadata_set "rocm-base-push-stable-tag" "0" - stable_digest=$(resolve_image_digest "${stable_tag}" 2>/dev/null || true) - if [[ -n "${stable_digest}" && "${stable_digest}" == "${selected_digest}" ]]; then - echo "--- :white_check_mark: Existing ROCm base is current" - echo "Stable ${stable_tag} already resolves to ${selected_digest}" - return 0 - fi - - promote_stable_base_tag "${immutable_ref}" "${stable_tag}" - - metadata_set "rocm-base-push-stable-tag" "${ROCM_BASE_STABLE_TAG_UPDATED}" - metadata_set "rocm-base-refresh" "1" - - echo "--- :white_check_mark: ROCm base image published" - echo "Use BASE_IMAGE=${immutable_ref} for downstream ROCm CI builds" -} - -main() { - if using_custom_rocm_dockerfiles; then - if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then - echo "Custom ROCm Dockerfiles require ROCM_BASE_REFRESH_SKIP=0" >&2 - return 2 - fi - validate_rocm_dockerfile "${DOCKERFILE}" || return $? - fi - metadata_set "rocm-base-refresh" "0" - - if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then - local stable_ref="${BASE_REPO}:base" - local stable_digest="" - - if ! stable_digest=$(resolve_image_digest "${stable_ref}"); then - echo "Could not pin explicitly skipped ROCm base: ${stable_ref}" >&2 - return 1 - fi - metadata_set "rocm-base-image" "${stable_ref}@${stable_digest}" - metadata_set "rocm-base-push-stable-tag" "0" - echo "ROCM_BASE_REFRESH_SKIP=1; pinned current stable base at ${stable_digest}" - return 0 - fi - build_base_image -} - -if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then - main "$@" -fi diff --git a/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py index af3de54f6020..3559afe404d4 100755 --- a/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py +++ b/.buildkite/scripts/rocm/select_aiter_nightly_wheel.py @@ -3,9 +3,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Select the AITER nightly wheel for a vLLM ROCm image. -aiter-nightly-overlay.sh runs this inside the stock ci_base, so it reads that -image's ROCm, torch and Python. It prints the wheel's URL, then the Buildkite -annotation naming it. +docker/Dockerfile.rocm runs this in the test image's AITER nightly step, so it +reads that image's ROCm, torch and Python. It prints the wheel's URL, then the +Buildkite annotation naming it. """ import argparse @@ -21,17 +21,21 @@ import regex as re -INDEX_URL = "https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/amd-aiter/" +INDEX_URL = ( + "https://rocm.frameworks-nightlies.amd.com/whl-multi-arch-staging/amd-aiter/" +) AITER_COMMIT_URL = "https://github.com/ROCm/aiter/commit/" STALE_AFTER = datetime.timedelta(days=1) # Example: -# amd_aiter-0.1.25+rocm7.2.3.torch2.12.ae984f5.d20261008-cp312-cp312-linux_x86_64.whl -# is AITER main at commit ae984f5, built 2026-10-08, for ROCm 7.2.3, torch 2.12 -# and Python 3.12. 0.1.25 is a guess at the next release, not a real one. +# amd_aiter-0.1.28+rocm10.0.0.torch2.13.f684e21.d20261029 +# -cp312-cp312-manylinux_2_28_x86_64.whl +# is AITER main at commit f684e21, built 2026-10-29, for ROCm 10.0.0, torch 2.13 +# and Python 3.12, and runs on glibc 2.28 or newer. 0.1.28 is a guess at the +# next release, not a real one. WHEEL_NAME = re.compile( r"^amd_aiter-(?P[^+-]+)\+(?P[^-]+)" - r"-(?Pcp\d+)-(?P=python)-linux_x86_64\.whl$" + r"-(?Pcp\d+)-(?P=python)-(?:manylinux_\d+_\d+|linux)_x86_64\.whl$" ) ROCM_FIELD = re.compile(r"(?:^|\.)rocm(\d+(?:\.\d+)*)(?:a\d+)?(?=\.|$)") TORCH_FIELD = re.compile(r"(?:^|\.)torch(\d+\.\d+)(?=\.|$)") @@ -235,7 +239,9 @@ def _today() -> datetime.date: return datetime.datetime.now(datetime.UTC).date() -# Not 1 or 2, which Buildkite retries; see aiter-nightly-overlay.sh. +# Its own status for "no wheel for this image". Dockerfile.rocm fails the image +# build on any failure here; bake reports that as exit 1, which the image-build +# step retries once. EXIT_NO_WHEEL = 10 diff --git a/.buildkite/scripts/rocm/smoke-test-image.sh b/.buildkite/scripts/rocm/smoke-test-image.sh deleted file mode 100755 index b155a0897d32..000000000000 --- a/.buildkite/scripts/rocm/smoke-test-image.sh +++ /dev/null @@ -1,84 +0,0 @@ -#!/usr/bin/env bash -# Fast structural smoke test for the full ROCm CI image. - -set -euo pipefail - -run_smoke_checks() { - local required_dir="" - - for required_dir in \ - /vllm-workspace \ - /vllm-workspace/tests \ - /vllm-workspace/src/vllm; do - if [[ ! -d "${required_dir}" ]]; then - echo "Missing directory: ${required_dir}" >&2 - return 1 - fi - done - if [[ ! -x /vllm-workspace/src/vllm/vllm-rs ]]; then - echo "Missing executable: /vllm-workspace/src/vllm/vllm-rs" >&2 - return 1 - fi - - command -v python3 - command -v uv - command -v pytest - - if ! command -v amd-smi >/dev/null 2>&1 \ - && ! command -v rocminfo >/dev/null 2>&1; then - echo "No ROCm CLI found in image" >&2 - return 1 - fi - - PYTHONDONTWRITEBYTECODE=1 python3 - <<'PY' -import torch -import vllm - -print(torch.__version__) -print(vllm.__version__) -PY - - echo "AMD image smoke OK" -} - -if [[ "${1:-}" == "--inside" ]]; then - run_smoke_checks - exit -fi -if (($#)); then - echo "Usage: $0 [--inside]" >&2 - exit 2 -fi - -if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then - base_refreshed="" - if command -v buildkite-agent >/dev/null 2>&1; then - base_refreshed="$(buildkite-agent meta-data get rocm-base-refresh 2>/dev/null || true)" - fi - if [[ "${base_refreshed}" != "1" ]]; then - echo "ROCM_CI_ARTIFACT_ONLY=1; no full image was built, skipping smoke test" - exit 0 - fi -fi - -smoke_marker="./build/rocm-smoke-export/vllm-smoke-ok" -expected_smoke_id="${BUILDKITE_BUILD_ID:-local}" -if [[ -f "${smoke_marker}" \ - && ( -z "${VLLM_CI_SMOKE_IMAGE:-}" \ - || "${VLLM_CI_SMOKE_IMAGE}" == "${IMAGE_TAG:-}" ) ]]; then - actual_smoke_id="$(< "${smoke_marker}")" - if [[ "${actual_smoke_id}" != "${expected_smoke_id}" ]]; then - echo "ROCm smoke marker belongs to ${actual_smoke_id}, not ${expected_smoke_id}" \ - >&2 - exit 1 - fi - echo "AMD image smoke OK (verified inside BuildKit)" - rm -f -- "${smoke_marker}" - rmdir -- "$(dirname "${smoke_marker}")" 2>/dev/null || true - exit -fi - -image_ref="${VLLM_CI_SMOKE_IMAGE:-${IMAGE_TAG:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}}" - -docker run --rm -i --network=none --entrypoint /bin/bash "${image_ref}" \ - -s -- --inside < "${BASH_SOURCE[0]}" diff --git a/.buildkite/scripts/rocm/stack.sh b/.buildkite/scripts/rocm/stack.sh deleted file mode 100644 index 14bfeb48a017..000000000000 --- a/.buildkite/scripts/rocm/stack.sh +++ /dev/null @@ -1,78 +0,0 @@ -#!/usr/bin/env bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# -# Derive everything the release pipeline needs about a ROCm stack from its base -# Dockerfile, so the stack's image is the single input to every ROCm script. -# -# Usage: source .buildkite/scripts/rocm/stack.sh -# e.g. docker/Dockerfile.rocm_base (ROCm 10.0) or docker/Dockerfile.rocm_72_base -# -# Exports: -# ROCM_STACK_BASE_DOCKERFILE the base Dockerfile -# ROCM_STACK_DOCKERFILE the matching vLLM Dockerfile (base name minus _base) -# ROCM_STACK_THEROCK 1 if the ROCm SDK comes from TheRock wheels -# ROCM_STACK_VERSION ROCm version, e.g. 10.0.0 or 7.2.3 -# ROCM_STACK_VARIANT wheel variant and Docker Hub tag flavor from major.minor, -# same rule as setup.py (rocm100, rocm72) -# ROCM_STACK_IMAGE_KEY hash of the base Dockerfile (base image cache key) -# ROCM_STACK_WHEEL_KEY hash of the base Dockerfile plus files it COPYs from -# the repo (base wheel cache key) -# ROCM_STACK_ECR_BASE ECR tag of the cached base image -# ROCM_STACK_ECR_IMAGE ECR tag of this commit's release image -# ROCM_STACK_BASE_WHEELS_DIR local dir for base wheels (also the artifact path) -# ROCM_STACK_VLLM_WHEEL_DIR local dir for the vLLM wheel (also the artifact path) - -rocm_stack_init() { - local base="$1" - if [[ ! -f "$base" ]]; then - echo "ROCm base Dockerfile not found: $base" >&2 - return 2 - fi - ROCM_STACK_BASE_DOCKERFILE="$base" - ROCM_STACK_DOCKERFILE="${base%_base}" - - local sdk - sdk=$(sed -nE 's/^ARG ROCM_SDK_VERSION=([0-9.]+).*/\1/p' "$base" | head -1) - if [[ -n "$sdk" ]]; then - ROCM_STACK_THEROCK=1 - ROCM_STACK_VERSION="$sdk" - else - ROCM_STACK_THEROCK=0 - # BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.2.3-complete -> 7.2.3 - ROCM_STACK_VERSION=$(sed -nE 's/^ARG BASE_IMAGE=.*:([0-9]+(\.[0-9]+)+).*/\1/p' "$base" | head -1) - fi - if [[ -z "$ROCM_STACK_VERSION" ]]; then - echo "Could not determine the ROCm version from $base" >&2 - return 2 - fi - - local major minor - IFS=. read -r major minor _ <<< "$ROCM_STACK_VERSION" - ROCM_STACK_VARIANT="rocm${major}${minor}" - - ROCM_STACK_IMAGE_KEY=$(sha256sum "$base" | cut -c1-16) - local inputs - inputs=$(grep -E '^COPY ' "$base" | grep -v -- '--from' | awk '{print $2}' || true) - if [[ -z "$inputs" ]]; then - ROCM_STACK_WHEEL_KEY="$ROCM_STACK_IMAGE_KEY" - else - # shellcheck disable=SC2086 - ROCM_STACK_WHEEL_KEY=$(sha256sum "$base" $inputs | sha256sum | cut -c1-16) - fi - - local ecr="public.ecr.aws/q9t5s3a7/vllm-release-repo" - ROCM_STACK_ECR_BASE="${ecr}:${ROCM_STACK_IMAGE_KEY}-${ROCM_STACK_VARIANT}-base" - ROCM_STACK_ECR_IMAGE="${ecr}:${BUILDKITE_COMMIT:-unknown}-${ROCM_STACK_VARIANT}" - ROCM_STACK_BASE_WHEELS_DIR="artifacts/${ROCM_STACK_VARIANT}-base-wheels" - ROCM_STACK_VLLM_WHEEL_DIR="artifacts/${ROCM_STACK_VARIANT}-vllm-wheel" - - export ROCM_STACK_BASE_DOCKERFILE ROCM_STACK_DOCKERFILE ROCM_STACK_THEROCK \ - ROCM_STACK_VERSION ROCM_STACK_VARIANT ROCM_STACK_IMAGE_KEY \ - ROCM_STACK_WHEEL_KEY ROCM_STACK_ECR_BASE ROCM_STACK_ECR_IMAGE \ - ROCM_STACK_BASE_WHEELS_DIR ROCM_STACK_VLLM_WHEEL_DIR -} - -if [[ $# -gt 0 ]]; then - rocm_stack_init "$1" -fi diff --git a/.buildkite/scripts/upload-rocm-wheels.sh b/.buildkite/scripts/upload-rocm-wheels.sh index 5cf32bcfb358..d303f16eae0e 100755 --- a/.buildkite/scripts/upload-rocm-wheels.sh +++ b/.buildkite/scripts/upload-rocm-wheels.sh @@ -4,8 +4,8 @@ # # Upload ROCm wheels to S3 with proper index generation # -# Usage: upload-rocm-wheels.sh -# base-dockerfile the ROCm stack, e.g. docker/Dockerfile.rocm_base (see rocm/stack.sh) +# Usage: upload-rocm-wheels.sh +# stack rocm (TheRock); dir the stack's wheels plus the vLLM wheel # # Required environment variables: # AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY (or IAM role) @@ -23,11 +23,11 @@ set -ex # ======== Configuration ======== BUCKET="${S3_BUCKET:-vllm-wheels}" -# shellcheck source=.buildkite/scripts/rocm/stack.sh -source "$(dirname "${BASH_SOURCE[0]}")/rocm/stack.sh" "${1:?Usage: $0 }" +STACK="${1:?Usage: $0 }" +WHEEL_DIR="${2:?Usage: $0 }" +VARIANT=$(basename "$WHEEL_DIR"/vllm-*.whl | grep -oE 'rocm[0-9]+' | tail -1) ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}" -WHEEL_DIRS="$ROCM_STACK_BASE_WHEELS_DIR $ROCM_STACK_VLLM_WHEEL_DIR" -WHEEL_SUBPATH="$ROCM_SUBPATH/${ROCM_STACK_VARIANT}-wheels" +WHEEL_SUBPATH="$ROCM_SUBPATH/${VARIANT:?no vLLM wheel in $WHEEL_DIR}-wheels" S3_WHEEL_PREFIX="s3://$BUCKET/$WHEEL_SUBPATH/" INDICES_OUTPUT_DIR="rocm-indices" @@ -61,9 +61,7 @@ source .buildkite/scripts/lib/manylinux.sh # Collect all wheels rm -rf all-rocm-wheels "$INDICES_OUTPUT_DIR" mkdir -p all-rocm-wheels -for dir in $WHEEL_DIRS; do - cp "$dir"/*.whl all-rocm-wheels/ 2>/dev/null || true -done +cp "$WHEEL_DIR"/*.whl all-rocm-wheels/ 2>/dev/null || true WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l) echo "Total wheels to upload: $WHEEL_COUNT" @@ -119,10 +117,10 @@ sed -i 's/import pybase64 as base64/import base64/g' tools/vllm-rocm/therock_whe INDEX_ARGS=(--version "$ROCM_SUBPATH" --wheel-dir "${WHEEL_SUBPATH#rocm/}") # TheRock's SDK, device kernels and torchvision/torchaudio are linked, not hosted -if [[ "$ROCM_STACK_THEROCK" == "1" ]]; then +if [[ "$STACK" == "rocm" ]]; then $PYTHON tools/vllm-rocm/therock_wheels.py external-links \ - --dockerfile "$ROCM_STACK_BASE_DOCKERFILE" \ - --torch-wheel "$(ls "$ROCM_STACK_BASE_WHEELS_DIR"/torch-*.whl)" > therock-external-links.txt + --dockerfile docker/Dockerfile.rocm_base \ + --torch-wheel "$(ls "$WHEEL_DIR"/torch-*.whl)" > therock-external-links.txt cat therock-external-links.txt INDEX_ARGS+=(--external-links therock-external-links.txt) fi @@ -191,11 +189,11 @@ echo "Wheels available at:" echo " s3://$BUCKET/$WHEEL_SUBPATH/" echo "" echo "Install command (by commit):" -echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/$ROCM_STACK_VARIANT/" +echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/$VARIANT/" echo "" if [[ "${NIGHTLY:-0}" == "1" ]]; then echo "Install command (nightly):" - echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/$ROCM_STACK_VARIANT/" + echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/$VARIANT/" fi echo "" echo "Wheel count: $WHEEL_COUNT" diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 6392e04f0540..a0253a49e2e6 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -14,12 +14,9 @@ # command(str): the single command to run for tests. incompatible with commands. # commands(list): the list of commands to run for the test. incompatible with command. # mirror_hardwares(list): selector tags for AMD pipelines that should include the test. -# dind(bool): when false, run the job directly in the AMD Kubernetes pod. -# When true or omitted, use the legacy Docker-in-Docker path. +# dind(bool): must be false: the job runs directly in the AMD Kubernetes pod (no Docker-in-Docker). # num_gpus(int): override the number of GPUs for the test. defaults to 1 GPU. currently supports 2,4. -# num_nodes(int): whether to simulate multi-node setup by launching multiple containers on one host, -# in this case, commands must be specified. the first command runs on the first host, the second -# command runs on the second host. +# num_nodes(int): not supported on AMD (num_nodes >= 2 fails). # timeout_in_minutes(int): sets a timeout for the step in minutes. if not specified, uses the default timeout. # parallelism(int): number of parallel jobs to run for this step. enables test sharding using $$BUILDKITE_PARALLEL_JOB # and $$BUILDKITE_PARALLEL_JOB_COUNT environment variables. @@ -121,36 +118,45 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - .buildkite/scripts/docker-build-metadata-args.sh - - .buildkite/scripts/ci-bake-rocm.sh - .buildkite/release-pipeline.yaml - docker/Dockerfile - docker/Dockerfile.cpu - docker/Dockerfile.rocm - docker/Dockerfile.rocm_base - - docker/ci-rocm.hcl - docker/docker-bake.hcl - docker/docker-bake-rocm.hcl + - .buildkite/scripts/rocm/build.sh + - .buildkite/scripts/rocm/publish.sh + - tools/vllm-rocm/therock_wheels.py - tests/tools/test_docker_build_metadata_args.py - commands: - - pytest -v -s tools/test_docker_build_metadata_args.py - -- label: ":amd: (MI300) Torch Stable ABI Audit" # TBD - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - no_gpu: true - optional: true - working_dir: "/vllm-workspace/" - source_file_dependencies: - - .buildkite/check-torch-abi.py - - csrc/ - - cmake/ - - setup.py - - vllm/platforms/rocm.py - commands: - - uv pip install --system torch-abi-audit==0.0.1 - - python3 /vllm-workspace/.buildkite/check-torch-abi.py + - tests/tools/test_rocm_build.py + - tests/tools/test_amd_ci_fast_exit.py + - .buildkite/scripts/hardware_ci/amd_ci_fast_exit.py + - tests/tools/test_kernel_symbol_map.py + - tools/ci/kernel_symbol_map.py + - tools/ci/rocm_kernel_symbol_map.py + commands: + - pytest -v -s tools/test_docker_build_metadata_args.py tools/test_rocm_build.py tools/test_amd_ci_fast_exit.py + - pytest -v -s tools/test_kernel_symbol_map.py + +# Disabled until torch-abi-audit transitive dependencies are locked. +# - label: ":amd: (MI300) Torch Stable ABI Audit" # TBD +# timeout_in_minutes: 35 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# no_gpu: true +# optional: true +# working_dir: "/vllm-workspace/" +# source_file_dependencies: +# - .buildkite/check-torch-abi.py +# - csrc/ +# - cmake/ +# - setup.py +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system torch-abi-audit==0.0.1 +# - python3 /vllm-workspace/.buildkite/check-torch-abi.py ######################################################################################################################################### # # @@ -207,22 +213,23 @@ steps: #---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------# -- label: ":amd: (MI250) Helion Kernels Shard %N" - timeout_in_minutes: 45 - parallelism: 2 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - dind: false - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/utils/import_utils.py - - tools/benchmark_helion_kernels.py - - tests/kernels/helion/ - - vllm/platforms/rocm.py - commands: - - pip install helion==1.4.0 - - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# Disabled until helion transitive dependencies are locked. +# - label: ":amd: (MI250) Helion Kernels Shard %N" +# timeout_in_minutes: 45 +# parallelism: 2 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] +# dind: false +# agent_pool: mi250_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/utils/import_utils.py +# - tools/benchmark_helion_kernels.py +# - tests/kernels/helion/ +# - vllm/platforms/rocm.py +# commands: +# - pip install helion==1.4.0 +# - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT #----------------------------------------------------- mi250 · models / language -----------------------------------------------------# @@ -291,7 +298,6 @@ steps: - export VLLM_GPU_MEMORY_UTILIZATION=0.8 - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA @@ -328,7 +334,13 @@ steps: - tests/v1/e2e/general/ - vllm/platforms/rocm.py commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_async_scheduling_accuracy.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2" + - >- + pytest -v -s + v1/e2e/general + --ignore v1/e2e/general/test_async_scheduling.py + --ignore v1/e2e/general/test_async_scheduling_accuracy.py + --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py + -k "not test_mamba_prefix_cache_mrv2" - label: ":amd: (MI250) E2E Scheduling" # TBD timeout_in_minutes: 55 @@ -377,7 +389,10 @@ steps: - vllm/v1/engine/ 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/engine + --ignore v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/test_tensor_ipc_queue.py - pytest -v -s config/test_mp_reducer.py @@ -453,7 +468,24 @@ steps: - tests/utils.py - tests/models/test_vision.py commands: - - case "$$BUILDKITE_PARALLEL_JOB" in 0) TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' && HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' ;; 1) pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && pytest models/language -v -s -m 'distributed(num_gpus=2)' && VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' ;; 2) pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' && pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' ;; *) exit 2 ;; esac + - >- + case "$$BUILDKITE_PARALLEL_JOB" in + 0) + TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' && + HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' ;; + 1) + pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && + pytest models/language -v -s -m 'distributed(num_gpus=2)' && + VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' && + pytest -v -s models/test_mimo_v2_attention_backend.py && + pytest -v -s models/multimodal/test_mimo_v2_omni.py ;; + 2) + pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' + --ignore models/multimodal/generation/test_whisper.py && + VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' && + pytest -v -s models/test_vision.py -m 'distributed(num_gpus=2)' ;; + *) exit 2 ;; + esac #---------------------------------------------------------- mi300 · compile ----------------------------------------------------------# @@ -584,7 +616,10 @@ steps: - vllm/v1/ optional: true commands: - - pytest -s -v compile/passes --ignore compile/passes/distributed + - >- + pytest -s -v + compile/passes + --ignore compile/passes/distributed - label: ":amd: (MI300) PyTorch Fullgraph" # TBD timeout_in_minutes: 90 @@ -873,7 +908,6 @@ steps: - tests/distributed/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - pytest -v -s distributed/test_elastic_ep.py - label: ":amd: (MI300) RayExecutorV2" # TBD @@ -1142,7 +1176,15 @@ steps: - tests/kernels/test_fused_qk_norm_rope_gate.py optional: true commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_fused_qk_norm_rope_gate.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - >- + pytest -v -s + kernels/core + --ignore=kernels/core/test_minimax_reduce_rms.py + kernels/test_fused_qk_norm_rope_gate.py + kernels/test_concat_mla_q.py + kernels/test_top_k_per_row.py + kernels/test_rocm_fp32_router_gemm.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) Quantization Kernels Shard %N" timeout_in_minutes: 120 @@ -1194,7 +1236,12 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --ignore=kernels/moe/test_modular_kernel_combinations.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - >- + pytest -v -s + kernels/moe + --ignore=kernels/moe/test_modular_oai_triton_moe.py + --ignore=kernels/moe/test_modular_kernel_combinations.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s kernels/moe/test_modular_kernel_combinations.py -k "not multigpu" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -1260,7 +1307,15 @@ steps: - vllm/platforms/rocm.py - vllm/utils/flashinfer.py commands: - - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - >- + pytest -v -s + kernels/attention + --ignore=kernels/attention/test_triton_unified_attention_diffkv.py + --ignore=kernels/attention/test_rocm_aiter_mla_decode.py + --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py + --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py + --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI300) MiniMax Reduce RMS Kernels" # TBD timeout_in_minutes: 45 @@ -1313,7 +1368,12 @@ steps: - vllm/platforms/rocm.py commands: # Run the ROCm AttnRes/KDA decode counterparts plus existing KDA coverage. - - pytest -v -s models/kimi_k3/test_amd_attn_res.py models/kimi_k3/test_amd_kda_decode.py models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py + - >- + pytest -v -s + models/kimi_k3/test_amd_attn_res.py + models/kimi_k3/test_amd_kda_decode.py + models/kimi_k3/test_kda.py + models/kimi_k3/test_kda_metadata.py - label: ":amd: (MI355) Kimi K3" # TBD timeout_in_minutes: 60 @@ -1334,7 +1394,7 @@ steps: # gated on gfx950, so this has to be MI355. The two kernels/ files the NVIDIA # step runs are CUDA-only (they import cuda.bindings.driver), so they are # deliberately left out here. - - pytest -v -s models/kimi_k3 + - pytest -v -s models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py - label: ":amd: (MI355) Kimi K3 Multi-GPU" # TBD timeout_in_minutes: 30 @@ -1355,7 +1415,7 @@ steps: # multi-GPU test anywhere in the directory without editing this list. # Four GPUs run the two tp4 cases in test_amd_latent_moe_runner.py; the tp8 # case and the 8/16-GPU cases in test_latent_moe_tail.py still skip. - - pytest -v -s -m "distributed" models/kimi_k3 + - pytest -v -s -m "distributed" models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py - label: ":amd: (MI355) MiniMax M3 Specialized Decode" timeout_in_minutes: 45 @@ -1511,8 +1571,6 @@ steps: - vllm/ - tests/models/language/generation commands: - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE 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)' #---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------# @@ -1720,20 +1778,21 @@ steps: - pytest -v -s plugins_tests/test_oot_registration_offline.py - pytest -v -s plugins_tests/lora_resolvers -- label: ":amd: (MI300) GGUF Plugin" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/plugins_tests/gguf - - vllm/platforms/rocm.py - commands: - - pip install "vllm-gguf-plugin >= 0.0.2" - - pytest -v -s plugins_tests/gguf +# Disabled until the GGUF plugin installation is locked. +# - label: ":amd: (MI300) GGUF Plugin" # TBD +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/model_executor/layers/quantization +# - tests/plugins_tests/gguf +# - vllm/platforms/rocm.py +# commands: +# - pip install "vllm-gguf-plugin >= 0.0.2" +# - pytest -v -s plugins_tests/gguf #------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------# @@ -1837,40 +1896,45 @@ steps: commands: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" + - >- + pytest -v -s + tool_use + --ignore=tool_use/mistral + --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice" #------------------------------------------------------- mi300 · quantization --------------------------------------------------------# -- label: ":amd: (MI300) Quantization Shard %N" - timeout_in_minutes: 115 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/layers/fused_moe/oracle/unquantized.py - - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - - tests/quantization - - tests/rocm/test_moe_weight_replay.py - - tests/rocm/test_moe_padding.py - optional: true - commands: - # temporary install here since we need nightly, will move to requirements/test.in - # after torchao 0.12 release, and pin a working version of torchao nightly here - # since torchao nightly is only compatible with torch nightly currently - # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now - # we can only upgrade after this is resolved - # TODO(jerryzh168): resolve the above comment - - uv pip install --system torchao==0.18.0 - - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi +# Disabled until torchao transitive dependencies are locked. +# - label: ":amd: (MI300) Quantization Shard %N" +# timeout_in_minutes: 115 +# parallelism: 4 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - csrc/ +# - vllm/model_executor/layers/quantization +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# - vllm/model_executor/layers/fused_moe/oracle/unquantized.py +# - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +# - tests/quantization +# - tests/rocm/test_moe_weight_replay.py +# - tests/rocm/test_moe_padding.py +# optional: true +# commands: +# # temporary install here since we need nightly, will move to requirements/test.in +# # after torchao 0.12 release, and pin a working version of torchao nightly here +# # since torchao nightly is only compatible with torch nightly currently +# # https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now +# # we can only upgrade after this is resolved +# # TODO(jerryzh168): resolve the above comment +# - uv pip install --system torchao==0.18.0 +# - uv pip install --system conch-triton-kernels +# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi +# - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi - label: ":amd: (MI300) Quantized Fusions" # TBD timeout_in_minutes: 45 @@ -1912,30 +1976,31 @@ steps: #------------------------------------------------------------ mi300 · misc ------------------------------------------------------------# -- label: ":amd: (MI300) Regression" # TBD - timeout_in_minutes: 25 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/test_regression - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - optional: true - commands: - - pip install 'modelscope<1.38' - - pytest -v -s test_regression.py +# Disabled until the modelscope installation is locked. +# - label: ":amd: (MI300) Regression" # TBD +# timeout_in_minutes: 25 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] +# dind: false +# agent_pool: mi300_1 +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/ +# - tests/test_regression +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/inputs/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/platforms/ +# - vllm/sampling_params.py +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# optional: true +# commands: +# - pip install 'modelscope<1.38' +# - pytest -v -s test_regression.py #--------------------------------------------------------- mi300 · ray_compat ---------------------------------------------------------# @@ -1974,7 +2039,6 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py @@ -2005,6 +2069,7 @@ steps: - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + - VLLM_TEST_MODEL=ibm-granite/granite-3.0-1b-a400m-base VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py - label: ":amd: (MI300) V1 Spec Decode Shard %N" # TBD timeout_in_minutes: 50 @@ -2073,7 +2138,10 @@ steps: - vllm/platforms/rocm.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/engine + --ignore v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/test_tensor_ipc_queue.py - label: ":amd: (MI300) Extract Hidden States Integration" # TBD @@ -2091,7 +2159,9 @@ steps: - vllm/model_executor/model_loader/ - vllm/model_executor/models/extract_hidden_states.py - vllm/model_executor/models/llama.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py @@ -2128,7 +2198,11 @@ steps: - vllm/platforms/rocm.py - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - >- + pytest -v -s + v1/attention + --ignore=v1/attention/test_sparse_mla_backends.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT @@ -2157,11 +2231,6 @@ steps: - vllm/utils/ - vllm/v1/ commands: - - "pip install \ - 'opentelemetry-sdk>=1.26.0' \ - 'opentelemetry-api>=1.26.0' \ - 'opentelemetry-exporter-otlp>=1.26.0' \ - 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing - pytest -v -s tracing @@ -2266,7 +2335,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" # TBD @@ -2282,7 +2350,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI300) Distributed DP Extended" # TBD @@ -2322,7 +2389,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" # TBD @@ -2339,7 +2405,6 @@ steps: - tests/v1/kv_connector/nixl_push_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI300) DP EP Distributed NixlConnector PD accuracy" # TBD @@ -2355,7 +2420,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" # TBD @@ -2371,7 +2435,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - 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: ":amd: (MI300) Hybrid SSM NixlConnector PD prefix cache" # TBD @@ -2389,7 +2452,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD accuracy" # TBD @@ -2408,7 +2470,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh - label: ":amd: (MI300) MultiConnector (Nixl+Offloading) PD edge cases" # TBD @@ -2427,7 +2488,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh ######################################################################################################################################### @@ -2657,7 +2717,11 @@ steps: - vllm/model_executor/models/llama.py commands: - amd-smi - - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py + - >- + pytest -v -s + tests/compile/passes/test_fusion_attn.py + tests/compile/passes/test_silu_mul_quant_fusion.py + tests/compile/passes/test_silu_mul_quant_manual_fusion.py - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py - label: ":amd: (MI355 DPX) Fusion E2E FP8 Config Sweep" @@ -2856,7 +2920,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh - label: ":amd: (MI355) Distributed NixlConnector PD accuracy" # TBD @@ -2872,7 +2935,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" @@ -2891,9 +2953,7 @@ steps: - vllm/v1/attention/backends/rocm_aiter_unified_attn.py - vllm/platforms/rocm.py - requirements/common.txt - - requirements/kv_connectors_rocm.txt commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py - label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" @@ -2913,9 +2973,7 @@ steps: - vllm/v1/attention/backends/rocm_aiter_unified_attn.py - vllm/platforms/rocm.py - requirements/common.txt - - requirements/kv_connectors_rocm.txt commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py - label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD @@ -2931,7 +2989,6 @@ steps: - tests/v1/kv_connector/nixl_integration/ - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" @@ -2947,7 +3004,6 @@ steps: - vllm/models/kimi_k3/ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ - - requirements/kv_connectors_rocm.txt commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - export ENABLE_HMA_FLAG='1' @@ -2960,7 +3016,6 @@ steps: - export PREFILL_BLOCK_SIZE='2048' - export DECODE_BLOCK_SIZE='2048' - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - label: ":amd: (MI355) Distributed Features" # TBD @@ -3084,7 +3139,6 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/kv_connectors_rocm.txt - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - vllm/model_executor/layers/fused_moe/ - vllm/model_executor/layers/quantization/ @@ -3094,7 +3148,6 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt # A mirror does not inherit the parent env, so it is repeated here in # the parent's order. The layouts must stay empty: unset, the script # defaults to HND, which the V4 indexer rejects. @@ -3219,7 +3272,15 @@ steps: - vllm/triton_utils/ - vllm/utils/ commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py + - >- + pytest -v -s + engine + test_sequence.py + test_config.py + test_logger.py + test_vllm_port.py + jit_monitor/test_hooks.py + jit_monitor/test_hooks_gpu.py #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------# @@ -3236,7 +3297,12 @@ steps: - tests/entrypoints/llm commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode + - >- + pytest -v -s + entrypoints/llm + --ignore=entrypoints/llm/test_generate.py + --ignore=entrypoints/llm/test_collective_rpc.py + --ignore=entrypoints/llm/offline_mode - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests @@ -3256,10 +3322,18 @@ steps: - tests/entrypoints/rl commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - >- + pytest -v -s + entrypoints/serve + --ignore=entrypoints/serve/dev/rpc + --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - >- + pytest -v -s + entrypoints/rl + --ignore=entrypoints/rl/weight_transfer + --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API completion)" # TBD timeout_in_minutes: 65 @@ -3274,8 +3348,17 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness + - >- + pytest -v -s + entrypoints/openai/completion + --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py + - >- + pytest -v -s + entrypoints/openai + --ignore=entrypoints/openai/completion + --ignore=entrypoints/openai/chat_completion + --ignore=entrypoints/openai/responses + --ignore=entrypoints/openai/correctness - label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API chat_completion)" # TBD timeout_in_minutes: 70 @@ -3498,6 +3581,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py @@ -3558,25 +3642,26 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep-mi355.txt -- label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" - timeout_in_minutes: 105 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/gdn/ - - tests/evals/qwen4_exp/ - - vllm/model_executor/layers/mamba/ - - vllm/models/qwen4_exp/ - - vllm/v1/attention/backends/short_conv_attn.py - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - uv pip install --system 'evalscope==1.10.0' - - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt +# Disabled until evalscope transitive dependencies are locked. +# - label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" +# timeout_in_minutes: 105 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - csrc/gdn/ +# - tests/evals/qwen4_exp/ +# - vllm/model_executor/layers/mamba/ +# - vllm/models/qwen4_exp/ +# - vllm/v1/attention/backends/short_conv_attn.py +# - vllm/platforms/rocm.py +# commands: +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - uv pip install --system 'evalscope==1.10.0' +# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt - label: ":amd: (MI355 DPX) LM Eval TurboQuant k8v4" timeout_in_minutes: 30 @@ -3683,6 +3768,25 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt +- label: ":amd: (MI300) LM Eval Octave KV Cache" + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - tests/evals/gsm8k/ + - tests/utils.py + - tests/conftest.py + - vllm/v1/attention/backends/rocm_octave_attn.py + - vllm/v1/attention/ops/rocm_octave.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt + - label: ":amd: (MI355) LM Eval Spec Decode (Fixed-length FP8)" # TBD timeout_in_minutes: 145 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -4055,6 +4159,8 @@ steps: - vllm/v1/kv_offload/ - vllm/v1/simple_kv_offload/ - tests/evals/gsm8k/test_gsm8k_offloading.py + - vllm/distributed/kv_transfer/kv_connector/v1/umbp/ + - tests/v1/kv_connector/umbp_test_utils.py - vllm/platforms/rocm.py commands: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" @@ -4310,7 +4416,11 @@ steps: - tests/kernels/test_gate_linear_rocm_dispatch.py - tests/kernels/test_bf16x3_router_gemm_rocm.py commands: - - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py + - >- + pytest -v -s + tests/kernels/test_rocm_fp32_router_gemm.py + tests/kernels/test_gate_linear_rocm_dispatch.py + tests/kernels/test_bf16x3_router_gemm_rocm.py - label: ":amd: (MI355 DPX) Attention DiffKV Kernels" timeout_in_minutes: 35 @@ -4440,13 +4550,21 @@ steps: - vllm/utils/flashinfer.py optional: true commands: - - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - -- label: ":amd: (MI355 DPX) MoE Kernels Shard %N" # TBD + - >- + pytest -v -s + kernels/attention + --ignore=kernels/attention/test_triton_unified_attention_diffkv.py + --ignore=kernels/attention/test_rocm_aiter_mla_decode.py + --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py + --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py + --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +- label: ":amd: (MI355) MoE Kernels Shard %N" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] dind: false - agent_pool: mi355_dpx + agent_pool: mi355_1 parallelism: 5 working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -4691,7 +4809,18 @@ steps: - tests/lora - vllm/platforms/rocm.py commands: - - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py + - >- + pytest -v -s + lora + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + --ignore=lora/test_chatglm3_tp.py + --ignore=lora/test_llama_tp.py + --ignore=lora/test_qwen3_with_multi_loras.py + --ignore=lora/test_olmoe_tp.py + --ignore=lora/test_deepseekv2_tp.py + --ignore=lora/test_gptoss_tp.py + --ignore=lora/test_qwen3moe_tp.py + --ignore=lora/test_qwen35_densemodel_lora.py #------------------------------------------------------ mi355 · models / basic -------------------------------------------------------# @@ -4710,9 +4839,15 @@ steps: - tests/models/test_hyv4_rocm.py - tests/models/test_vision.py - tests/models/test_deepseek_v4_vl_rocm.py - - tests/models/test_deepseek_v4_vl_stream_weights.py + - tests/models/test_deepseek_v4_vl.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/test_deepseek_v4_vl_stream_weights.py + - >- + pytest -v -s + models/test_terratorch.py + models/transformers/test_backend.py + models/test_registry.py + models/test_deepseek_v4_vl_rocm.py + models/test_deepseek_v4_vl.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge @@ -4729,10 +4864,41 @@ steps: - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ commands: - - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py + - >- + pytest -v -s + models/qwen4_exp/test_qsa_amd.py + models/qwen4_exp/test_qsa_prepare_amd.py # PLE and QSA import conflicting HC registrations in a shared process. - pytest -v -s models/qwen4_exp/test_ple.py +- label: ":amd: (MI355 DPX) Step5" + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/models/step5/ + - tests/models/step5/ + - vllm/model_executor/models/step3p7.py + - vllm/model_executor/models/step3_vl.py + - vllm/model_executor/models/step_vl.py + - vllm/config/compilation.py + - vllm/model_executor/layers/indexer_topk.py + - vllm/transformers_utils/configs/step5.py + - vllm/model_executor/warmup/step5_dsa_warmup.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/triton_utils/ + - vllm/platforms/rocm.py + - tests/v1/attention/utils.py + - tests/models/test_initialization.py + - tests/models/registry.py + commands: + - pytest -v -s models/step5 + - pytest -v -s models/test_initialization.py -k Step5ForConditionalGeneration + - label: ":amd: (MI355 DPX) GLM5Next" timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -4833,11 +4999,6 @@ steps: - vllm/ - tests/models/language/generation commands: - - export GIT_CONFIG_COUNT=1 - - export GIT_CONFIG_KEY_0=http.version - - export GIT_CONFIG_VALUE_0=HTTP/1.1 - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: ":amd: (MI355 DPX) Granite Language Model Compatibility" @@ -4853,11 +5014,6 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/language/generation commands: - - export GIT_CONFIG_COUNT=1 - - export GIT_CONFIG_KEY_0=http.version - - export GIT_CONFIG_VALUE_0=HTTP/1.1 - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview' - label: ":amd: (MI355 DPX) Language Models (Extended Generation)" # TBD @@ -4871,8 +5027,6 @@ steps: - tests/models/language/generation optional: true commands: - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - label: ":amd: (MI355 DPX) Language Models (Extended Pooling) Shard %N" @@ -4897,6 +5051,7 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py @@ -4960,7 +5115,11 @@ steps: - tests/models/multimodal/generation - tests/models/multimodal/test_mapping.py commands: - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py + - >- + pytest -v -s + models/multimodal/generation + -m 'not core_model' + --ignore models/multimodal/generation/test_common.py - pytest -v -s models/multimodal/test_mapping.py - label: ":amd: (MI355 DPX) Multimodal Models (Extended Generation 3)" # TBD @@ -5000,7 +5159,19 @@ steps: - vllm/ - tests/models/multimodal commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - >- + pytest -v -s + models/multimodal + -m core_model + --ignore models/multimodal/generation/test_common.py + --ignore models/multimodal/generation/test_ultravox.py + --ignore models/multimodal/generation/test_qwen2_5_vl.py + --ignore models/multimodal/generation/test_qwen2_vl.py + --ignore models/multimodal/generation/test_whisper.py + --ignore models/multimodal/generation/test_mm_prefix_lm.py + --ignore models/multimodal/generation/test_memory_leak.py + --ignore models/multimodal/generation/test_vit_cudagraph.py + --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model @@ -5038,28 +5209,29 @@ steps: #------------------------------------------------------- mi355 · quantization --------------------------------------------------------# -- label: ":amd: (MI355 DPX) Quantization Shard %N" - timeout_in_minutes: 115 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - tests/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/layers/fused_moe/oracle/unquantized.py - - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - - tests/rocm/test_moe_weight_replay.py - optional: true - commands: - - uv pip install --system torchao==0.18.0 - - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi +# Disabled until torchao transitive dependencies are locked. +# - label: ":amd: (MI355 DPX) Quantization Shard %N" +# timeout_in_minutes: 115 +# parallelism: 4 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - csrc/ +# - vllm/model_executor/layers/quantization +# - tests/quantization +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# - vllm/model_executor/layers/fused_moe/oracle/unquantized.py +# - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +# - tests/rocm/test_moe_weight_replay.py +# optional: true +# commands: +# - uv pip install --system torchao==0.18.0 +# - uv pip install --system conch-triton-kernels +# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi #--------------------------------------------------------- mi355 · samplers ----------------------------------------------------------# @@ -5102,7 +5274,6 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py @@ -5135,7 +5306,9 @@ steps: - vllm/v1/attention/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/kernels/linear/ @@ -5364,69 +5537,71 @@ steps: - pytest -v -s v1/simple_kv_offload optional: true -- label: ":amd: (MI355 DPX) V1 KV Connectors Shard %N" - timeout_in_minutes: 60 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - working_dir: /vllm-workspace/tests - source_file_dependencies: - - vllm/ - - tests/v1/core - - tests/v1/executor - - tests/v1/kv_offload - - tests/v1/simple_kv_offload - - tests/v1/worker - - tests/v1/streaming_input - - tests/v1/kv_connector/unit - - tests/v1/ec_connector/unit - - tests/v1/metrics - - tests/entrypoints/openai/correctness/test_lmeval.py - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/entrypoints/pooling/ - - vllm/inputs/ - - vllm/lora/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/outputs.py - - vllm/platforms/ - - vllm/pooling_params.py - - vllm/profiler/ - - vllm/sampling_params.py - - vllm/tokenizers/ - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - optional: true +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355 DPX) V1 KV Connectors Shard %N" +# timeout_in_minutes: 60 +# parallelism: 4 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# working_dir: /vllm-workspace/tests +# source_file_dependencies: +# - vllm/ +# - tests/v1/core +# - tests/v1/executor +# - tests/v1/kv_offload +# - tests/v1/simple_kv_offload +# - tests/v1/worker +# - tests/v1/streaming_input +# - tests/v1/kv_connector/unit +# - tests/v1/ec_connector/unit +# - tests/v1/metrics +# - tests/entrypoints/openai/correctness/test_lmeval.py +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/entrypoints/pooling/ +# - vllm/inputs/ +# - vllm/lora/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/outputs.py +# - vllm/platforms/ +# - vllm/pooling_params.py +# - vllm/profiler/ +# - vllm/sampling_params.py +# - vllm/tokenizers/ +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB +# - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB +# optional: true -- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" - timeout_in_minutes: 30 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - num_gpus: 1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ - - vllm/distributed/kv_transfer/kv_connector/v1/base.py - - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py - - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py - - tests/v1/kv_connector/unit/test_moriio_ - - tests/v1/kv_connector/unit/utils.py - - requirements/kv_connectors_rocm.txt - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" +# timeout_in_minutes: 30 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# num_gpus: 1 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ +# - vllm/distributed/kv_transfer/kv_connector/v1/base.py +# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +# - tests/v1/kv_connector/unit/test_moriio_ +# - tests/v1/kv_connector/unit/utils.py +# - requirements/kv_connectors_rocm.txt +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py - label: ":amd: (MI355 DPX) UMBP Connector Unit Tests" timeout_in_minutes: 30 @@ -5442,33 +5617,61 @@ steps: - vllm/v1/kv_cache_interface.py - tests/v1/kv_connector/unit/test_umbp_ - tests/v1/kv_connector/umbp_test_utils.py + - tests/v1/kv_connector/unit/utils.py + - tests/v1/kv_connector/conftest.py + - docker/Dockerfile.rocm_base commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn # Fail loudly instead of skipping every test if the image lacks mori. - python3 -c "import mori" - pytest -v -s -rs v1/kv_connector/unit/test_umbp_*.py -- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_2 - num_gpus: 2 - parallelism: 6 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ - - vllm/distributed/kv_transfer/kv_connector/v1/base.py - - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py - - vllm/model_executor/model_loader/dummy_loader.py - - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py - - tests/v1/kv_connector/moriio_integration/ - - requirements/kv_connectors_rocm.txt - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_2 +# num_gpus: 2 +# parallelism: 6 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ +# - vllm/distributed/kv_transfer/kv_connector/v1/base.py +# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +# - vllm/model_executor/model_loader/dummy_loader.py +# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +# - tests/v1/kv_connector/moriio_integration/ +# - requirements/kv_connectors_rocm.txt +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +# The same P/D between different TP sizes (WRITE TP2 -> TP1, READ of MLA +# TP1 <-> TP2), which needs 3 GPUs: 8 cases, ~36 min on MI355. +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP" +# timeout_in_minutes: 45 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_4 +# num_gpus: 4 +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ +# - vllm/distributed/kv_transfer/kv_connector/v1/base.py +# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +# - vllm/model_executor/model_loader/dummy_loader.py +# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +# - tests/v1/kv_connector/moriio_integration/ +# - requirements/kv_connectors_rocm.txt +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k heterogeneous_tp - label: ":amd: (MI355 DPX) V1 Metrics + LM Eval" timeout_in_minutes: 65 @@ -5593,7 +5796,10 @@ steps: - vllm/platforms/rocm.py commands: - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py - - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async + - >- + pytest -v -s + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async - label: ":amd: (MI355 DPX) Extract Hidden States Integration Single GPU" timeout_in_minutes: 45 @@ -5614,6 +5820,26 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s -m 'not distributed' tests/v1/kv_connector/extract_hidden_states_integration +- label: ":amd: (MI355 DPX) UMBP Connector Integration" + timeout_in_minutes: 30 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + num_gpus: 1 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/umbp/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - tests/v1/kv_connector/umbp_integration/ + - tests/v1/kv_connector/umbp_test_utils.py + - tests/v1/kv_connector/conftest.py + - tests/v1/kv_connector/rocm_pd_accuracy_utils.py + - docker/Dockerfile.rocm_base + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/umbp_integration + - label: ":amd: (MI355 DPX) V1 Attention Shard %N" # TBD timeout_in_minutes: 125 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -5632,8 +5858,17 @@ steps: - vllm/models/deepseek_v32/amd/rocm.py optional: true commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --ignore=v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py v1/attention/test_sparse_mla_backends.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - >- + pytest -v -s + v1/attention + --ignore=v1/attention/test_mla_backends.py + --ignore=v1/attention/test_sparse_mla_backends.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - >- + VLLM_ROCM_USE_AITER=1 pytest -v -s + v1/attention/test_mla_backends.py + v1/attention/test_sparse_mla_backends.py + --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: ":amd: (MI355 DPX) CUDAGraph" timeout_in_minutes: 55 @@ -5717,7 +5952,13 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/conftest.py commands: - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi + - >- + if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then + pytest -v -s + v1/e2e/spec_decode/acceptance_rates/mtp_other/ + --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" + --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; + fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi @@ -5867,30 +6108,31 @@ steps: #----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# -- label: ":amd: (MI355 DPX) Regression" # TBD - timeout_in_minutes: 25 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - dind: false - agent_pool: mi355_dpx - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/test_regression - - vllm/config/ - - vllm/distributed/ - - vllm/engine/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - commands: - - pip install 'modelscope<1.38' - - pytest -v -s test_regression.py +# Disabled until the modelscope installation is locked. +# - label: ":amd: (MI355 DPX) Regression" # TBD +# timeout_in_minutes: 25 +# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] +# dind: false +# agent_pool: mi355_dpx +# optional: true +# working_dir: "/vllm-workspace/tests" +# source_file_dependencies: +# - vllm/ +# - tests/test_regression +# - vllm/config/ +# - vllm/distributed/ +# - vllm/engine/ +# - vllm/inputs/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/platforms/ +# - vllm/sampling_params.py +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# commands: +# - pip install 'modelscope<1.38' +# - pytest -v -s test_regression.py - label: ':amd: (MI355) DeepSeek V4 MoE Integration' timeout_in_minutes: 20 @@ -5978,3 +6220,4 @@ steps: - vllm/envs.py commands: - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300 + - pytest -v -s tests/jit_monitor/test_flydsl_contract.py diff --git a/.buildkite/test_areas/attention.yaml b/.buildkite/test_areas/attention.yaml index 9b957353ba10..41866227f342 100644 --- a/.buildkite/test_areas/attention.yaml +++ b/.buildkite/test_areas/attention.yaml @@ -33,7 +33,10 @@ steps: - vllm/platforms/rocm.py - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_sparse_mla_backends.py + - >- + pytest -v -s + v1/attention + --ignore=v1/attention/test_sparse_mla_backends.py - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_sparse_mla_backends.py - label: ":nvidia: (B200) V1 Attention" @@ -67,11 +70,19 @@ steps: - vllm/platforms/rocm.py - vllm/models/deepseek_v32/amd/rocm.py commands: - - pytest -v -s v1/attention --ignore=v1/attention/test_mla_backends.py --ignore=v1/attention/test_sparse_mla_backends.py - - VLLM_ROCM_USE_AITER=1 pytest -v -s v1/attention/test_mla_backends.py v1/attention/test_sparse_mla_backends.py + - >- + pytest -v -s + v1/attention + --ignore=v1/attention/test_mla_backends.py + --ignore=v1/attention/test_sparse_mla_backends.py + - >- + VLLM_ROCM_USE_AITER=1 pytest -v -s + v1/attention/test_mla_backends.py + v1/attention/test_sparse_mla_backends.py - label: ":nvidia: (DGX Spark) [Nightly] FlashInfer NVFP4 XQA" key: flashinfer-nvfp4-xqa-dgx-spark + automatic_shard: true timeout_in_minutes: 30 device: dgx-spark optional: true @@ -82,4 +93,4 @@ steps: - vllm/v1/attention/backends/flashinfer.py - tests/v1/attention/test_attention_backends.py commands: - - pytest -v -s v1/attention/test_attention_backends.py -k test_flashinfer_xqa_nvfp4_decode_correctness + - pytest -v -s v1/attention/test_attention_backends.py -k nvfp4_decode_correctness diff --git a/.buildkite/test_areas/basic_correctness.yaml b/.buildkite/test_areas/basic_correctness.yaml index f294301a7b1f..7cf48f6467fa 100644 --- a/.buildkite/test_areas/basic_correctness.yaml +++ b/.buildkite/test_areas/basic_correctness.yaml @@ -26,6 +26,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Basic Correctness CuMem" key: basic-correctness-cumem + automatic_shard: true timeout_in_minutes: 68 device: h200_18gb source_file_dependencies: *basic-correctness-deps @@ -34,6 +35,7 @@ steps: - pytest -v -s basic_correctness/memory/cumem mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Basic Correctness CuMem" dind: false device: mi355_dpx @@ -43,6 +45,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Basic Correctness Sleep Mode" key: basic-correctness-sleep-mode + automatic_shard: true timeout_in_minutes: 68 device: h200_18gb source_file_dependencies: *basic-correctness-deps @@ -51,6 +54,7 @@ steps: - pytest -v -s basic_correctness/memory/sleep_mode mirror: amd: + automatic_shard: true label: ":amd: (MI300) Basic Correctness Sleep Mode" dind: false device: mi300_1 @@ -60,6 +64,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Basic Correctness CPU Offload" key: basic-correctness-cpu-offload + automatic_shard: true timeout_in_minutes: 68 device: h200_18gb source_file_dependencies: *basic-correctness-deps @@ -68,6 +73,7 @@ steps: - pytest -v -s basic_correctness/cpu_offload mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Basic Correctness CPU Offload" dind: false device: mi355_dpx @@ -77,6 +83,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Basic Correctness Prefetch Offload" key: basic-correctness-prefetch-offload + automatic_shard: true timeout_in_minutes: 68 device: h200_18gb source_file_dependencies: *basic-correctness-deps @@ -85,6 +92,7 @@ steps: - pytest -v -s basic_correctness/prefetch_offload mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Basic Correctness Prefetch Offload" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/benchmarks.yaml b/.buildkite/test_areas/benchmarks.yaml index 7c95577e0373..1eca056cbb9c 100644 --- a/.buildkite/test_areas/benchmarks.yaml +++ b/.buildkite/test_areas/benchmarks.yaml @@ -16,6 +16,7 @@ steps: - pytest -v -s benchmarks/ mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Benchmarks CLI" dind: false device: mi355_dpx @@ -34,7 +35,7 @@ steps: - benchmarks/attention_benchmarks/ - vllm/v1/attention/ commands: - - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" + - python3 benchmarks/attention_benchmarks/benchmark.py --backends FLASH_ATTN FLASHINFER --batch-specs "8q1s1k" mirror: amd: label: ":amd: (MI355) Attention Benchmark Smoke" diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 46c966a4e5b4..46d980df398b 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H100) Distributed Compile" key: distributed-compile-unit-tests-2xh100 + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/" device: h100 @@ -17,6 +18,7 @@ steps: - pytest -s -v tests/compile/passes/distributed mirror: amd: + automatic_shard: true label: ':amd: (MI300) Distributed Compile' dind: false device: mi300_2 @@ -36,6 +38,7 @@ steps: - label: ":nvidia: (B200) Fusion and Compile" key: fusion-and-compile-unit-tests-2xb200 + automatic_shard: true timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: b200-k8s @@ -65,6 +68,7 @@ steps: - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) Fusion and Compile FP8" dind: false device: mi355_2 @@ -86,11 +90,16 @@ steps: commands: - amd-smi # MLA static FP8 and distributed AR/RMS are owned by the existing ROCm jobs. - - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py + - >- + pytest -v -s + tests/compile/passes/test_fusion_attn.py + tests/compile/passes/test_silu_mul_quant_fusion.py + tests/compile/passes/test_silu_mul_quant_manual_fusion.py - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py - label: ":nvidia: (H100) Fusion E2E Quick" key: fusion-e2e-quick-h100 + automatic_shard: true timeout_in_minutes: 25 working_dir: "/vllm-workspace/" device: h100 @@ -109,6 +118,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)" mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Fusion E2E Quick" dind: false device: mi355_dpx @@ -131,6 +141,7 @@ steps: - label: ":nvidia: (H100) Fusion E2E Config Sweep" key: fusion-e2e-config-sweep-h100 + automatic_shard: true timeout_in_minutes: 25 working_dir: "/vllm-workspace/" device: h100 @@ -150,6 +161,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3" mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Fusion E2E Config Sweep" dind: false device: mi355_dpx @@ -208,6 +220,7 @@ steps: - label: ":nvidia: (H100) Fusion E2E TP2 Quick" key: fusion-e2e-tp2-quick-h100 + automatic_shard: true timeout_in_minutes: 35 working_dir: "/vllm-workspace/" device: h100 @@ -224,6 +237,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))" mirror: amd: + automatic_shard: true label: ":amd: (MI355) Fusion E2E TP2 Quick" dind: false device: mi355_2 @@ -266,6 +280,7 @@ steps: - label: ":nvidia: (H100) Fusion E2E TP2 AR-RMS Config Sweep" key: fusion-e2e-tp2-ar-rms-config-sweep-h100 + automatic_shard: true timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: h100 @@ -285,6 +300,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3" mirror: amd: + automatic_shard: true label: ":amd: (MI355) Fusion E2E TP2 AR-RMS FP8 Config Sweep" dind: false device: mi355_2 @@ -311,6 +327,7 @@ steps: - label: ":nvidia: (B200) Fusion E2E TP2" key: fusion-e2e-tp2-b200 + automatic_shard: true timeout_in_minutes: 45 working_dir: "/vllm-workspace/" device: b200-k8s @@ -329,6 +346,7 @@ steps: - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4" mirror: amd: + automatic_shard: true label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites" dind: false device: mi355_2 @@ -356,6 +374,7 @@ steps: - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)' key: fusion-e2e-tp2-ar-rms-inductor-partition-amd + automatic_shard: true dind: false device: mi355_2 num_devices: 2 @@ -377,6 +396,7 @@ steps: - label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Dynamo Partition)' key: fusion-e2e-tp2-ar-rms-dynamo-partition-amd + automatic_shard: true dind: false device: mi355_2 num_devices: 2 diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index b3bfa9f09708..bc4afc0d9d36 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.buildkite/test_areas/cuda.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) CUDA Platform" key: platform-tests + automatic_shard: true timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: @@ -19,6 +20,7 @@ steps: - pytest -v -s cuda/test_cuda_compatibility_path.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) CUDA Platform" dind: false device: mi355_dpx @@ -54,6 +56,7 @@ steps: - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py mirror: amd: + automatic_shard: true label: ":amd: (MI250) CUDAGraph" dind: false device: mi250_1 diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 1a6edf2652ce..30560237fed7 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -1,209 +1,215 @@ group: Disaggregated depends_on: - image-build -steps: -- label: ":nvidia: (L4) [Nightly] Distributed NixlConnector PD accuracy" - key: distributed-nixlconnector-pd-accuracy-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 55 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - - GPU_MEMORY_UTILIZATION=0.6 SLEEP_WAKE=1 bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI355) Distributed NixlConnector PD accuracy" - dind: false - device: mi355_4 - timeout_in_minutes: 60 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +steps: [] +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] Distributed NixlConnector PD accuracy" +# key: distributed-nixlconnector-pd-accuracy-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 55 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# - GPU_MEMORY_UTILIZATION=0.6 SLEEP_WAKE=1 bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI355) Distributed NixlConnector PD accuracy" +# dind: false +# device: mi355_4 +# timeout_in_minutes: 60 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: ":nvidia: (L4) [Nightly] Distributed FlashInfer NixlConnector PD accuracy" - key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 55 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" - dind: false - device: mi355_4 - num_devices: 4 - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/tests" - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - tests/v1/kv_connector/rocm_pd_accuracy_utils.py - - tests/utils.py - - vllm/v1/attention/backends/rocm_aiter_unified_attn.py - - vllm/platforms/rocm.py - - requirements/common.txt - - requirements/kv_connectors_rocm.txt - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - # Preserve all six model/TP pairs with native AITER attention. - - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] Distributed FlashInfer NixlConnector PD accuracy" +# key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 55 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy" +# dind: false +# device: mi355_4 +# num_devices: 4 +# timeout_in_minutes: 60 +# working_dir: "/vllm-workspace/tests" +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - tests/v1/kv_connector/rocm_pd_accuracy_utils.py +# - tests/utils.py +# - vllm/v1/attention/backends/rocm_aiter_unified_attn.py +# - vllm/platforms/rocm.py +# - requirements/common.txt +# - requirements/kv_connectors_rocm.txt +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# # Preserve all six model/TP pairs with native AITER attention. +# - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py -- label: ":nvidia: (L4) [Nightly] Push NixlConnector PP prefill PD accuracy" - key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - tests/v1/kv_connector/nixl_push_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" - dind: false - device: mi300_4 - timeout_in_minutes: 55 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - tests/v1/kv_connector/nixl_push_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] Push NixlConnector PP prefill PD accuracy" +# key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 30 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - tests/v1/kv_connector/nixl_push_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI300) Push NixlConnector PP prefill PD accuracy" +# dind: false +# device: mi300_4 +# timeout_in_minutes: 55 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - tests/v1/kv_connector/nixl_push_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh -- label: ":nvidia: (L4) [Nightly] DP EP Distributed NixlConnector PD accuracy" - key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" - dind: false - device: mi355_4 - timeout_in_minutes: 40 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] DP EP Distributed NixlConnector PD accuracy" +# key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 30 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" +# dind: false +# device: mi355_4 +# timeout_in_minutes: 40 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: ":nvidia: (L4) [Nightly] CrossLayer KV layout Distributed NixlConnector PD accuracy" - key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 55 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" - dind: false - device: mi300_4 - timeout_in_minutes: 60 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] CrossLayer KV layout Distributed NixlConnector PD accuracy" +# key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 55 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI300) CrossLayer KV layout Distributed NixlConnector PD accuracy" +# dind: false +# device: mi300_4 +# timeout_in_minutes: 60 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - CROSS_LAYERS_BLOCKS=True ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh -- label: ":nvidia: (L4) [Nightly] Hybrid SSM NixlConnector PD accuracy" - key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" - dind: false - device: mi300_4 - timeout_in_minutes: 80 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - 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 +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] Hybrid SSM NixlConnector PD accuracy" +# key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 60 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" +# dind: false +# device: mi300_4 +# timeout_in_minutes: 80 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - 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 # Per-PR core of the nightly-only 4-GPU sweeps above: hetero-TP (both # directions), MoE DP-EP decode, push-mode PP prefill, sleep/wake, MLA, @@ -211,290 +217,298 @@ steps: # KV sharing on a 2-GPU pod. TP>1 configs share GPUs between prefill and # decode (gpu-memory-utilization 0.35 per engine); TP1/TP1 architecture # configs keep one dedicated GPU per side. -- label: ":nvidia: (L4) NixlConnector PD accuracy core" - key: nixlconnector-pd-accuracy-core-2-gpus - timeout_in_minutes: 90 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 2 - # No AMD mirror: the MI300/MI355 mirrors above are optional: false and - # already cover these paths on PRs. - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - tests/v1/kv_connector/nixl_push_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_core_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) NixlConnector PD accuracy core" +# key: nixlconnector-pd-accuracy-core-2-gpus +# timeout_in_minutes: 90 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 2 +# # No AMD mirror: the MI300/MI355 mirrors above are optional: false and +# # already cover these paths on PRs. +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - tests/v1/kv_connector/nixl_push_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/config_sweep_accuracy_core_test.sh -- label: ":nvidia: (L4) NixlConnector PD edge case" - key: nixlconnector-pd-edge-cases-2-gpus - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/tests" - device: l4 - 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 - mirror: - amd: - label: ":amd: (MI300) NixlConnector PD edge case" - dind: false - device: mi300_2 - timeout_in_minutes: 65 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/v1/core/sched/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) NixlConnector PD edge case" +# key: nixlconnector-pd-edge-cases-2-gpus +# timeout_in_minutes: 40 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# 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 +# mirror: +# amd: +# label: ":amd: (MI300) NixlConnector PD edge case" +# dind: false +# device: mi300_2 +# timeout_in_minutes: 65 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/v1/core/sched/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - PREFILL_GPU_ID=0 DECODE_GPU_ID=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_edge_case_test.sh -- label: ":nvidia: (L4) Hybrid SSM NixlConnector PD prefix cache" - key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus - timeout_in_minutes: 25 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 2 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/v1/core/sched/ - - vllm/v1/core/kv_cache_coordinator.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh - mirror: - amd: - label: ":amd: (MI355) Hybrid SSM NixlConnector PD prefix cache" - dind: false - device: mi355_2 - timeout_in_minutes: 25 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/v1/core/sched/ - - vllm/v1/core/kv_cache_coordinator.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) Hybrid SSM NixlConnector PD prefix cache" +# key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus +# timeout_in_minutes: 25 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 2 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/v1/core/sched/ +# - vllm/v1/core/kv_cache_coordinator.py +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) Hybrid SSM NixlConnector PD prefix cache" +# dind: false +# device: mi355_2 +# timeout_in_minutes: 25 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/v1/core/sched/ +# - vllm/v1/core/kv_cache_coordinator.py +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh -- label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy" - key: multiconnector-nixl-offloading-pd-accuracy-2-gpus - timeout_in_minutes: 40 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 2 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh - mirror: - amd: - label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD accuracy" - dind: false - device: mi355_2 - timeout_in_minutes: 40 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - tests/v1/kv_connector/nixl_integration/ - - requirements/kv_connectors_rocm.txt - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD accuracy" +# key: multiconnector-nixl-offloading-pd-accuracy-2-gpus +# timeout_in_minutes: 40 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 2 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD accuracy" +# dind: false +# device: mi355_2 +# timeout_in_minutes: 40 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ +# - tests/v1/kv_connector/nixl_integration/ +# - requirements/kv_connectors_rocm.txt +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh -- label: ":nvidia: (A100) NixlConnector PD + Spec Decode acceptance" - key: nixlconnector-pd-spec-decode-acceptance-2-gpus - timeout_in_minutes: 45 - device: a100 - working_dir: "/vllm-workspace/tests" - num_devices: 2 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/v1/worker/kv_connector_model_runner_mixin.py - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh - mirror: - amd: - label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" - dind: false - device: mi355_2 - timeout_in_minutes: 45 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/v1/worker/kv_connector_model_runner_mixin.py - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (A100) NixlConnector PD + Spec Decode acceptance" +# key: nixlconnector-pd-spec-decode-acceptance-2-gpus +# timeout_in_minutes: 45 +# device: a100 +# working_dir: "/vllm-workspace/tests" +# num_devices: 2 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/v1/worker/kv_connector_model_runner_mixin.py +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) NixlConnector PD + Spec Decode acceptance" +# dind: false +# device: mi355_2 +# timeout_in_minutes: 45 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/v1/worker/kv_connector_model_runner_mixin.py +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - KV_CACHE_MEMORY_BYTES=8G ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh -- label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD edge cases" - key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus - timeout_in_minutes: 25 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 2 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - tests/v1/kv_connector/nixl_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh - mirror: - amd: - label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD edge cases" - dind: false - device: mi355_2 - timeout_in_minutes: 25 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) MultiConnector (Nixl+Offloading) PD edge cases" +# key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus +# timeout_in_minutes: 25 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 2 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ +# - tests/v1/kv_connector/nixl_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) MultiConnector (Nixl+Offloading) PD edge cases" +# dind: false +# device: mi355_2 +# timeout_in_minutes: 25 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh # P TP 4 - D DPEP 4 test case for DSv4-Flash -- label: ":nvidia: (H200) [Nightly] DSv4-Flash Disaggregated DP EP" - key: dsv4-flash-disaggregated - timeout_in_minutes: 60 - device: h200 - optional: true - working_dir: "/vllm-workspace/tests" - num_devices: 8 - env: - ENABLE_HMA_FLAG: "1" - DP_EP: "1" - GPU_MEMORY_UTILIZATION: "0.85" - PREFILLER_TP_SIZE: "4" - DECODER_TP_SIZE: "4" - PREFILL_BLOCK_SIZE: "256" - DECODE_BLOCK_SIZE: "256" - MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" - VLLM_ENGINE_READY_TIMEOUT_S: "1800" - VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" - # The V4 indexer only supports block-outermost layouts; let it pick. - PREFILLER_KV_LAYOUT: "" - DECODER_KV_LAYOUT: "" - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - mirror: - amd: - label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" - dind: false - device: mi355_8 - timeout_in_minutes: 85 - depends_on: - - image-build-amd - source_file_dependencies: - - requirements/kv_connectors_rocm.txt - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/models/deepseek_v4/ - - vllm/v1/attention/ - - tests/v1/kv_connector/nixl_integration/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - # A mirror does not inherit the parent env, so it is repeated here in - # the parent's order. The layouts must stay empty: unset, the script - # defaults to HND, which the V4 indexer rejects. - - >- - VLLM_ROCM_USE_AITER=1 - ENABLE_HMA_FLAG=1 - DP_EP=1 - GPU_MEMORY_UTILIZATION=0.85 - PREFILLER_TP_SIZE=4 - DECODER_TP_SIZE=4 - PREFILL_BLOCK_SIZE=256 - DECODE_BLOCK_SIZE=256 - MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash - VLLM_ENGINE_READY_TIMEOUT_S=1800 - VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8 - PREFILLER_KV_LAYOUT= - DECODER_KV_LAYOUT= - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (H200) [Nightly] DSv4-Flash Disaggregated DP EP" +# key: dsv4-flash-disaggregated +# timeout_in_minutes: 60 +# device: h200 +# optional: true +# working_dir: "/vllm-workspace/tests" +# num_devices: 8 +# env: +# ENABLE_HMA_FLAG: "1" +# DP_EP: "1" +# GPU_MEMORY_UTILIZATION: "0.85" +# PREFILLER_TP_SIZE: "4" +# DECODER_TP_SIZE: "4" +# PREFILL_BLOCK_SIZE: "256" +# DECODE_BLOCK_SIZE: "256" +# MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" +# VLLM_ENGINE_READY_TIMEOUT_S: "1800" +# VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" +# # The V4 indexer only supports block-outermost layouts; let it pick. +# PREFILLER_KV_LAYOUT: "" +# DECODER_KV_LAYOUT: "" +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" +# dind: false +# device: mi355_8 +# timeout_in_minutes: 85 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - requirements/kv_connectors_rocm.txt +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/model_executor/layers/fused_moe/ +# - vllm/model_executor/layers/quantization/ +# - vllm/models/deepseek_v4/ +# - vllm/v1/attention/ +# - tests/v1/kv_connector/nixl_integration/ +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# # A mirror does not inherit the parent env, so it is repeated here in +# # the parent's order. The layouts must stay empty: unset, the script +# # defaults to HND, which the V4 indexer rejects. +# - >- +# VLLM_ROCM_USE_AITER=1 +# ENABLE_HMA_FLAG=1 +# DP_EP=1 +# GPU_MEMORY_UTILIZATION=0.85 +# PREFILLER_TP_SIZE=4 +# DECODER_TP_SIZE=4 +# PREFILL_BLOCK_SIZE=256 +# DECODE_BLOCK_SIZE=256 +# MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash +# VLLM_ENGINE_READY_TIMEOUT_S=1800 +# VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8 +# PREFILLER_KV_LAYOUT= +# DECODER_KV_LAYOUT= +# bash v1/kv_connector/nixl_integration/run_accuracy_test.sh -- label: ":nvidia: (H200) [Nightly] Kimi-Linear-48B-A3B Disaggregated DP EP" - key: kimi-linear-disaggregated - timeout_in_minutes: 60 - device: h200 - optional: true - working_dir: "/vllm-workspace/tests" - num_devices: 4 - env: - ENABLE_HMA_FLAG: "1" - DP_EP: "1" - GPU_MEMORY_UTILIZATION: "0.9" - VLLM_SSM_CONV_STATE_LAYOUT: "DS" - PREFILLER_TP_SIZE: "2" - DECODER_TP_SIZE: "2" - MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct" - # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. - PREFILL_BLOCK_SIZE: "2048" - DECODE_BLOCK_SIZE: "2048" - VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768" - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh - mirror: - amd: - label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" - dind: false - device: mi355_4 - num_devices: 4 - timeout_in_minutes: 105 - working_dir: /vllm-workspace/tests - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/platforms/rocm.py - - vllm/models/kimi_k3/ - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - tests/v1/kv_connector/nixl_integration/ - - requirements/kv_connectors_rocm.txt - env: - VLLM_WORKER_MULTIPROC_METHOD: "spawn" - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (H200) [Nightly] Kimi-Linear-48B-A3B Disaggregated DP EP" +# key: kimi-linear-disaggregated +# timeout_in_minutes: 60 +# device: h200 +# optional: true +# working_dir: "/vllm-workspace/tests" +# num_devices: 4 +# env: +# ENABLE_HMA_FLAG: "1" +# DP_EP: "1" +# GPU_MEMORY_UTILIZATION: "0.9" +# VLLM_SSM_CONV_STATE_LAYOUT: "DS" +# PREFILLER_TP_SIZE: "2" +# DECODER_TP_SIZE: "2" +# MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct" +# # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. +# PREFILL_BLOCK_SIZE: "2048" +# DECODE_BLOCK_SIZE: "2048" +# VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768" +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) Kimi-Linear-48B-A3B Disaggregated DP EP" +# dind: false +# device: mi355_4 +# num_devices: 4 +# timeout_in_minutes: 105 +# working_dir: /vllm-workspace/tests +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/platforms/rocm.py +# - vllm/models/kimi_k3/ +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - tests/v1/kv_connector/nixl_integration/ +# - requirements/kv_connectors_rocm.txt +# env: +# VLLM_WORKER_MULTIPROC_METHOD: "spawn" +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/disaggregated_mooncake.yaml b/.buildkite/test_areas/disaggregated_mooncake.yaml index 87c004ecfb5e..d5ced93f34b2 100644 --- a/.buildkite/test_areas/disaggregated_mooncake.yaml +++ b/.buildkite/test_areas/disaggregated_mooncake.yaml @@ -2,50 +2,51 @@ group: Disaggregated Mooncake depends_on: - image-build steps: -- label: ":nvidia: (B200) [Nightly] Distributed MooncakeConnector PD accuracy" - key: distributed-mooncakeconnector-pd-accuracy-4-gpus - # Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 10 - working_dir: "/vllm-workspace/tests" - # Runs on b200-k8s because it is the only NVIDIA pool that can carry mooncake - # transfers: its pods are the only ones granted IPC_LOCK, which RDMA needs to - # pin memory, and the only ones using host networking. Elsewhere topology - # discovery finds no usable HCA and the tcp fallback then advertises the - # docker bridge address, so every transfer fails to connect. - device: b200-k8s - num_devices: 4 - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/ - - tests/v1/kv_connector/mooncake_integration/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh - mirror: - amd: - label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" - optional: true - dind: false - device: mi355_4 - num_devices: 4 - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/ - - tests/v1/kv_connector/mooncake_integration/ - - tests/v1/kv_connector/rocm_pd_accuracy_utils.py - - tests/utils.py - - examples/disaggregated/mooncake_connector/ - - vllm/v1/attention/backends/rocm_aiter_unified_attn.py - - vllm/platforms/rocm.py - - requirements/common.txt - - requirements/kv_connectors_rocm.txt - commands: - - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - # Same full accuracy/TP sweep, using same-host HIP GPU transport. - - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (B200) [Nightly] Distributed MooncakeConnector PD accuracy" +# key: distributed-mooncakeconnector-pd-accuracy-4-gpus +# # Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 10 +# working_dir: "/vllm-workspace/tests" +# # Runs on b200-k8s because it is the only NVIDIA pool that can carry mooncake +# # transfers: its pods are the only ones granted IPC_LOCK, which RDMA needs to +# # pin memory, and the only ones using host networking. Elsewhere topology +# # discovery finds no usable HCA and the tcp fallback then advertises the +# # docker bridge address, so every transfer fails to connect. +# device: b200-k8s +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/ +# - tests/v1/kv_connector/mooncake_integration/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh +# mirror: +# amd: +# label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy" +# optional: true +# dind: false +# device: mi355_4 +# num_devices: 4 +# timeout_in_minutes: 30 +# working_dir: "/vllm-workspace/tests" +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/ +# - tests/v1/kv_connector/mooncake_integration/ +# - tests/v1/kv_connector/rocm_pd_accuracy_utils.py +# - tests/utils.py +# - examples/disaggregated/mooncake_connector/ +# - vllm/v1/attention/backends/rocm_aiter_unified_attn.py +# - vllm/platforms/rocm.py +# - requirements/common.txt +# - requirements/kv_connectors_rocm.txt +# commands: +# - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt +# # Same full accuracy/TP sweep, using same-host HIP GPU transport. +# - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py - label: ":nvidia: (L4) Mooncake EC TCP E2E" key: mooncake-ec-tcp-e2e-2-gpus diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 303798f25fb2..ef3d5e339f09 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (L4) Distributed Comm Ops" key: distributed-comm-ops + automatic_shard: true timeout_in_minutes: 25 working_dir: "/vllm-workspace/tests" device: l4 @@ -13,6 +14,7 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/distributed commands: + - pytest -v -s distributed/test_localization.py - pytest -v -s distributed/test_comm_ops.py - pytest -v -s distributed/test_cpu_all_reduce.py - pytest -v -s distributed/test_shm_broadcast.py @@ -20,6 +22,7 @@ steps: - pytest -v -s distributed/test_shm_storage.py mirror: amd: + automatic_shard: true label: ':amd: (MI300) Distributed Comm Ops' dind: false device: mi300_2 @@ -39,6 +42,7 @@ steps: - label: ":nvidia: (L4) Sharded RDT Weight Transfer" key: sharded-rdt-weight-transfer + automatic_shard: true device: l4 timeout_in_minutes: 20 working_dir: "/vllm-workspace/tests" @@ -54,6 +58,7 @@ steps: - pytest -v -s distributed/test_sharded_rdt_trainer.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Sharded RDT Weight Transfer" dind: false device: mi355_dpx @@ -129,6 +134,7 @@ steps: - label: ":nvidia: (L4) Distributed Compile + RPC" key: distributed-compile-rpc-tests-2-gpus + automatic_shard: true timeout_in_minutes: 65 working_dir: "/vllm-workspace/tests" device: l4 @@ -143,14 +149,17 @@ steps: - vllm/v1/engine/ - vllm/v1/worker/ - tests/compile/test_wrapper.py + - tests/compile/test_aot_compile.py - tests/entrypoints/llm/test_collective_rpc.py commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 - pytest -v -s entrypoints/llm/test_collective_rpc.py - pytest -v -s ./compile/test_wrapper.py + - pytest -v -s ./compile/test_aot_compile.py::test_load_on_another_device mirror: amd: + automatic_shard: true label: ":amd: (MI355) Distributed Compile + RPC" dind: false device: mi355_2 @@ -175,6 +184,7 @@ steps: - label: ":nvidia: (L4) Distributed Torchrun + Shutdown" key: distributed-torchrun-shutdown-tests-2-gpus + automatic_shard: true timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" device: l4 @@ -199,6 +209,7 @@ steps: - pytest -v -s v1/worker/test_worker_memory_snapshot.py mirror: amd: + automatic_shard: true label: ':amd: (MI300) Distributed Torchrun + Shutdown' dind: false device: mi300_2 @@ -223,88 +234,89 @@ steps: - HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - pytest -v -s v1/worker/test_worker_memory_snapshot.py -- label: ":nvidia: (L4) [Nightly] Distributed Torchrun + Examples" - key: distributed-torchrun-examples-4-gpus - # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 45 - working_dir: "/vllm-workspace" - device: l4 - num_devices: 4 - source_file_dependencies: - - vllm/distributed/ - - "!vllm/distributed/kv_transfer/" - - tests/distributed/test_torchrun_example.py - - tests/distributed/test_torchrun_example_moe.py - - examples/rl/ - - examples/features/data_parallel/data_parallel_offline.py - commands: - # https://github.com/NVIDIA/nccl/issues/1838 - - export NCCL_CUMEM_HOST_ENABLE=0 - # test with torchrun tp=2 and external_dp=2 - - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py - # test with torchrun tp=2 and pp=2 - - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py - # test with torchrun tp=4 and dp=1 - - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - # test with torchrun tp=2, pp=2 and dp=1 - - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - # test with torchrun tp=1 and dp=4 with ep - - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - # test with torchrun tp=2 and dp=2 with ep - - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - # test with internal dp - - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager - # rlhf examples (each launches its own `vllm serve` and tears it down) - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py - # sharded RDT: 2 FSDP2 trainer ranks -> 2 vLLM DP ranks with EP. Its data - # plane is NIXL, which ships only in the kv-connector set, and the script - # also drops the nixl-cu wheel variant that does not match this image. Kept - # last so the commands above run in an unmodified environment. RDT_MODEL swaps - # the script's reference MoE for a tiny one: no weights are downloaded either - # way, but every trainer rank builds the model before sharding it, and the - # default does not fit an L4. - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py - mirror: - amd: - # Run on PRs even when the L4 step is optional (nightly only). - optional: false - label: ":amd: (MI355) Distributed Torchrun + Examples" - dind: false - device: mi355_4 - timeout_in_minutes: 55 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/ - - "!vllm/distributed/kv_transfer/" - - tests/distributed/test_torchrun_example.py - - tests/distributed/test_torchrun_example_moe.py - - examples/rl/ - - examples/features/data_parallel/data_parallel_offline.py - - .buildkite/scripts/install-kv-connectors.sh - - requirements/kv_connectors_rocm.txt - - requirements/test/rocm.in - - requirements/test/rocm.txt - - vllm/platforms/rocm.py - commands: - # Sharded RDT needs Ray >= 2.56; keep the ROCm image's dependencies. - - uv pip install --system --no-deps 'ray==2.56.1' - - export NCCL_CUMEM_HOST_ENABLE=0 - - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py - - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py - - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py - # Loading weights from the shared HF cache can exceed the default 300s --timeout. - - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager --timeout 900 - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py - - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (L4) [Nightly] Distributed Torchrun + Examples" +# key: distributed-torchrun-examples-4-gpus +# # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 45 +# working_dir: "/vllm-workspace" +# device: l4 +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/ +# - "!vllm/distributed/kv_transfer/" +# - tests/distributed/test_torchrun_example.py +# - tests/distributed/test_torchrun_example_moe.py +# - examples/rl/ +# - examples/features/data_parallel/data_parallel_offline.py +# commands: +# # https://github.com/NVIDIA/nccl/issues/1838 +# - export NCCL_CUMEM_HOST_ENABLE=0 +# # test with torchrun tp=2 and external_dp=2 +# - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py +# # test with torchrun tp=2 and pp=2 +# - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py +# # test with torchrun tp=4 and dp=1 +# - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# # test with torchrun tp=2, pp=2 and dp=1 +# - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# # test with torchrun tp=1 and dp=4 with ep +# - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# # test with torchrun tp=2 and dp=2 with ep +# - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# # test with internal dp +# - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager +# # rlhf examples (each launches its own `vllm serve` and tears it down) +# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py +# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py +# # sharded RDT: 2 FSDP2 trainer ranks -> 2 vLLM DP ranks with EP. Its data +# # plane is NIXL, which ships only in the kv-connector set, and the script +# # also drops the nixl-cu wheel variant that does not match this image. Kept +# # last so the commands above run in an unmodified environment. RDT_MODEL swaps +# # the script's reference MoE for a tiny one: no weights are downloaded either +# # way, but every trainer rank builds the model before sharding it, and the +# # default does not fit an L4. +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py +# mirror: +# amd: +# # Run on PRs even when the L4 step is optional (nightly only). +# optional: false +# label: ":amd: (MI355) Distributed Torchrun + Examples" +# dind: false +# device: mi355_4 +# timeout_in_minutes: 55 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/ +# - "!vllm/distributed/kv_transfer/" +# - tests/distributed/test_torchrun_example.py +# - tests/distributed/test_torchrun_example_moe.py +# - examples/rl/ +# - examples/features/data_parallel/data_parallel_offline.py +# - .buildkite/scripts/install-kv-connectors.sh +# - requirements/kv_connectors_rocm.txt +# - requirements/test/rocm.in +# - requirements/test/rocm.txt +# - vllm/platforms/rocm.py +# commands: +# # Sharded RDT needs Ray >= 2.56; keep the ROCm image's dependencies. +# - uv pip install --system --no-deps 'ray==2.56.1' +# - export NCCL_CUMEM_HOST_ENABLE=0 +# - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py +# - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py +# - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py +# # Loading weights from the shared HF cache can exceed the default 300s --timeout. +# - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager --timeout 900 +# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_nccl.py +# - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_http_ipc.py +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - RDT_MODEL=hmellor/tiny-random-DeepseekV2ForCausalLM python3 examples/rl/rlhf_sharded_rdt_small_ep.py - label: ":nvidia: (L4) [Nightly] Distributed DP Extended" key: distributed-dp-tests-4-gpus @@ -367,6 +379,7 @@ steps: - label: ":nvidia: (L4) [Nightly] Distributed Compile + Comm" key: distributed-compile-comm-4-gpus + automatic_shard: true # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. optional: true timeout_in_minutes: 70 @@ -390,6 +403,7 @@ steps: - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node mirror: amd: + automatic_shard: true # Run on PRs even when the L4 step is optional (nightly only). optional: false label: ':amd: (MI355) Distributed Compile + Comm' @@ -414,6 +428,7 @@ steps: - label: ":nvidia: (H100) [Nightly] Distributed DP + EP" key: distributed-tests-8xh100 + automatic_shard: true timeout_in_minutes: 20 optional: true device: h100 @@ -428,14 +443,18 @@ steps: - vllm/v1/executor/uniproc_executor.py - vllm/v1/worker/gpu_worker.py - tests/distributed/test_mnnvl_alltoall.py + - tests/distributed/test_eplb_execute.py commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 # test with torchrun tp=2 and dp=4 with ep - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep + # The L4 EPLB Execution job has no NVLink, so it skips symm_mem. + - pytest -v -s distributed/test_eplb_execute.py -k symm_mem mirror: amd: + automatic_shard: true label: ":amd: (MI355) Distributed DP + EP" dind: false device: mi355_8 @@ -460,43 +479,34 @@ steps: - export TORCH_CPP_LOG_LEVEL=INFO - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep -- label: ":nvidia: (A100) [Nightly] Distributed" - key: distributed-tests-4xa100 - device: a100 +# Key kept from the former A100 step's mirror for job history and retries. +- label: ':amd: (MI355) Distributed' + key: amd-distributed-tests-4xa100 + dind: false + device: mi355_4 optional: true num_devices: 4 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" + - tests/models/test_vision.py + - vllm/platforms/rocm.py + - tests/distributed/ + - tests/basic_correctness/ + - tests/lora/test_mixtral.py commands: # NOTE: don't test llama model here, it seems hf implementation is buggy # see https://github.com/vllm-project/vllm/pull/5689 for details - pytest -v -s distributed/test_custom_all_reduce.py - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s distributed/test_weight_transfer.py -k multiple_sources + - TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py - mirror: - amd: - label: ':amd: (MI355) Distributed' - dind: false - device: mi355_4 - timeout_in_minutes: 40 - working_dir: /vllm-workspace/tests - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/ - - tests/models/test_vision.py - - vllm/platforms/rocm.py - - tests/distributed/ - - tests/basic_correctness/ - - tests/lora/test_mixtral.py - commands: - - pytest -v -s distributed/test_custom_all_reduce.py - - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - - TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - - pytest -v -s -x lora/test_mixtral.py - - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' + - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' - label: ":nvidia: (H100) [Nightly] Distributed Features" key: distributed-tests-2xh100-2xmi300 @@ -570,6 +580,7 @@ steps: - label: ":nvidia: (B200) [Nightly] Distributed" key: distributed-tests-2xb200 + automatic_shard: true device: b200-k8s optional: true working_dir: "/vllm-workspace/" @@ -582,6 +593,7 @@ steps: - pytest -v -s tests/distributed/test_mnnvl_alltoall.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) Distributed AgRs All2All" dind: false device: mi355_2 @@ -604,6 +616,7 @@ steps: - label: ":nvidia: (B200) GEMM-RS/AR" key: gemm-rs-ar-2xb200 + automatic_shard: true timeout_in_minutes: 20 device: b200-k8s working_dir: "/vllm-workspace/" @@ -616,6 +629,7 @@ steps: - pytest -v -s tests/kernels/test_gemm_rs_ar.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) GEMM-RS-AR" dind: false device: mi355_4 @@ -660,7 +674,15 @@ steps: - tests/distributed/ - examples/features/data_parallel/data_parallel_offline.py commands: - - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" + - >- + ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG + "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && + NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && + python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && + VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" + "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && + NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && + python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code" - label: ":nvidia: (L4) [Nightly] Pipeline + Context Parallelism" key: pipeline-context-parallelism-4-gpus @@ -681,16 +703,20 @@ steps: - vllm/v1/worker/gpu_worker.py - tests/distributed/ - tests/v1/distributed/test_pp_dp_v2.py + - vllm/benchmarks/forward/ + - tests/benchmarks/test_forward_cli.py commands: - pytest -v -s distributed/test_pp_cudagraph.py - pytest -v -s distributed/test_pipeline_parallel.py - pytest -v -s v1/distributed/test_pp_dp_v2.py + - pytest -v -s benchmarks/test_forward_cli.py::test_bench_forward_pipeline mirror: amd: # Run on PRs even when the L4 step is optional (nightly only). optional: false label: ":amd: (MI250) Pipeline + Context Parallelism" dind: false + automatic_shard: true device: mi250_4 timeout_in_minutes: 90 depends_on: @@ -760,11 +786,13 @@ steps: - label: ":nvidia: (L4) Distributed Preload" key: preload-2-gpus + automatic_shard: true timeout_in_minutes: 15 working_dir: "/vllm-workspace/tests" device: l4 num_devices: 2 source_file_dependencies: + - vllm/entrypoints/cli/preload.py - vllm/model_executor/model_loader/weight_cache/ - vllm/model_executor/layers/fused_moe/expert_map_manager.py - tests/model_executor/model_loader/test_weight_cache.py @@ -772,6 +800,7 @@ steps: - pytest -v -s model_executor/model_loader/test_weight_cache.py -m 'distributed(num_gpus=2)' mirror: amd: + automatic_shard: true label: ":amd: (MI355) Distributed Preload" dind: false device: mi355_2 @@ -780,6 +809,7 @@ steps: depends_on: - image-build-amd source_file_dependencies: + - vllm/entrypoints/cli/preload.py - vllm/model_executor/model_loader/weight_cache/ - vllm/model_executor/layers/fused_moe/expert_map_manager.py - tests/model_executor/model_loader/test_weight_cache.py diff --git a/.buildkite/test_areas/docker.yaml b/.buildkite/test_areas/docker.yaml index 4c80f26ca52e..8e08f158c64b 100644 --- a/.buildkite/test_areas/docker.yaml +++ b/.buildkite/test_areas/docker.yaml @@ -9,12 +9,25 @@ steps: source_file_dependencies: - .buildkite/release-pipeline.yaml - .buildkite/scripts/docker-build-metadata-args.sh + - .buildkite/scripts/rocm/build.sh + - .buildkite/scripts/rocm/publish.sh + - tools/vllm-rocm/therock_wheels.py - docker/Dockerfile - docker/Dockerfile.cpu + - docker/Dockerfile.rocm + - docker/Dockerfile.rocm_base - docker/docker-bake.hcl + - docker/docker-bake-rocm.hcl - tests/tools/test_docker_build_metadata_args.py + - tests/tools/test_rocm_build.py + - tests/tools/test_amd_ci_fast_exit.py + - .buildkite/scripts/hardware_ci/amd_ci_fast_exit.py + - tests/tools/test_kernel_symbol_map.py + - tools/ci/kernel_symbol_map.py + - tools/ci/rocm_kernel_symbol_map.py commands: - - pytest -v -s tools/test_docker_build_metadata_args.py + - pytest -v -s tools/test_docker_build_metadata_args.py tools/test_rocm_build.py tools/test_amd_ci_fast_exit.py + - pytest -v -s tools/test_kernel_symbol_map.py mirror: amd: label: ":amd: (MI250) Docker Build Metadata" @@ -26,9 +39,16 @@ steps: - image-build-amd source_file_dependencies: - .buildkite/scripts/docker-build-metadata-args.sh - - .buildkite/scripts/ci-bake-rocm.sh - docker/Dockerfile.rocm - docker/Dockerfile.rocm_base - - docker/ci-rocm.hcl - docker/docker-bake-rocm.hcl + - .buildkite/scripts/rocm/build.sh + - .buildkite/scripts/rocm/publish.sh + - tools/vllm-rocm/therock_wheels.py - tests/tools/test_docker_build_metadata_args.py + - tests/tools/test_rocm_build.py + - tests/tools/test_amd_ci_fast_exit.py + - .buildkite/scripts/hardware_ci/amd_ci_fast_exit.py + - tests/tools/test_kernel_symbol_map.py + - tools/ci/kernel_symbol_map.py + - tools/ci/rocm_kernel_symbol_map.py diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 7f42f3cb0019..f2dc9e6041fc 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Engine Core" key: engine + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -26,11 +27,20 @@ steps: - tests/jit_monitor/test_hooks.py - tests/jit_monitor/test_hooks_gpu.py commands: - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py + - >- + pytest -v -s + engine + test_sequence.py + test_config.py + test_logger.py + test_vllm_port.py + jit_monitor/test_hooks.py + jit_monitor/test_hooks_gpu.py mirror: amd: label: ":amd: (MI355 DPX) Engine Core" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 40 depends_on: @@ -48,10 +58,14 @@ steps: - 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/engine + --ignore v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/test_tensor_ipc_queue.py mirror: amd: + automatic_shard: true label: ":amd: (MI250) V1 Engine" dind: false device: mi250_1 @@ -62,7 +76,10 @@ steps: - vllm/platforms/rocm.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/engine + --ignore v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/test_tensor_ipc_queue.py - pytest -v -s config/test_mp_reducer.py @@ -90,6 +107,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) E2E Scheduling Accuracy" key: e2e-scheduling-accuracy-1-gpu + automatic_shard: true timeout_in_minutes: 53 device: h200_18gb source_file_dependencies: @@ -100,6 +118,7 @@ steps: mirror: amd: label: ":amd: (MI250) E2E Scheduling Accuracy" + automatic_shard: true dind: false device: mi250_1 timeout_in_minutes: 55 @@ -117,13 +136,26 @@ steps: - vllm/v1/ - tests/v1/e2e/general/ commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_async_scheduling_accuracy.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2" + - >- + pytest -v -s + v1/e2e/general + --ignore v1/e2e/general/test_async_scheduling.py + --ignore v1/e2e/general/test_async_scheduling_accuracy.py + --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py + -k "not test_mamba_prefix_cache_mrv2" mirror: amd: + automatic_shard: true label: ":amd: (MI250) E2E Core" dind: false commands: - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_async_scheduling_accuracy.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2" + - >- + pytest -v -s + v1/e2e/general + --ignore v1/e2e/general/test_async_scheduling.py + --ignore v1/e2e/general/test_async_scheduling_accuracy.py + --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py + -k "not test_mamba_prefix_cache_mrv2" device: mi250_1 timeout_in_minutes: 50 depends_on: @@ -137,15 +169,21 @@ steps: - label: ":nvidia: (H200 MIG 35GB) E2E Core Large Memory" device: h200_35gb key: e2e-core-large-memory + automatic_shard: true timeout_in_minutes: 50 source_file_dependencies: - vllm/v1/ - tests/v1/e2e/general/ commands: - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py - - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async_flashinfer + - >- + pytest -v -s + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async_flashinfer mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) E2E Core Large Memory" dind: false device: mi355_dpx @@ -160,7 +198,10 @@ steps: - vllm/platforms/rocm.py commands: - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py - - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async + - >- + pytest -v -s + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 + v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async - label: ":nvidia: (H100) ReplaySSM E2E" device: h100 @@ -183,9 +224,22 @@ steps: - vllm/v1/worker/gpu_worker.py commands: - pytest -v -s v1/e2e/test_replayssm_decode.py + mirror: + amd: + automatic_shard: true + label: ":amd: (MI355) ReplaySSM E2E" + dind: false + device: mi355_2 + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/platforms/rocm.py - label: ":nvidia: (L4) [Nightly] V1 E2E" key: v1-e2e-2-gpus + automatic_shard: true timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability optional: true device: l4 @@ -216,8 +270,10 @@ steps: pytest -v -s v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism + v1/e2e/spec_decode/draft_model/test_structured_output_pp.py::test_specdec_pp_structured_output mirror: amd: + automatic_shard: true label: ":amd: (MI300) V1 E2E" dind: false device: mi300_2 @@ -227,6 +283,7 @@ steps: - label: ":nvidia: (B200) [Nightly] V1 E2E Hybrid Chunked Prefill" key: v1-e2e-hybrid-chunked-prefill-b200 + automatic_shard: true timeout_in_minutes: 35 device: b200-k8s num_devices: 1 @@ -240,6 +297,7 @@ steps: - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) V1 E2E Hybrid Chunked Prefill" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 1247e485f2d4..72b993425e12 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -63,7 +63,12 @@ steps: - tests/entrypoints/llm commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode + - >- + pytest -v -s + entrypoints/llm + --ignore=entrypoints/llm/test_generate.py + --ignore=entrypoints/llm/test_collective_rpc.py + --ignore=entrypoints/llm/offline_mode - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests mirror: @@ -89,7 +94,11 @@ steps: - tests/entrypoints/scale_out commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - >- + pytest -v -s + entrypoints/serve + --ignore=entrypoints/serve/dev/rpc + --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc; fi - pytest -v -s entrypoints/scale_out --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: @@ -112,8 +121,17 @@ steps: - tests/entrypoints/openai commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py - - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness + - >- + pytest -v -s + entrypoints/openai/completion + --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py + - >- + pytest -v -s + entrypoints/openai + --ignore=entrypoints/openai/completion + --ignore=entrypoints/openai/chat_completion + --ignore=entrypoints/openai/responses + --ignore=entrypoints/openai/correctness mirror: amd: label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API completion)" @@ -138,6 +156,7 @@ steps: - pytest -v -s entrypoints/openai/chat_completion mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Entrypoints Integration (OpenAI API chat_completion)" dind: false device: mi355_dpx @@ -167,6 +186,7 @@ steps: - pytest -v -s entrypoints/cohere mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Entrypoints Integration (API Server Generate)" dind: false device: mi355_dpx @@ -236,6 +256,7 @@ steps: amd: label: ":amd: (MI355 DPX) Entrypoints Integration (Multimodal)" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 55 depends_on: @@ -267,6 +288,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) OpenAI API Correctness" key: openai-api-correctness + automatic_shard: true timeout_in_minutes: 20 device: h200_18gb source_file_dependencies: @@ -276,6 +298,7 @@ steps: - pytest -s entrypoints/openai/correctness/ mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) OpenAI API Correctness" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 8fc86e65b98e..e9c13dad6348 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) EPLB Algorithm" key: eplb-algorithm + automatic_shard: true timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -18,6 +19,7 @@ steps: - pytest -v -s distributed/test_eplb_utils.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) EPLB Algorithm" dind: false device: mi355_dpx @@ -32,6 +34,7 @@ steps: - label: ":nvidia: (L4) [Nightly] EPLB Execution" key: eplb-execution + automatic_shard: true # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. optional: true timeout_in_minutes: 25 @@ -46,6 +49,7 @@ steps: - pytest -v -s distributed/test_eplb_spec_decode.py mirror: amd: + automatic_shard: true # Run on PRs even when the L4 step is optional (nightly only). optional: false label: ":amd: (MI355) EPLB Execution" @@ -60,21 +64,22 @@ steps: - tests/distributed/test_eplb_spec_decode.py - vllm/platforms/rocm.py -- label: ":nvidia: (H100) [Nightly] Elastic EP Scaling" - key: elastic-ep-scaling-test - # Nightly only: 4-GPU H100 pods wait behind smaller jobs on per-PR runs. - optional: true - timeout_in_minutes: 40 - device: h100 - working_dir: "/vllm-workspace/tests" - num_devices: 4 - source_file_dependencies: - - vllm/distributed/ - - "!vllm/distributed/kv_transfer/" - - vllm/engine/ - - vllm/v1/executor/ - - vllm/compilation/ - - tests/distributed/ - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - pytest -v -s distributed/test_elastic_ep.py +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (H100) [Nightly] Elastic EP Scaling" +# key: elastic-ep-scaling-test +# # Nightly only: 4-GPU H100 pods wait behind smaller jobs on per-PR runs. +# optional: true +# timeout_in_minutes: 40 +# device: h100 +# working_dir: "/vllm-workspace/tests" +# num_devices: 4 +# source_file_dependencies: +# - vllm/distributed/ +# - "!vllm/distributed/kv_transfer/" +# - vllm/engine/ +# - vllm/v1/executor/ +# - vllm/compilation/ +# - tests/distributed/ +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - pytest -v -s distributed/test_elastic_ep.py diff --git a/.buildkite/test_areas/fault_tolerance.yaml b/.buildkite/test_areas/fault_tolerance.yaml index 042a296edfbd..972b405822cd 100644 --- a/.buildkite/test_areas/fault_tolerance.yaml +++ b/.buildkite/test_areas/fault_tolerance.yaml @@ -2,13 +2,86 @@ group: Fault Tolerance depends_on: - image-build steps: -- label: ":nvidia: (H100) Fault Tolerance E2E" - key: fault-tolerance-e2e-2xh100 - timeout_in_minutes: 35 - device: h100 +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (H100) Fault Tolerance E2E" +# key: fault-tolerance-e2e-4xh100 +# timeout_in_minutes: 35 +# device: h100 +# num_devices: 4 +# 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/distributed/kv_transfer/ +# - vllm/v1/engine/ +# - vllm/v1/executor/ +# - 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_engine_core_sentinel.py +# - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py +# mirror: +# amd: +# label: ":amd: (MI355) MoRI EP Healthy Baseline" +# dind: false +# device: mi355_2 +# num_devices: 2 +# timeout_in_minutes: 10 +# working_dir: "/vllm-workspace/tests" +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - csrc/moe/ +# - csrc/rocm/ +# - tests/kernels/moe/ +# - vllm/model_executor/layers/fused_moe/ +# - vllm/model_executor/layers/quantization/ +# - vllm/distributed/ +# - vllm/config/ +# - vllm/forward_context.py +# - vllm/v1/worker/workspace.py +# - vllm/utils/import_utils.py +# - vllm/utils/math_utils.py +# - vllm/utils/torch_utils.py +# - vllm/platforms/ +# - vllm/_aiter_ops.py +# # TODO: Add fault injection, detection and recovery once MoRI supports them. +# # This mirror currently checks healthy EP execution only. +# commands: +# - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph +# - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]' + +# Keep the AMD job: its commands do not run this installation. +- label: ':amd: (MI355) MoRI EP Healthy Baseline' + key: amd-fault-tolerance-e2e-2xh100 + automatic_shard: true + timeout_in_minutes: 10 + device: mi355_2 num_devices: 2 - working_dir: "/vllm-workspace/tests" + working_dir: /vllm-workspace/tests source_file_dependencies: + - csrc/moe/ + - csrc/rocm/ + - tests/kernels/moe/ + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/ + - vllm/config/ + - vllm/forward_context.py + - vllm/v1/worker/workspace.py + - vllm/utils/import_utils.py + - vllm/utils/math_utils.py + - vllm/utils/torch_utils.py + - vllm/platforms/ + - vllm/_aiter_ops.py - vllm/v1/fault_tolerance/ - vllm/v1/worker/sentinel/ - vllm/entrypoints/serve/fault_tolerance/ @@ -21,39 +94,9 @@ steps: - 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_engine_core_sentinel.py - - pytest -v -s v1/fault_tolerance/test_fault_tolerance_e2e.py - mirror: - amd: - label: ":amd: (MI355) MoRI EP Healthy Baseline" - dind: false - device: mi355_2 - num_devices: 2 - timeout_in_minutes: 10 - working_dir: "/vllm-workspace/tests" - depends_on: - - image-build-amd - source_file_dependencies: - - csrc/moe/ - - csrc/rocm/ - - tests/kernels/moe/ - - vllm/model_executor/layers/fused_moe/ - - vllm/model_executor/layers/quantization/ - - vllm/distributed/ - - vllm/config/ - - vllm/forward_context.py - - vllm/v1/worker/workspace.py - - vllm/utils/import_utils.py - - vllm/utils/math_utils.py - - vllm/utils/torch_utils.py - - vllm/platforms/ - - vllm/_aiter_ops.py - # TODO: Add fault injection, detection and recovery once MoRI supports them. - # This mirror currently checks healthy EP execution only. - commands: - - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph - - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]' + - pytest -v -s v1/fault_tolerance/test_expert_reload.py + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs kernels/moe/test_moe_layer.py::test_moe_layer_mori_graph + - VLLM_TEST_ENABLE_MORI_MOE_LAYER=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -rs 'kernels/moe/test_moe_layer.py::test_moe_layer[False-mori_high_throughput-2-1-True]' --subtests='[32-512-512-8-2-bfloat16-None-False-False-False-False-mori_high_throughput-2-2-1]' + dind: false + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/jit_monitor.yaml b/.buildkite/test_areas/jit_monitor.yaml index 9e40c9d1affc..2e606a988ca0 100644 --- a/.buildkite/test_areas/jit_monitor.yaml +++ b/.buildkite/test_areas/jit_monitor.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 35GB) No Runtime JITs E2E" key: jit-monitor-no-runtime-jit + automatic_shard: true device: h200_35gb timeout_in_minutes: 45 source_file_dependencies: @@ -22,6 +23,7 @@ steps: - pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread mirror: amd: + automatic_shard: true # Focused MRV2 regression; the broader model/backend matrix is pending. label: ':amd: (MI355 DPX) MRV2 Sampler JIT Warmup' dind: false @@ -41,3 +43,4 @@ steps: - vllm/envs.py commands: - pytest -v -s tests/jit_monitor/test_no_runtime_jit_rocm.py --timeout=300 + - pytest -v -s tests/jit_monitor/test_flydsl_contract.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 44811a4085c1..c9c62aafd63d 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) vLLM IR" key: vllm-ir-tests + automatic_shard: true timeout_in_minutes: 35 device: h200_18gb working_dir: "/vllm-workspace/" @@ -15,6 +16,7 @@ steps: - pytest -v -s tests/kernels/ir mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) vLLM IR" dind: false device: mi355_dpx @@ -38,23 +40,31 @@ steps: source_file_dependencies: - csrc/ - tests/kernels/core + - vllm/kernels/triton/ - tests/kernels/test_concat_mla_q.py - tests/kernels/test_fused_qk_norm_rope_gate.py + - vllm/model_executor/layers/fused_qk_norm_rope.py commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_fused_qk_norm_rope_gate.py + - >- + pytest -v -s + kernels/core + --ignore=kernels/core/test_minimax_reduce_rms.py + kernels/test_concat_mla_q.py + kernels/test_fused_qk_norm_rope_gate.py mirror: amd: label: ":amd: (MI300) Core Operation Kernels" dind: false automatic_shard: true device: mi300_1 - timeout_in_minutes: 45 + timeout_in_minutes: 90 working_dir: "/vllm-workspace/tests" depends_on: - image-build-amd source_file_dependencies: - csrc/ - tests/kernels/core + - vllm/kernels/triton/ - tests/kernels/test_top_k_per_row.py - tests/kernels/test_concat_mla_q.py - tests/kernels/test_rocm_fp32_router_gemm.py @@ -64,11 +74,20 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py - tests/kernels/test_fused_qk_norm_rope_gate.py + - vllm/model_executor/layers/fused_qk_norm_rope.py commands: - - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_fused_qk_norm_rope_gate.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py + - >- + pytest -v -s + kernels/core + --ignore=kernels/core/test_minimax_reduce_rms.py + kernels/test_fused_qk_norm_rope_gate.py + kernels/test_concat_mla_q.py + kernels/test_top_k_per_row.py + kernels/test_rocm_fp32_router_gemm.py - label: ":nvidia: (H100) MiniMax Reduce RMS Kernels" key: kernels-minimax-reduce-rms-test-2-gpus + automatic_shard: true timeout_in_minutes: 20 num_devices: 2 device: h100 @@ -84,6 +103,7 @@ steps: - pytest -v -s kernels/core/test_minimax_reduce_rms.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) MiniMax Reduce RMS Kernels" dind: false device: mi355_2 @@ -102,6 +122,7 @@ steps: - label: ":nvidia: (H100) DeepSeek V4 Kernels" key: deepseek-v4-kernel-test-h100 + automatic_shard: true timeout_in_minutes: 30 device: h100 source_file_dependencies: @@ -110,12 +131,13 @@ steps: - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations commands: - - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - pytest -v -s kernels/test_top_k_per_row.py mirror: amd: label: ":amd: (MI355 DPX) DeepSeek V4 Fused Kernels" dind: false + automatic_shard: true device: mi355_dpx num_devices: 1 timeout_in_minutes: 60 @@ -130,8 +152,21 @@ steps: commands: - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +- label: ":nvidia: (H100) BF16x3 Router GEMM Kernels" + key: kernels-bf16x3-router-gemm-h100 + automatic_shard: true + timeout_in_minutes: 20 + device: h100 + source_file_dependencies: + - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_sm90.py + - vllm/model_executor/layers/fused_moe/router/gate_linear.py + - tests/kernels/test_bf16x3_router_gemm_cutedsl.py + commands: + - pytest -v -s kernels/test_bf16x3_router_gemm_cutedsl.py + - label: ":nvidia: (B200) DeepSeek V4 Kernels" key: deepseek-v4-kernel-test-b200 + automatic_shard: true timeout_in_minutes: 20 device: b200-k8s source_file_dependencies: @@ -144,12 +179,18 @@ steps: - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - tests/kernels/moe/test_deep_gemm_mega_moe.py - tests/models/test_deepseek_v4_mega_moe.py + - vllm/model_executor/kernels/attention/dsa/ + - vllm/model_executor/layers/sparse_mqa_indexer.py + - tests/model_executor/layers/test_mla_short_prefill_indexer.py commands: - - pytest -v -s kernels/test_fused_deepseek_v4_*.py + - pytest -v -s kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py - pytest -v -s kernels/moe/test_deep_gemm_mega_moe.py - pytest -v -s models/test_deepseek_v4_mega_moe.py + # DeepSeek V4.1 sparse MQA logits (MXFP4 and FP8 indexer cache); SM100 only. + - pytest -v -s model_executor/layers/test_mla_short_prefill_indexer.py -k sparse_mqa_logits mirror: amd: + automatic_shard: true label: ":amd: (MI355) DeepSeek V4 MoE Integration" dind: false device: mi355_1 @@ -182,6 +223,7 @@ steps: # ptxas override. - label: ":nvidia: (B200) mHC Kernels" key: kernels-mhc-test-b200 + automatic_shard: true timeout_in_minutes: 25 device: b200-k8s source_file_dependencies: @@ -193,6 +235,7 @@ steps: - pytest -v -s kernels/mhc mirror: amd: + automatic_shard: true label: ":amd: (MI355) mHC Kernels" dind: false device: mi355_1 @@ -208,6 +251,7 @@ steps: - label: ":nvidia: (B200) FLA Ops Kernels" key: kernels-fla-ops-test-b200 + automatic_shard: true timeout_in_minutes: 25 device: b200-k8s source_file_dependencies: @@ -219,6 +263,7 @@ steps: - pytest -v -s kernels/fla mirror: amd: + automatic_shard: true label: ":amd: (MI355) FLA Ops Kernels" dind: false device: mi355_1 @@ -232,11 +277,38 @@ steps: - tests/kernels/ - vllm/platforms/rocm.py +- label: ":nvidia: (B200) MiniMax M3 MegaMoE" + key: minimax-m3-mega-moe-test-b200 + automatic_shard: true + timeout_in_minutes: 20 + device: b200-k8s + source_file_dependencies: + - cmake/external_projects/deepgemm.cmake + - vllm/models/minimax_m3/nvidia/ + - vllm/utils/deep_gemm.py + - tests/models/test_minimax_m3_mega_moe.py + commands: + - pytest -v -s models/test_minimax_m3_mega_moe.py + +- label: ":nvidia: (B200) MiniMax-M3 Model Tests" + key: minimax-m3-model-tests-b200 + automatic_shard: true + timeout_in_minutes: 20 + device: b200-k8s + source_file_dependencies: + - vllm/models/minimax_m3/ + - vllm/cute_utils/ + - vllm/utils/flashinfer.py + - tests/models/minimax_m3/ + commands: + - pytest -v -s models/minimax_m3 --ignore=models/minimax_m3/amd_specialized_decode + # 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_bf16x3_router_gemm_cutedsl runs in Kernels (B200); test_ll_bf16_gemm and -# test_ll_fp32w_gemm run in Gate Linear CuTeDSL Kernels (B200)). +# (test_bf16x3_router_gemm_cutedsl runs in Kernels (B200) and BF16x3 Router GEMM +# Kernels (H100); test_ll_bf16_gemm and test_ll_fp32w_gemm run in Gate Linear +# CuTeDSL Kernels (B200)). - label: ":nvidia: (B200) Miscellaneous Kernels" key: kernels-root-misc-test-b200 automatic_shard: true @@ -316,8 +388,9 @@ steps: --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated - pytest -v -s kernels/test_fused_muse_spark_qknorm_rope_kv_insert.py -k rocm -- label: ":nvidia: (L4) Attention Kernels Shard %N" +- label: ":nvidia: (L4) Attention Kernels" key: kernels-attention-test + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 30 @@ -330,15 +403,14 @@ steps: - vllm/utils/flashinfer.py - tests/kernels/attention commands: - - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 7 + - pytest -v -s kernels/attention mirror: amd: - label: ":amd: (MI300) Attention Kernels Shard %N" + label: ":amd: (MI300) Attention Kernels" + automatic_shard: true dind: false device: mi300_1 timeout_in_minutes: 90 - parallelism: 6 depends_on: - image-build-amd source_file_dependencies: @@ -351,10 +423,11 @@ steps: - vllm/envs.py - vllm/platforms/rocm.py commands: - - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/attention --ignore=kernels/attention/test_triton_unified_attention_diffkv.py --ignore=kernels/attention/test_rocm_aiter_mla_decode.py --ignore=kernels/attention/test_rocm_aiter_mla_decode_metadata.py --ignore=kernels/attention/test_rocm_aiter_mla_fp8_support.py --ignore=kernels/attention/test_rocm_aiter_mla_op_registration.py - label: ":nvidia: (H100) Attention DiffKV Kernels" key: kernels-attention-diffkv-test-h100 + automatic_shard: true timeout_in_minutes: 20 device: h100 num_devices: 1 @@ -367,6 +440,7 @@ steps: - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Attention DiffKV Kernels" dind: false device: mi355_dpx @@ -386,6 +460,7 @@ steps: - label: ":nvidia: (H100) MLA Kernel Test" key: kernels-flashmla-test-h100 + automatic_shard: true timeout_in_minutes: 25 device: h100 num_devices: 1 @@ -407,6 +482,7 @@ steps: - pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) MLA Kernels" dind: false device: mi355_dpx @@ -475,8 +551,9 @@ steps: commands: - pytest -v -s kernels/quantization -k "not test_nvfp4_dense_emulation" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT -- label: ":nvidia: (L4) MoE Kernels Shard %N" +- label: ":nvidia: (L4) MoE Kernels" key: kernels-moe-test + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 50 @@ -485,16 +562,18 @@ steps: - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ + - vllm/models/muse_spark/common/ + - vllm/models/muse_spark/nvidia/ - vllm/distributed/device_communicators/ - vllm/envs.py - vllm/config commands: - - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 5 + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py mirror: amd: - label: ":amd: (MI300) MoE Kernels Shard %N" + label: ":amd: (MI300) MoE Kernels" + automatic_shard: true dind: false device: mi300_1 timeout_in_minutes: 55 @@ -503,6 +582,8 @@ steps: - csrc/moe/ - tests/kernels/moe - vllm/model_executor/layers/fused_moe/ + - vllm/models/muse_spark/common/ + - vllm/models/muse_spark/nvidia/ - vllm/distributed/device_communicators/ - vllm/envs.py - vllm/config @@ -525,6 +606,7 @@ steps: - tests/kernels/mamba - vllm/model_executor/layers/mamba/checkpoint.py - vllm/model_executor/layers/mamba/gdn + - vllm/models/qwen3_5/ - vllm/model_executor/layers/mamba/ops - vllm/v1/attention/backends/gdn_attn.py commands: @@ -544,6 +626,7 @@ steps: - label: ":nvidia: (H100) DeepGEMM Kernels" key: kernels-deepgemm-test-h100 + automatic_shard: true timeout_in_minutes: 35 device: h100 num_devices: 1 @@ -570,6 +653,7 @@ steps: - label: ":nvidia: (B200) MiniMax-M3 Attention Kernels" key: kernels-minimax-m3-attention-b200 + automatic_shard: true timeout_in_minutes: 30 working_dir: "/vllm-workspace/" device: b200-k8s @@ -582,11 +666,42 @@ steps: - tests/kernels/attention/test_minimax_m3.py commands: - pytest -v -s tests/kernels/attention/test_minimax_m3.py + mirror: + amd: + automatic_shard: true + label: ":amd: (MI355 DPX) MiniMax-M3 Attention Kernels" + dind: false + device: mi355_dpx + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + +- label: ":nvidia: (H100) MiniMax-M3 Sparse Prefill Kernels" + key: kernels-minimax-m3-sparse-prefill-h100 + automatic_shard: true + timeout_in_minutes: 15 + working_dir: "/vllm-workspace/" + device: h100 + num_devices: 1 + source_file_dependencies: + - vllm/models/minimax_m3/common/ops/sparse_attn.py + - vllm/models/minimax_m3/common/sparse_attention.py + - vllm/models/minimax_m3/nvidia/ops/sparse_prefill.py + - tests/v1/attention/utils.py + - tests/kernels/attention/test_minimax_m3.py + - tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py + commands: + - pytest -v -s tests/kernels/attention/test_minimax_m3.py -k prefill_sparse_attention_correctness + - pytest -v -s tests/kernels/test_minimax_m3_sparse_attn_fp8_scale.py -- label: ":nvidia: (B200) Kernels Shard %N" +- label: ":nvidia: (B200) Kernels" key: kernels-b200 + automatic_shard: true timeout_in_minutes: 30 - parallelism: 3 working_dir: "/vllm-workspace/" device: b200-k8s # optional: true @@ -611,7 +726,7 @@ steps: - vllm/platforms/cuda.py - vllm/cute_utils/ - vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/ - - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py + - vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_sm100.py - cmake/external_projects/fmha_sm100.cmake - vllm/models/minimax_m3/common/sparse_attention.py - vllm/models/minimax_m3/nvidia/ @@ -629,57 +744,47 @@ steps: - tests/kernels/mamba/test_gdn_prefill_cutedsl.py - tests/kernels/test_bf16x3_router_gemm_cutedsl.py - tests/kernels/test_top_k_per_row.py + - csrc/fused_mqa_topk/ + - tests/kernels/attention/test_fused_mqa_topk.py - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py - tests/kernels/quantization/test_block_fp8.py - tests/models/quantization/test_nvfp4.py commands: - nvidia-smi - - | - set -eu - case "$$BUILDKITE_PARALLEL_JOB" in - 0) - pytest -v -s tests/kernels/attention/test_attention_selector.py - pytest -v -s tests/kernels/attention/test_cache.py -k nvfp4 - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py - pytest -v -s tests/kernels/test_top_k_per_row.py - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py - ;; - 1) - python3 examples/basic/offline_inference/chat.py - pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py - pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py - pytest -v -s tests/models/quantization/test_nvfp4.py - ;; - 2) - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2' - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8 or cutlass_subset' - pytest -v -s tests/kernels/quantization/test_block_fp8.py::test_w8a8_block_fp8_cutlass_matmul - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py - pytest -v -s tests/kernels/moe/test_flashinfer.py - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py - ;; - *) - echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" >&2 - exit 1 - ;; - esac + - python3 examples/basic/offline_inference/chat.py + - pytest -v -s tests/kernels/attention/test_attention_selector.py + - pytest -v -s tests/kernels/attention/test_cache.py -k nvfp4 + - pytest -v -s tests/kernels/attention/test_flashinfer_trtllm_attention.py + - pytest -v -s tests/kernels/attention/test_flashinfer_mla_decode.py + - pytest -v -s tests/kernels/test_top_k_per_row.py + - pytest -v -s tests/kernels/attention/test_fused_mqa_topk.py + - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py + - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py + - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py + - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py + - pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py + - pytest -v -s tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py + - pytest -v -s tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py + - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py + - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py + - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py + - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py + - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py + - pytest -v -s tests/models/quantization/test_nvfp4.py + - pytest -v -s tests/kernels/attention/test_flashinfer.py -k 'not num_heads2' + - pytest -v -s tests/kernels/attention/test_cutlass_mla_decode.py + - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8 or cutlass_subset' + - pytest -v -s tests/kernels/quantization/test_block_fp8.py::test_w8a8_block_fp8_cutlass_matmul + - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py + - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py + - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py + - pytest -v -s tests/kernels/moe/test_flashinfer.py + - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py + - pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py mirror: amd: - label: ":amd: (MI355) NVFP4 Dense Emulation Shard %N" + automatic_shard: true + label: ":amd: (MI355) NVFP4 Dense Emulation" dind: false device: mi355_1 num_devices: 1 @@ -702,10 +807,11 @@ steps: # activation/KV-cache layouts and CUDA-only kernels need AMD backends. commands: - amd-smi - - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s kernels/quantization/test_nvfp4_emulation.py -k test_nvfp4_dense_emulation - label: ":nvidia: (B200) Gate Linear CuTeDSL Kernels" key: kernels-gate-linear-cutedsl-b200 + automatic_shard: true timeout_in_minutes: 15 working_dir: "/vllm-workspace/" device: b200-k8s @@ -727,6 +833,7 @@ steps: - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) Gate Linear Router Kernels" dind: false device: mi355_1 @@ -742,75 +849,82 @@ steps: - tests/kernels/test_gate_linear_rocm_dispatch.py - tests/kernels/test_bf16x3_router_gemm_rocm.py commands: - - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py + - >- + pytest -v -s + tests/kernels/test_rocm_fp32_router_gemm.py + tests/kernels/test_gate_linear_rocm_dispatch.py + tests/kernels/test_bf16x3_router_gemm_rocm.py -- label: ":nvidia: (DGX Spark) [Nightly] b12x Linear Kernels" - key: b12x-linear-kernels-dgx-spark-nightly - timeout_in_minutes: 60 - device: dgx-spark - optional: true - num_devices: 1 - depends_on: - - arm64-image-build - source_file_dependencies: - - setup.py - - vllm/config/kernel.py - - vllm/envs.py - - vllm/model_executor/kernels/linear/ - - vllm/model_executor/layers/fused_moe/b12x.py - - vllm/model_executor/layers/fused_moe/config.py - - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py - - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py - - vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/ - - vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py - - vllm/model_executor/layers/quantization/modelopt.py - - vllm/model_executor/layers/quantization/mxfp4.py - - vllm/model_executor/layers/quantization/utils/b12x_moe.py - - vllm/model_executor/warmup/b12x_warmup.py - - vllm/utils/b12x.py - - tests/kernels/moe/test_b12x.py - - tests/kernels/quantization/nvfp4_utils.py - - tests/model_executor/kernels/test_b12x_linear.py - - tests/model_executor/test_b12x_warmup.py - - tests/kernels/quantization/test_block_fp8.py - - tests/quantization/reference_mxfp4.py - - tests/quantization/test_auto_round.py - commands: - - uv pip install --system b12x==1.5.1 - - pytest -v -s model_executor/kernels/test_b12x_linear.py - model_executor/test_b12x_warmup.py - kernels/quantization/test_block_fp8.py - kernels/moe/test_b12x.py - quantization/test_auto_round.py -k b12x +# Disabled until b12x transitive dependencies are locked. +# - label: ":nvidia: (DGX Spark) [Nightly] b12x Linear Kernels" +# key: b12x-linear-kernels-dgx-spark-nightly +# timeout_in_minutes: 60 +# device: dgx-spark +# optional: true +# num_devices: 1 +# depends_on: +# - arm64-image-build +# source_file_dependencies: +# - setup.py +# - vllm/config/kernel.py +# - vllm/envs.py +# - vllm/model_executor/kernels/linear/ +# - vllm/model_executor/layers/fused_moe/b12x.py +# - vllm/model_executor/layers/fused_moe/config.py +# - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +# - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +# - vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/ +# - vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py +# - vllm/model_executor/layers/quantization/modelopt.py +# - vllm/model_executor/layers/quantization/mxfp4.py +# - vllm/model_executor/layers/quantization/utils/b12x_moe.py +# - vllm/model_executor/warmup/b12x_warmup.py +# - vllm/utils/b12x.py +# - tests/kernels/moe/test_b12x.py +# - tests/kernels/quantization/nvfp4_utils.py +# - tests/model_executor/kernels/test_b12x_linear.py +# - tests/model_executor/test_b12x_warmup.py +# - tests/kernels/quantization/test_block_fp8.py +# - tests/quantization/reference_mxfp4.py +# - tests/quantization/test_auto_round.py +# commands: +# - uv pip install --system b12x==1.5.1 +# - pytest -v -s model_executor/kernels/test_b12x_linear.py +# model_executor/test_b12x_warmup.py +# kernels/quantization/test_block_fp8.py +# kernels/moe/test_b12x.py +# quantization/test_auto_round.py -k b12x -- label: ":nvidia: (H100) Helion Kernels Shard %N" - key: kernels-helion-test - timeout_in_minutes: 115 - device: h100 - source_file_dependencies: - - vllm/utils/import_utils.py - - tests/kernels/helion/ - commands: - - pip install helion==1.4.0 - - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - parallelism: 5 - mirror: - amd: - label: ":amd: (MI250) Helion Kernels Shard %N" - dind: false - device: mi250_1 - timeout_in_minutes: 45 - # Explicit so the base step's count does not multiply the mirror's replicas. - parallelism: 2 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/platforms/rocm.py - - tools/benchmark_helion_kernels.py +# Disabled until helion transitive dependencies are locked. +# - label: ":nvidia: (H100) Helion Kernels Shard %N" +# key: kernels-helion-test +# timeout_in_minutes: 115 +# device: h100 +# source_file_dependencies: +# - vllm/utils/import_utils.py +# - tests/kernels/helion/ +# commands: +# - pip install helion==1.4.0 +# - pytest -v -s kernels/helion/ --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT +# parallelism: 5 +# mirror: +# amd: +# label: ":amd: (MI250) Helion Kernels Shard %N" +# dind: false +# device: mi250_1 +# timeout_in_minutes: 45 +# # Explicit so the base step's count does not multiply the mirror's replicas. +# parallelism: 2 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/platforms/rocm.py +# - tools/benchmark_helion_kernels.py - label: ":nvidia: (H100) [Nightly] FP8 MoE Kernels" key: kernels-fp8-moe-test-1xh100 + automatic_shard: true timeout_in_minutes: 40 device: h100 num_devices: 1 @@ -828,6 +942,7 @@ steps: - label: ":nvidia: (H100) [Nightly] DeepEP FP8 MoE Kernels" key: kernels-fp8-moe-test-2xh100 + automatic_shard: true timeout_in_minutes: 45 device: h100 num_devices: 2 @@ -838,6 +953,7 @@ steps: mirror: amd: label: ":amd: (MI355) DeepEP FP8 MoE Kernels" + automatic_shard: true dind: false device: mi355_2 timeout_in_minutes: 55 @@ -857,6 +973,7 @@ steps: - label: ":nvidia: (B200) [Nightly] FP4 MoE Kernels" key: kernels-fp4-moe-test-b200 + automatic_shard: true timeout_in_minutes: 25 device: b200-k8s num_devices: 1 @@ -868,6 +985,7 @@ steps: - pytest -v -s kernels/moe/test_ocp_mx_moe.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) FP4 and INT4 MoE Kernels" dind: false device: mi355_dpx @@ -899,6 +1017,7 @@ steps: - label: ":nvidia: (H100) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-h100s + automatic_shard: true timeout_in_minutes: 30 device: h100 num_devices: 2 @@ -916,6 +1035,7 @@ steps: amd: label: ":amd: (MI355) FusedMoE Layer Kernels" dind: false + automatic_shard: true device: mi355_2 timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" @@ -942,6 +1062,7 @@ steps: - label: ":nvidia: (B200) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-b200s + automatic_shard: true timeout_in_minutes: 45 device: b200-k8s num_devices: 2 @@ -958,6 +1079,7 @@ steps: - pytest -v -s kernels/moe/test_deepep_v2_moe.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) MoRI EP Numerics" dind: false device: mi355_2 @@ -991,6 +1113,7 @@ steps: - label: ":amd: (MI355 DPX) Native Quantization Kernels" key: amd-native-quantization-kernels-mi355 + automatic_shard: true dind: false device: mi355_dpx num_devices: 1 @@ -1019,6 +1142,7 @@ steps: - label: ":amd: (MI355 DPX) FP8 MoE Kernels" key: amd-fp8-moe-kernels-mi355 + automatic_shard: true dind: false device: mi355_dpx num_devices: 1 @@ -1049,6 +1173,7 @@ steps: - label: ":amd: (MI355 DPX) Kernels" key: amd-kernels-mi355 + automatic_shard: true dind: false device: mi355_dpx num_devices: 1 @@ -1099,4 +1224,12 @@ steps: - tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py - tests/kernels/models/amd/test_dsv41_mono_gfx942.py commands: - - pytest -v -s tests/kernels/models/amd/test_dsv41_gfx942_topk.py tests/kernels/models/amd/test_dsv41_gfx942_cand_logits.py tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py tests/kernels/models/amd/test_dsv41_mono_gfx942.py + - >- + pytest -v -s + tests/kernels/models/amd/test_dsv41_gfx942_topk.py + tests/kernels/models/amd/test_dsv41_gfx942_cand_logits.py + tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py + tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py + tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py + tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py + tests/kernels/models/amd/test_dsv41_mono_gfx942.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 84a4d98a6ec9..6ecfd895cc4a 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -15,6 +15,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.yaml mirror: amd: + automatic_shard: true label: ":amd: (MI300) LM Eval Small Models" dind: false device: mi300_1 @@ -33,6 +34,7 @@ steps: - label: ":amd: (MI300) LM Eval Small Models Harness" key: amd-lm-eval-small-models-harness + automatic_shard: true dind: false device: mi300_1 timeout_in_minutes: 35 @@ -55,6 +57,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) LM Eval Watermarking" device: h200_35gb key: lm-eval-watermarking + automatic_shard: true # Recorded runs took 80–140 seconds timeout_in_minutes: 10 source_file_dependencies: @@ -69,6 +72,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) LM Eval Watermarking" dind: false device: mi355_dpx @@ -90,19 +94,6 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark.txt -# - label: LM Eval Large Models (4xA100) -# key: lm-eval-large-models-4xa100 -# device: a100 -# optional: true -# num_devices: 4 -# working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" -# source_file_dependencies: -# - csrc/ -# - vllm/model_executor/layers/quantization -# commands: -# - export VLLM_WORKER_MULTIPROC_METHOD=spawn -# - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 - - label: ":nvidia: (H100) [Nightly] LM Eval Large Models" key: lm-eval-large-models-4xh100 automatic_shard: true @@ -178,6 +169,7 @@ steps: - label: ":nvidia: (B200) [Nightly] LM Eval Small Models Distributed" key: lm-eval-small-models-distributed-2xb200 + automatic_shard: true timeout_in_minutes: 120 device: b200-k8s num_devices: 2 @@ -191,6 +183,7 @@ steps: mirror: amd: label: ":amd: (MI355) LM Eval Small Models Distributed" + automatic_shard: true dind: false device: mi355_2 timeout_in_minutes: 145 @@ -265,6 +258,7 @@ steps: - label: ":nvidia: (H100) LM Eval Watermark Feature Combination" key: lm-eval-dspark-watermark-2xh100 + automatic_shard: true device: h100 num_devices: 2 # Local TP2 runs took 84–379 seconds, including startup. @@ -287,9 +281,14 @@ steps: autorun_on_main: true commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt - - pytest -s -v watermarking/test_compatibility.py watermarking/test_engine_compatibility.py watermarking/test_parallel_compatibility.py + - >- + pytest -s -v + watermarking/test_compatibility.py + watermarking/test_engine_compatibility.py + watermarking/test_parallel_compatibility.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) LM Eval Watermark Feature Combination" dind: false device: mi355_2 @@ -444,11 +443,13 @@ steps: - label: ":nvidia: (B200) [Nightly] LM Eval Qwen3.5 Models" key: lm-eval-qwen3-5-models-2xb200 + automatic_shard: true timeout_in_minutes: 45 device: b200-k8s optional: true num_devices: 2 source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen3_next.py @@ -459,12 +460,14 @@ steps: mirror: amd: label: ":amd: (MI355) LM Eval Qwen3.5 Models" + automatic_shard: true dind: false device: mi355_2 timeout_in_minutes: 75 depends_on: - image-build-amd source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py @@ -668,6 +671,7 @@ steps: - label: ":nvidia: (H200) [Nightly] LM Eval Nemotron 3 Super" key: lm-eval-nemotron-3-super-8xh200 + automatic_shard: true timeout_in_minutes: 35 device: h200 optional: true @@ -676,6 +680,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200-nemotron-3-super.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355) LM Eval Nemotron 3 Super" dind: false device: mi355_4 @@ -712,6 +717,7 @@ steps: - label: ":amd: (MI355) [Nightly] LM Eval Large Models" key: amd-lm-eval-large-models + automatic_shard: true dind: false device: mi355_8 optional: true @@ -727,18 +733,16 @@ steps: - label: ":nvidia: (H100) [Nightly] MoE Refactor Integration TEMPORARY" key: moe-refactor-integration-test-h100-temporary + automatic_shard: true device: h100 optional: true num_devices: 2 - parallelism: 4 commands: - # Timing-balanced static split of the former config-h100.txt (measured on - # build 83851): shards ~12.4/13.6/12.4/14.5m of eval time. New configs must - # be added to exactly one config-h100-shard-*.txt. - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt mirror: amd: - label: ":amd: (MI355) MoE Refactor Integration Shard %N" + automatic_shard: true + label: ":amd: (MI355) MoE Refactor Integration" dind: false device: mi355_2 num_devices: 2 @@ -755,23 +759,23 @@ steps: env: VLLM_WORKER_MULTIPROC_METHOD: "spawn" commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm-shard-$$BUILDKITE_PARALLEL_JOB.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-rocm.txt -- label: ":nvidia: (B200) [Nightly] MoE Refactor Integration TEMPORARY Shard %N" +- label: ":nvidia: (B200) [Nightly] MoE Refactor Integration TEMPORARY" key: moe-refactor-integration-test-b200-temporary + automatic_shard: true device: b200-k8s optional: true num_devices: 2 - parallelism: 4 commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200-shard-$$BUILDKITE_PARALLEL_JOB.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355) MoE Refactor BF16 Integration" dind: false device: mi355_2 num_devices: 2 - parallelism: 1 timeout_in_minutes: 105 working_dir: /vllm-workspace/tests depends_on: @@ -802,6 +806,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355) MoE Refactor BF16 DP Integration" dind: false device: mi355_2 @@ -829,6 +834,7 @@ steps: - label: ":nvidia: (L4) [Nightly] LM Eval Humming TEMPORARY" key: lm-eval-humming-l4 + automatic_shard: true optional: true timeout_in_minutes: 90 device: l4 @@ -846,45 +852,13 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-l4.txt -- label: ":nvidia: (A100) [Nightly] LM Eval Humming Default TEMPORARY Shard %N" - key: lm-eval-humming-a100-default - optional: true - timeout_in_minutes: 75 - device: a100 - num_devices: 1 - parallelism: 3 - source_file_dependencies: - - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming/ - - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/ - - vllm/model_executor/kernels/linear/ - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt - -- label: ":nvidia: (A100) [Nightly] LM Eval Humming Activation INT8 TEMPORARY Shard %N" - key: lm-eval-humming-a100-act - optional: true - timeout_in_minutes: 45 - device: a100 - num_devices: 1 - parallelism: 2 - source_file_dependencies: - - vllm/model_executor/layers/quantization/humming.py - - vllm/model_executor/layers/quantization/utils/humming/ - - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py - - vllm/model_executor/layers/fused_moe/oracle/ - - vllm/model_executor/kernels/linear/ - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-act-int8-shard-$$BUILDKITE_PARALLEL_JOB.txt - -- label: ":nvidia: (H100) [Nightly] LM Eval Humming Default TEMPORARY Shard %N" +- label: ":nvidia: (H100) [Nightly] LM Eval Humming Default TEMPORARY" key: lm-eval-humming-h100-default + automatic_shard: true optional: true timeout_in_minutes: 70 device: h100 num_devices: 1 - parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming/ @@ -892,10 +866,11 @@ steps: - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100.txt - label: ":nvidia: (H100) [Nightly] LM Eval Humming Activation FP8/INT8 TEMPORARY" key: lm-eval-humming-h100-act + automatic_shard: true optional: true timeout_in_minutes: 90 device: h100 @@ -912,6 +887,7 @@ steps: - label: ":nvidia: (H100) [Nightly] LM Eval Humming DeepEP TEMPORARY" key: lm-eval-humming-h100-deepep + automatic_shard: true optional: true timeout_in_minutes: 60 device: h100 @@ -932,6 +908,7 @@ steps: - label: ":nvidia: (B200) [Nightly] LM Eval Humming Default TEMPORARY" key: lm-eval-humming-b200-default + automatic_shard: true optional: true timeout_in_minutes: 50 device: b200-k8s @@ -963,6 +940,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant k8v4" key: lm-eval-turboquant-k8v4 + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -974,6 +952,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-k8v4.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) LM Eval TurboQuant k8v4" dind: false device: mi355_dpx @@ -1003,6 +982,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant t4nc" key: lm-eval-turboquant-t4nc + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -1016,6 +996,7 @@ steps: amd: label: ":amd: (MI355 DPX) LM Eval TurboQuant t4nc" dind: false + automatic_shard: true device: mi355_dpx num_devices: 1 timeout_in_minutes: 30 @@ -1043,6 +1024,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant k3v4nc" key: lm-eval-turboquant-k3v4nc + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -1056,6 +1038,7 @@ steps: amd: label: ":amd: (MI355 DPX) LM Eval TurboQuant k3v4nc" dind: false + automatic_shard: true device: mi355_dpx num_devices: 1 timeout_in_minutes: 30 @@ -1083,6 +1066,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) LM Eval TurboQuant t3nc" key: lm-eval-turboquant-t3nc + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -1096,6 +1080,7 @@ steps: amd: label: ":amd: (MI355 DPX) LM Eval TurboQuant t3nc" dind: false + automatic_shard: true device: mi355_dpx num_devices: 1 timeout_in_minutes: 30 @@ -1121,6 +1106,28 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant-t3nc-rocm.txt +# Octave kernels are built only for gfx942 and gfx1100, so this runs on MI300. +- label: ":amd: (MI300) LM Eval Octave KV Cache" + key: amd-lm-eval-octave-mi300 + automatic_shard: true + dind: false + device: mi300_1 + num_devices: 1 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - tests/evals/gsm8k/ + - tests/utils.py + - tests/conftest.py + - vllm/v1/attention/backends/rocm_octave_attn.py + - vllm/v1/attention/ops/rocm_octave.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-octave.txt + - label: ":nvidia: (H100) [Nightly] GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-2xh100 timeout_in_minutes: 35 @@ -1159,6 +1166,7 @@ steps: - label: ":nvidia: (B200) [Nightly] GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-2xb200 + automatic_shard: true timeout_in_minutes: 30 device: b200-k8s optional: true @@ -1171,6 +1179,7 @@ steps: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355) GPQA Eval (GPT-OSS)" dind: false device: mi355_2 @@ -1191,62 +1200,65 @@ steps: commands: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx950.txt -- label: ":nvidia: (B200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval" - key: accuracy-eval-qwen3-8-flash-next-fp8-b200 - timeout_in_minutes: 60 - device: b200-k8s - optional: true - num_devices: 4 - source_file_dependencies: - - csrc/gdn/ - - tests/evals/qwen4_exp/ - - vllm/model_executor/layers/mamba/ - - vllm/models/qwen4_exp/ - - vllm/v1/attention/backends/short_conv_attn.py - commands: - - uv pip install --system 'evalscope==1.10.0' - - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-b200.txt +# Disabled until evalscope transitive dependencies are locked. +# - label: ":nvidia: (B200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval" +# key: accuracy-eval-qwen3-8-flash-next-fp8-b200 +# timeout_in_minutes: 60 +# device: b200-k8s +# optional: true +# num_devices: 4 +# source_file_dependencies: +# - csrc/gdn/ +# - tests/evals/qwen4_exp/ +# - vllm/model_executor/layers/mamba/ +# - vllm/models/qwen4_exp/ +# - vllm/v1/attention/backends/short_conv_attn.py +# commands: +# - uv pip install --system 'evalscope==1.10.0' +# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-b200.txt -- label: ":nvidia: (H200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval" - key: accuracy-eval-qwen3-8-flash-next-fp8-h200 - timeout_in_minutes: 60 - device: h200 - optional: true - num_devices: 4 - source_file_dependencies: - - csrc/gdn/ - - tests/evals/qwen4_exp/ - - vllm/model_executor/layers/mamba/ - - vllm/models/qwen4_exp/ - - vllm/v1/attention/backends/short_conv_attn.py - commands: - - uv pip install --system 'evalscope==1.10.0' - - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt - mirror: - amd: - label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" - dind: false - device: mi355_4 - num_devices: 4 - timeout_in_minutes: 105 - working_dir: /vllm-workspace/tests - depends_on: - - image-build-amd - source_file_dependencies: - - csrc/gdn/ - - tests/evals/qwen4_exp/ - - vllm/model_executor/layers/mamba/ - - vllm/models/qwen4_exp/ - - vllm/v1/attention/backends/short_conv_attn.py - - vllm/platforms/rocm.py - env: - VLLM_WORKER_MULTIPROC_METHOD: "spawn" - commands: - - uv pip install --system 'evalscope==1.10.0' - - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt +# Disabled until evalscope transitive dependencies are locked. +# - label: ":nvidia: (H200) [Nightly] Qwen3.8-Flash-Next-FP8 Accuracy Eval" +# key: accuracy-eval-qwen3-8-flash-next-fp8-h200 +# timeout_in_minutes: 60 +# device: h200 +# optional: true +# num_devices: 4 +# source_file_dependencies: +# - csrc/gdn/ +# - tests/evals/qwen4_exp/ +# - vllm/model_executor/layers/mamba/ +# - vllm/models/qwen4_exp/ +# - vllm/v1/attention/backends/short_conv_attn.py +# commands: +# - uv pip install --system 'evalscope==1.10.0' +# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt +# mirror: +# amd: +# label: ":amd: (MI355) Qwen3.8-Flash-Next-FP8 Accuracy Eval" +# dind: false +# device: mi355_4 +# num_devices: 4 +# timeout_in_minutes: 105 +# working_dir: /vllm-workspace/tests +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - csrc/gdn/ +# - tests/evals/qwen4_exp/ +# - vllm/model_executor/layers/mamba/ +# - vllm/models/qwen4_exp/ +# - vllm/v1/attention/backends/short_conv_attn.py +# - vllm/platforms/rocm.py +# env: +# VLLM_WORKER_MULTIPROC_METHOD: "spawn" +# commands: +# - uv pip install --system 'evalscope==1.10.0' +# - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-rocm.txt - label: ":nvidia: (DGX Spark) [Nightly] GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-spark + automatic_shard: true timeout_in_minutes: 35 device: dgx-spark optional: true @@ -1261,6 +1273,7 @@ steps: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) GPQA Eval (GPT-OSS TP1)" # The Spark config is a backend-neutral GPT-OSS TP1 baseline. dind: false @@ -1284,6 +1297,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) LM Eval KV-Offload" key: kv-offload-small + automatic_shard: true timeout_in_minutes: 30 device: h200_35gb source_file_dependencies: @@ -1296,6 +1310,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) LM Eval KV-Offload" dind: false device: mi355_dpx @@ -1309,10 +1324,13 @@ steps: - vllm/v1/kv_offload/ - vllm/v1/simple_kv_offload/ - tests/evals/gsm8k/test_gsm8k_offloading.py + - vllm/distributed/kv_transfer/kv_connector/v1/umbp/ + - tests/v1/kv_connector/umbp_test_utils.py - vllm/platforms/rocm.py - label: ":nvidia: (H100) LM Eval KV-Offload Medium" key: kv-offload-medium + automatic_shard: true timeout_in_minutes: 45 device: h100 num_devices: 2 @@ -1326,6 +1344,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" mirror: amd: + automatic_shard: true label: ":amd: (MI355) LM Eval KV-Offload Medium" dind: false device: mi355_2 @@ -1343,6 +1362,7 @@ steps: - label: ":nvidia: (H100) [Nightly] LM Eval KV-Offload Large" key: kv-offload-large + automatic_shard: true # Nightly only: 4-GPU H100 pods wait behind smaller jobs on per-PR runs. optional: true timeout_in_minutes: 40 @@ -1358,6 +1378,7 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" mirror: amd: + automatic_shard: true label: ":amd: (MI355) LM Eval KV-Offload Large" dind: false device: mi355_4 @@ -1377,6 +1398,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) MRCR Eval Small Models" key: mrcr-eval-small-models + automatic_shard: true device: h200_35gb timeout_in_minutes: 25 source_file_dependencies: @@ -1385,6 +1407,7 @@ steps: - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) MRCR Eval Small Models" dind: false device: mi355_dpx @@ -1402,6 +1425,7 @@ steps: - label: ":amd: (MI355) [Nightly] LM Eval Large Models Harness" key: amd-lm-eval-large-models-harness + automatic_shard: true dind: false device: mi355_8 optional: true @@ -1427,6 +1451,7 @@ steps: - label: ":amd: (MI355) [Nightly] LM Eval Large Models ROCm Harness" key: amd-lm-eval-large-models-rocm-harness + automatic_shard: true dind: false device: mi355_8 optional: true diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index 568f8a1198e5..74770b4a4d89 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -14,7 +14,17 @@ steps: # Which tests share a shard varies, and forking an engine is unsafe once an # earlier test initialized CUDA in a way vLLM's fork check misses. - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s lora --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py + - >- + pytest -v -s + lora + --ignore=lora/test_chatglm3_tp.py + --ignore=lora/test_llama_tp.py + --ignore=lora/test_qwen3_with_multi_loras.py + --ignore=lora/test_olmoe_tp.py + --ignore=lora/test_deepseekv2_tp.py + --ignore=lora/test_gptoss_tp.py + --ignore=lora/test_qwen3moe_tp.py + --ignore=lora/test_qwen35_densemodel_lora.py mirror: amd: label: ":amd: (MI355 DPX) LoRA" @@ -31,17 +41,18 @@ steps: - image-build-amd -- label: ":nvidia: (L4) [Nightly] LoRA TP (Distributed) %N" +- label: ":nvidia: (L4) [Nightly] LoRA TP (Distributed)" key: lora-tp-distributed + automatic_shard: true # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. optional: true timeout_in_minutes: 25 device: l4 num_devices: 4 - parallelism: 4 source_file_dependencies: - vllm/lora - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/models/transformers/moe.py - tests/lora commands: # FIXIT: find out which code initialize cuda before running the test @@ -51,21 +62,20 @@ steps: - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True # There is some Tensor Parallelism related processing logic in LoRA that # requires multi-GPU testing for validation. - # Timing-balanced static split from per-file walls in builds 83921/83936. - - case "$$BUILDKITE_PARALLEL_JOB" in 0|1|2|3) ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" && exit 1;; esac - - test "$$BUILDKITE_PARALLEL_JOB" != "3" || pytest -v -s -x lora/test_chatglm3_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_llama_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "3" || pytest -v -s -x lora/test_qwen3_with_multi_loras.py - - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_olmoe_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_gptoss_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "1" || pytest -v -s -x lora/test_qwen35_densemodel_lora.py - - test "$$BUILDKITE_PARALLEL_JOB" != "2" || pytest -v -s -x lora/test_gemma4_tp.py - - test "$$BUILDKITE_PARALLEL_JOB" != "0" || pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp + - pytest -v -s -x lora/test_chatglm3_tp.py + - pytest -v -s -x lora/test_llama_tp.py + - pytest -v -s -x lora/test_qwen3_with_multi_loras.py + - pytest -v -s -x lora/test_olmoe_tp.py + - pytest -v -s -x lora/test_gptoss_tp.py + - pytest -v -s -x lora/test_qwen35_densemodel_lora.py + - pytest -v -s -x lora/test_gemma4_tp.py + - pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp mirror: amd: # Run on PRs even when the L4 step is optional (nightly only). optional: false - label: ":amd: (MI300) LoRA TP (Distributed) Shard %N" + automatic_shard: true + label: ":amd: (MI300) LoRA TP (Distributed)" dind: false device: mi300_4 working_dir: "/vllm-workspace/tests" @@ -75,17 +85,17 @@ steps: source_file_dependencies: - vllm/lora - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/models/transformers/moe.py - vllm/platforms/rocm.py - tests/lora commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn # expandable_segments breaks custom all-reduce IPC on ROCm. - - case "$$BUILDKITE_PARALLEL_JOB" in 0|1|2|3) ;; *) echo "unexpected BUILDKITE_PARALLEL_JOB=$$BUILDKITE_PARALLEL_JOB" && exit 1;; esac - - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_chatglm3_tp.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_llama_tp.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "3" ]; then pytest -v -s -x lora/test_qwen3_with_multi_loras.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_olmoe_tp.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_gptoss_tp.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s -x lora/test_qwen35_densemodel_lora.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s -x lora/test_gemma4_tp.py; fi - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp; fi + - pytest -v -s -x lora/test_chatglm3_tp.py + - pytest -v -s -x lora/test_llama_tp.py + - pytest -v -s -x lora/test_qwen3_with_multi_loras.py + - pytest -v -s -x lora/test_olmoe_tp.py + - pytest -v -s -x lora/test_gptoss_tp.py + - pytest -v -s -x lora/test_qwen35_densemodel_lora.py + - pytest -v -s -x lora/test_gemma4_tp.py + - pytest -v -s -x lora/test_sequence_classification.py::test_batched_loras_tp diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index a81db91830ef..3f162810a3d4 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -40,6 +40,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) V1 Logits + Oracle" key: v1-logits-oracle + automatic_shard: true timeout_in_minutes: 45 device: h200_18gb source_file_dependencies: *v1-sample-logits-deps @@ -55,6 +56,7 @@ steps: amd: label: ":amd: (MI355 DPX) V1 Logits + Oracle" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 70 depends_on: @@ -63,6 +65,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) V1 Core" device: h200_35gb key: v1-core + automatic_shard: true timeout_in_minutes: 45 source_file_dependencies: &v1-core-kv-metrics-deps - vllm/config/ @@ -98,6 +101,7 @@ steps: - pytest -v -s -m 'not cpu_test' v1/core mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) V1 Core" dind: false device: mi355_dpx @@ -108,6 +112,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) V1 Executor + Worker" device: h200_18gb key: v1-executor-worker + automatic_shard: true timeout_in_minutes: 45 source_file_dependencies: *v1-core-kv-metrics-deps commands: @@ -120,6 +125,7 @@ steps: amd: label: ":amd: (MI355 DPX) V1 Executor + Worker" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 60 depends_on: @@ -144,56 +150,59 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors" - device: h200_35gb - key: v1-kv-connectors - automatic_shard: true - timeout_in_minutes: 45 - source_file_dependencies: *v1-core-kv-metrics-deps - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit - mirror: - amd: - label: ":amd: (MI355 DPX) V1 KV Connectors" - automatic_shard: true - dind: false - device: mi355_dpx - timeout_in_minutes: 60 - depends_on: - - image-build-amd +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors" +# device: h200_35gb +# key: v1-kv-connectors +# automatic_shard: true +# timeout_in_minutes: 45 +# source_file_dependencies: *v1-core-kv-metrics-deps +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit +# - pytest -v -s -m 'not cpu_test' v1/ec_connector/unit +# mirror: +# amd: +# label: ":amd: (MI355 DPX) V1 KV Connectors" +# automatic_shard: true +# dind: false +# device: mi355_dpx +# timeout_in_minutes: 60 +# depends_on: +# - image-build-amd # The shared KV Connectors steps above exclude vllm/distributed/kv_transfer/ from # their triggers, and CUDA skips MoRIIO's mori/aiter-gated tests. Run the whole # MoRIIO suite on ROCm whenever the connector or its tests change. -- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" - key: amd-moriio-connector-unit-mi355 - dind: false - device: mi355_dpx - num_devices: 1 - timeout_in_minutes: 30 - working_dir: /vllm-workspace/tests - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ - - vllm/distributed/kv_transfer/kv_connector/v1/base.py - - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py - - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py - - tests/v1/kv_connector/unit/test_moriio_ - - tests/v1/kv_connector/unit/utils.py - - requirements/kv_connectors_rocm.txt - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" +# key: amd-moriio-connector-unit-mi355 +# dind: false +# device: mi355_dpx +# num_devices: 1 +# timeout_in_minutes: 30 +# working_dir: /vllm-workspace/tests +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ +# - vllm/distributed/kv_transfer/kv_connector/v1/base.py +# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +# - tests/v1/kv_connector/unit/test_moriio_ +# - tests/v1/kv_connector/unit/utils.py +# - requirements/kv_connectors_rocm.txt +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py # UMBP is built on mori, so its tests skip on CUDA. Run them on MI355 whenever # the connector or its tests change. - label: ":amd: (MI355 DPX) UMBP Connector Unit Tests" key: amd-umbp-connector-unit-mi355 + automatic_shard: true dind: false device: mi355_dpx num_devices: 1 @@ -207,64 +216,92 @@ steps: - vllm/v1/kv_cache_interface.py - tests/v1/kv_connector/unit/test_umbp_ - tests/v1/kv_connector/umbp_test_utils.py + - tests/v1/kv_connector/unit/utils.py + - tests/v1/kv_connector/conftest.py + - docker/Dockerfile.rocm_base commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/unit/test_umbp_*.py + - pytest -v -s -rs v1/kv_connector/unit/test_umbp_async_lookup.py v1/kv_connector/unit/test_umbp_core.py v1/kv_connector/unit/test_umbp_embedded.py v1/kv_connector/unit/test_umbp_events.py v1/kv_connector/unit/test_umbp_hybrid.py -# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate -# GPUs over xGMI, compared exactly against a direct serve. -- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" - key: amd-moriio-write-dummy-e2e-mi355 +- label: ":amd: (MI355 DPX) UMBP Connector Integration" + key: amd-umbp-connector-integration-mi355 dind: false - device: mi355_2 - num_devices: 2 - # ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the - # largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports - # for its servers and proxy and its own compile-cache dir. - parallelism: 6 - timeout_in_minutes: 45 + device: mi355_dpx + num_devices: 1 + timeout_in_minutes: 30 working_dir: /vllm-workspace/tests depends_on: - image-build-amd source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/umbp/ - vllm/distributed/kv_transfer/kv_connector/v1/base.py - - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py - - vllm/model_executor/model_loader/dummy_loader.py - - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py - - tests/v1/kv_connector/moriio_integration/ - - requirements/kv_connectors_rocm.txt + - tests/v1/kv_connector/umbp_integration/ + - tests/v1/kv_connector/umbp_test_utils.py + - tests/v1/kv_connector/conftest.py + - tests/v1/kv_connector/rocm_pd_accuracy_utils.py + - docker/Dockerfile.rocm_base commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - pytest -v -s -rs v1/kv_connector/umbp_integration -# The same P/D with prefill TP2 writing into decode TP1, which needs 3 GPUs. -- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP" - key: amd-moriio-write-dummy-hetero-tp-mi355 - dind: false - device: mi355_4 - num_devices: 4 - timeout_in_minutes: 45 - working_dir: /vllm-workspace/tests - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ - - vllm/distributed/kv_transfer/kv_connector/v1/base.py - - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py - - vllm/model_executor/model_loader/dummy_loader.py - - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py - - tests/v1/kv_connector/moriio_integration/ - - requirements/kv_connectors_rocm.txt - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k heterogeneous_tp +# Single-node MoRIIO P/D on dummy weights, WRITE and READ mode: prefill and +# decode on separate GPUs over xGMI, compared exactly against a direct serve. +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" +# key: amd-moriio-write-dummy-e2e-mi355 +# dind: false +# device: mi355_2 +# num_devices: 2 +# # ~3.5-5 min per case; pytest-shard hashes test ids, so with 28 cases the +# # largest of 6 shards holds 7 (~32 min on MI355). Every case picks free ports +# # for its servers and proxy and its own compile-cache dir. +# parallelism: 6 +# timeout_in_minutes: 45 +# working_dir: /vllm-workspace/tests +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ +# - vllm/distributed/kv_transfer/kv_connector/v1/base.py +# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +# - vllm/model_executor/model_loader/dummy_loader.py +# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +# - tests/v1/kv_connector/moriio_integration/ +# - requirements/kv_connectors_rocm.txt +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k "not heterogeneous_tp" --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + +# The same P/D between different TP sizes (WRITE TP2 -> TP1, READ of MLA +# TP1 <-> TP2), which needs 3 GPUs: 8 cases, ~36 min on MI355. +# Disabled until the KV connectors installation is locked. +# - label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight Heterogeneous TP" +# key: amd-moriio-write-dummy-hetero-tp-mi355 +# dind: false +# device: mi355_4 +# num_devices: 4 +# timeout_in_minutes: 45 +# working_dir: /vllm-workspace/tests +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ +# - vllm/distributed/kv_transfer/kv_connector/v1/base.py +# - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +# - vllm/model_executor/model_loader/dummy_loader.py +# - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +# - tests/v1/kv_connector/moriio_integration/ +# - requirements/kv_connectors_rocm.txt +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py -k heterogeneous_tp - label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval" device: h200_35gb key: v1-metrics-lmeval + automatic_shard: true timeout_in_minutes: 45 env: # GitHub intermittently rejects HTTP/2 upload-pack requests from the H200 @@ -282,6 +319,7 @@ steps: - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) V1 Metrics + LM Eval" dind: false device: mi355_dpx @@ -313,13 +351,13 @@ steps: - tests/v1 - tests/watermarking - examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py + - rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json device: cpu-small commands: # split the test to avoid interference - 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/worker/test_gpu_worker.py::test_cpu_worker_without_triton - pytest -v -s v1/worker/test_cpu_worker_kv_budget.py - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py @@ -330,6 +368,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Extract Hidden States Integration" key: extract-hidden-states-integration + automatic_shard: true timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace" @@ -344,6 +383,7 @@ steps: - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Extract Hidden States Integration Single GPU" dind: false device: mi355_dpx @@ -366,6 +406,7 @@ steps: - label: ":nvidia: (L4) Extract Hidden States Integration" key: extract-hidden-states-integration-2-gpus + automatic_shard: true timeout_in_minutes: 20 device: l4 num_devices: 2 @@ -381,6 +422,7 @@ steps: - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration mirror: amd: + automatic_shard: true label: ":amd: (MI355) Extract Hidden States Integration" dind: false device: mi355_2 @@ -395,7 +437,9 @@ steps: - vllm/model_executor/model_loader/ - vllm/model_executor/models/extract_hidden_states.py - vllm/model_executor/models/llama.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/registry.py - vllm/transformers_utils/configs/extract_hidden_states.py @@ -408,37 +452,38 @@ steps: - tests/v1/kv_connector/extract_hidden_states_integration/ - vllm/platforms/rocm.py -- label: ":nvidia: (H200 MIG 18GB) Regression" - key: regression - timeout_in_minutes: 30 - device: h200_18gb - source_file_dependencies: - - vllm/config/ - - vllm/distributed/ - - "!vllm/distributed/kv_transfer/" - - vllm/engine/ - - vllm/inputs/ - - vllm/model_executor/ - - vllm/multimodal/ - - vllm/platforms/ - - vllm/sampling_params.py - - vllm/transformers_utils/ - - vllm/utils/ - - vllm/v1/ - - tests/test_regression - commands: - - pip install 'modelscope<1.38' - - pytest -v -s test_regression.py - working_dir: "/vllm-workspace/tests" # optional - mirror: - amd: - label: ":amd: (MI355 DPX) Regression" - dind: false - device: mi355_dpx - timeout_in_minutes: 25 - depends_on: - - image-build-amd - working_dir: /vllm-workspace/tests +# Disabled until the modelscope installation is locked. +# - label: ":nvidia: (H200 MIG 18GB) Regression" +# key: regression +# timeout_in_minutes: 30 +# device: h200_18gb +# source_file_dependencies: +# - vllm/config/ +# - vllm/distributed/ +# - "!vllm/distributed/kv_transfer/" +# - vllm/engine/ +# - vllm/inputs/ +# - vllm/model_executor/ +# - vllm/multimodal/ +# - vllm/platforms/ +# - vllm/sampling_params.py +# - vllm/transformers_utils/ +# - vllm/utils/ +# - vllm/v1/ +# - tests/test_regression +# commands: +# - pip install 'modelscope<1.38' +# - pytest -v -s test_regression.py +# working_dir: "/vllm-workspace/tests" # optional +# mirror: +# amd: +# label: ":amd: (MI355 DPX) Regression" +# dind: false +# device: mi355_dpx +# timeout_in_minutes: 25 +# depends_on: +# - image-build-amd +# working_dir: /vllm-workspace/tests - label: ":nvidia: (H200 MIG 35GB) Examples" device: h200_35gb @@ -488,6 +533,7 @@ steps: - label: ":nvidia: (L4) Metrics, Tracing" key: metrics-tracing-2-gpus + automatic_shard: true timeout_in_minutes: 25 device: l4 num_devices: 2 @@ -508,15 +554,11 @@ steps: - tests/v1/tracing - tests/tracing/ commands: - - "pip install \ - 'opentelemetry-sdk>=1.26.0' \ - 'opentelemetry-api>=1.26.0' \ - 'opentelemetry-exporter-otlp>=1.26.0' \ - 'opentelemetry-semantic-conventions-ai>=0.4.1'" - pytest -v -s v1/tracing - pytest -v -s tracing mirror: amd: + automatic_shard: true label: ":amd: (MI300) Metrics, Tracing" dind: false device: mi300_2 @@ -562,6 +604,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Async Engine, Inputs, Utils, Worker" device: h200_35gb key: async-engine-inputs-utils-worker + automatic_shard: true timeout_in_minutes: 25 source_file_dependencies: - vllm/assets/ @@ -589,6 +632,7 @@ steps: amd: label: ":amd: (MI355 DPX) Async Engine, Inputs, Utils, Worker" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 50 depends_on: @@ -600,8 +644,129 @@ steps: - pytest -v -s multimodal/test_cache.py::test_sleep_wake_preserves_mm_cache_consistency - pytest -v -s utils_ -- label: ":computer: (CPU) Params, Env, Tokenizers, Parser" - key: cpu-params-env-tokenizers-parser +# Disabled until the cmake installation is locked. +# - label: ":computer: (CPU) Params, Env, Tokenizers, Parser" +# key: cpu-params-env-tokenizers-parser +# depends_on: +# - image-build-cpu +# timeout_in_minutes: 35 +# source_file_dependencies: &cpu-async-engine-deps +# - vllm/assets/ +# - vllm/config/ +# - vllm/device_allocator/ +# - vllm/connections.py +# - vllm/engine/arg_utils.py +# - vllm/renderers/chat_utils/ +# - vllm/entrypoints/mcp/ +# - vllm/entrypoints/openai/chat_completion/protocol.py +# - vllm/entrypoints/openai/completion/protocol.py +# - vllm/entrypoints/openai/responses/protocol.py +# - vllm/entrypoints/generate/base/protocol.py +# - vllm/env_override.py +# - vllm/envs.py +# - vllm/exceptions.py +# - vllm/inputs/ +# - vllm/logging_utils/ +# - vllm/logprobs.py +# - vllm/model_executor/layers/quantization/quark/ +# - vllm/model_executor/model_loader/tensorizer.py +# - vllm/models/muse_spark/ +# - vllm/multimodal/ +# - vllm/outputs.py +# - vllm/parser/ +# - vllm/platforms/ +# - vllm/pooling_params.py +# - vllm/ray/ +# - vllm/reasoning/ +# - vllm/renderers/ +# - vllm/sampling_params.py +# - vllm/scalar_type.py +# - vllm/tokenizers/ +# - vllm/tool_parsers/ +# - vllm/transformers_utils/ +# - vllm/triton_utils/ +# - vllm/utils/ +# - vllm/v1/ +# - vllm/version.py +# - cmake/utils.cmake +# - cmake/hipify.py +# - tests/test_uvicorn_logging.py +# - tests/test_attention_backend_registry.py +# - tests/test_audio_media_size_precheck.py +# - tests/test_cmake_utils.py +# - tests/test_connections.py +# - tests/test_embedded_commit.py +# - tests/test_envs.py +# - tests/test_force_first_config.py +# - tests/test_fxgraphcache_pickle_patch.py +# - tests/test_logprobs.py +# - tests/test_outputs.py +# - tests/test_pcp_dp.py +# - tests/test_pooling_params.py +# - tests/test_quantization_revision_pin.py +# - tests/test_ray_env.py +# - tests/test_ray_env_utils.py +# - tests/test_request_input_bounds.py +# - tests/test_sampling_params.py +# - tests/test_scalartype.py +# - tests/test_tensorizer_revision_pin.py +# - tests/test_triton_utils.py +# - tests/test_version.py +# - tests/test_zen_cpu_platform_detection.py +# - tests/multimodal +# - tests/renderers +# - tests/standalone_tests/lazy_imports.py +# - tests/standalone_tests/test_precompiled_wheel.py +# - setup.py +# - tests/reasoning +# - tools/pre_commit/check_test_tethering.py +# - tools/pre_commit/test_tethering_allowlist.txt +# - tests/tools/test_check_test_tethering.py +# - tools/pre_commit/check_agent_files.py +# - tests/tools/test_check_agent_files.py +# - tools/pre_commit/check_trust_remote_code.py +# - tests/tools/test_check_trust_remote_code.py +# - tests/tool_parsers +# - tests/tokenizers_ +# - tests/parser +# - tests/transformers_utils +# - tests/config +# - tests/device_allocator +# device: cpu-small +# commands: +# - pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none +# - python3 standalone_tests/lazy_imports.py +# - python3 standalone_tests/test_precompiled_wheel.py +# - pytest -v -s device_allocator +# - pytest -v -s test_uvicorn_logging.py +# - pytest -v -s test_attention_backend_registry.py +# - pytest -v -s test_audio_media_size_precheck.py +# - pytest -v -s test_cmake_utils.py +# - pytest -v -s test_connections.py +# - pytest -v -s test_embedded_commit.py +# - pytest -v -s test_envs.py +# - pytest -v -s test_force_first_config.py +# - pytest -v -s test_fxgraphcache_pickle_patch.py +# - pytest -v -s test_logprobs.py +# - pytest -v -s test_outputs.py +# - pytest -v -s test_pcp_dp.py +# - pytest -v -s test_pooling_params.py +# - pytest -v -s test_quantization_revision_pin.py +# - pytest -v -s test_ray_env.py +# - pytest -v -s test_ray_env_utils.py +# - pytest -v -s test_request_input_bounds.py +# - pytest -v -s test_sampling_params.py +# - pytest -v -s test_scalartype.py +# - pytest -v -s test_tensorizer_revision_pin.py +# - pytest -v -s test_triton_utils.py +# - pytest -v -s test_version.py +# - pytest -v -s test_zen_cpu_platform_detection.py +# - pytest -v -s tokenizers_ +# - pytest -v -s parser +# - pytest -v -s transformers_utils + +- label: ":computer: (CPU) Multimodal + Config" + key: cpu-multimodal-config depends_on: - image-build-cpu timeout_in_minutes: 35 @@ -644,6 +809,7 @@ steps: - vllm/v1/ - vllm/version.py - cmake/utils.cmake + - cmake/hipify.py - tests/test_uvicorn_logging.py - tests/test_attention_backend_registry.py - tests/test_audio_media_size_precheck.py @@ -688,45 +854,6 @@ steps: - tests/device_allocator device: cpu-small commands: - - pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none - - python3 standalone_tests/lazy_imports.py - - python3 standalone_tests/test_precompiled_wheel.py - - pytest -v -s device_allocator - - pytest -v -s test_uvicorn_logging.py - - pytest -v -s test_attention_backend_registry.py - - pytest -v -s test_audio_media_size_precheck.py - - pytest -v -s test_cmake_utils.py - - pytest -v -s test_connections.py - - pytest -v -s test_embedded_commit.py - - pytest -v -s test_envs.py - - pytest -v -s test_force_first_config.py - - pytest -v -s test_fxgraphcache_pickle_patch.py - - pytest -v -s test_logprobs.py - - pytest -v -s test_outputs.py - - pytest -v -s test_pcp_dp.py - - pytest -v -s test_pooling_params.py - - pytest -v -s test_quantization_revision_pin.py - - pytest -v -s test_ray_env.py - - pytest -v -s test_ray_env_utils.py - - pytest -v -s test_request_input_bounds.py - - pytest -v -s test_sampling_params.py - - pytest -v -s test_scalartype.py - - pytest -v -s test_tensorizer_revision_pin.py - - pytest -v -s test_triton_utils.py - - pytest -v -s test_version.py - - pytest -v -s test_zen_cpu_platform_detection.py - - pytest -v -s tokenizers_ - - pytest -v -s parser - - pytest -v -s transformers_utils - -- label: ":computer: (CPU) Multimodal + Config" - key: cpu-multimodal-config - depends_on: - - image-build-cpu - timeout_in_minutes: 35 - source_file_dependencies: *cpu-async-engine-deps - device: cpu-small - commands: - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s config - pytest -v -s tools/test_check_test_tethering.py @@ -754,40 +881,30 @@ steps: commands: - pytest -v -s tool_parsers -- label: ":nvidia: (A100) Batch Invariance" - key: batch-invariance-a100 +# Key kept from the former A100 step's mirror for job history and retries. +- label: ":amd: (MI250) Batch Invariance" + key: amd-batch-invariance-a100 + dind: false + device: mi250_1 timeout_in_minutes: 60 - device: a100 + env: + VLLM_GPU_MEMORY_UTILIZATION: "0.8" + VLLM_TEST_CLEAN_GPU_MEMORY: "1" + depends_on: + - image-build-amd source_file_dependencies: - - vllm/v1/attention - - vllm/model_executor/layers - - vllm/model_executor/kernels/linear - - vllm/model_executor/determinism - - tests/v1/determinism/ + - vllm/v1/attention + - vllm/model_executor/layers + - vllm/model_executor/kernels/linear + - vllm/model_executor/determinism + - tests/v1/determinism/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - docker/Dockerfile.rocm_base commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pip install pytest-timeout pytest-forked - pytest -v -s v1/determinism/test_batch_invariance.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - mirror: - amd: - label: ":amd: (MI250) Batch Invariance" - dind: false - device: mi250_1 - timeout_in_minutes: 60 - env: - VLLM_GPU_MEMORY_UTILIZATION: "0.8" - VLLM_TEST_CLEAN_GPU_MEMORY: "1" - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/v1/attention - - vllm/model_executor/layers - - vllm/model_executor/kernels/linear - - vllm/model_executor/determinism - - tests/v1/determinism/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - label: ":nvidia: (H100) Batch Invariance" key: batch-invariance-h100 @@ -803,7 +920,6 @@ steps: - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py @@ -829,7 +945,6 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py @@ -852,7 +967,6 @@ steps: - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py @@ -867,6 +981,7 @@ steps: amd: label: ":amd: (MI355) Batch Invariance" dind: false + automatic_shard: true device: mi355_1 timeout_in_minutes: 60 depends_on: @@ -882,7 +997,6 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - 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 - pytest -v -s v1/determinism/test_lora_shrink_batch_invariant.py @@ -915,6 +1029,7 @@ steps: - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + - VLLM_TEST_MODEL=ibm-granite/granite-3.0-1b-a400m-base VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py mirror: amd: label: ":amd: (MI300) Batch Invariance (Distributed)" @@ -937,10 +1052,12 @@ steps: - VLLM_TEST_TP_SIZE=4 VLLM_TEST_DISABLE_CUSTOM_AR=1 pytest -v -s v1/determinism/test_batch_invariance.py -k bs1_vs_bsN - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py + - VLLM_TEST_MODEL=ibm-granite/granite-3.0-1b-a400m-base VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py - label: ":nvidia: (H200 MIG 35GB) [Nightly] Acceptance Length (Large Models)" device: h200_35gb key: acceptance-length-test-large-models + automatic_shard: true timeout_in_minutes: 20 gpu: h100 optional: true @@ -954,6 +1071,7 @@ steps: - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Acceptance Length (Large Models)" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index d6c239e4f646..cfa925b2d4c9 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -2,55 +2,84 @@ group: Model Executor depends_on: - image-build steps: -- label: ":nvidia: (H200 MIG 35GB) Model Executor" - device: h200_35gb - key: model-executor +# Disabled until the apt packages installation is locked. +# - label: ":nvidia: (H200 MIG 35GB) Model Executor" +# device: h200_35gb +# key: model-executor +# automatic_shard: true +# timeout_in_minutes: 60 +# source_file_dependencies: +# - vllm/engine/arg_utils.py +# - vllm/config/model.py +# - vllm/model_executor +# - vllm/model_executor/warmup +# - tests/model_executor +# - tests/model_executor/test_jit_warmup.py +# - tests/model_executor/test_jit_warmup_cutedsl_launcher.py +# - tests/model_executor/test_jit_warmup_triton_launcher.py +# - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py +# commands: +# - apt-get update && apt-get install -y curl libsodium23 +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA +# # init surfaces a stack instead of silently stalling. +# - export PYTHONFAULTHANDLER=1 +# # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA +# # init) fails fast with a traceback instead of running until the global +# # build timeout. The `thread` method also handles hangs inside C/CUDA +# # calls that the signal method cannot interrupt. +# - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread +# - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread +# mirror: +# amd: +# label: ":amd: (MI300) Model Executor" +# dind: false +# device: mi300_1 +# timeout_in_minutes: 75 +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/engine/arg_utils.py +# - vllm/config/model.py +# - vllm/model_executor +# - vllm/model_executor/warmup +# - tests/model_executor +# - tests/model_executor/test_jit_warmup.py +# - tests/model_executor/test_jit_warmup_cutedsl_launcher.py +# - tests/model_executor/test_jit_warmup_triton_launcher.py +# - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# commands: +# - export VLLM_WORKER_MULTIPROC_METHOD=spawn +# - export PYTHONFAULTHANDLER=1 +# - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread +# - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread + +# Keep the AMD job: its commands do not run this installation. +- label: ':amd: (MI300) Model Executor' + device: mi300_1 + key: amd-model-executor automatic_shard: true - timeout_in_minutes: 60 + timeout_in_minutes: 75 source_file_dependencies: - vllm/engine/arg_utils.py - vllm/config/model.py - vllm/model_executor - vllm/model_executor/warmup + - vllm/models/deepseek_v32/ - tests/model_executor - tests/model_executor/test_jit_warmup.py - tests/model_executor/test_jit_warmup_cutedsl_launcher.py - tests/model_executor/test_jit_warmup_triton_launcher.py - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py commands: - - apt-get update && apt-get install -y curl libsodium23 - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA - # init surfaces a stack instead of silently stalling. - - export PYTHONFAULTHANDLER=1 - # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA - # init) fails fast with a traceback instead of running until the global - # build timeout. The `thread` method also handles hangs inside C/CUDA - # calls that the signal method cannot interrupt. - - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread - - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread - mirror: - amd: - label: ":amd: (MI300) Model Executor" - dind: false - device: mi300_1 - timeout_in_minutes: 75 - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/engine/arg_utils.py - - vllm/config/model.py - - vllm/model_executor - - vllm/model_executor/warmup - - tests/model_executor - - tests/model_executor/test_jit_warmup.py - - tests/model_executor/test_jit_warmup_cutedsl_launcher.py - - tests/model_executor/test_jit_warmup_triton_launcher.py - - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export PYTHONFAULTHANDLER=1 - - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread - - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export PYTHONFAULTHANDLER=1 + - pytest -v -s model_executor -m '(not slow_test) and (not distributed)' --timeout=900 --timeout-method=thread + - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread + dind: false + depends_on: + - image-build-amd diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index d4880005f4be..03fbd875f782 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Basic Models (Initialization)" key: basic-models-tests-initialization + automatic_shard: true timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: @@ -19,6 +20,7 @@ steps: - pytest -v -s models/test_beam_construction.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Basic Models (Initialization)" dind: false device: mi355_dpx @@ -26,9 +28,10 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) Basic Models (Extra Initialization) Shard %N" +- label: ":nvidia: (H200 MIG 35GB) Basic Models (Extra Initialization)" device: h200_35gb key: basic-models-tests-extra-initialization + automatic_shard: true timeout_in_minutes: 30 source_file_dependencies: - vllm/model_executor/models/ @@ -38,15 +41,14 @@ steps: # Only when vLLM model source is modified - test initialization of a large # subset of supported models (the complement of the small subset in the above # test.) Also run if model initialization test file is modified - - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - parallelism: 14 + - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' mirror: amd: - label: ":amd: (MI355 DPX) Basic Models (Extra Initialization) Shard %N" + label: ":amd: (MI355 DPX) Basic Models (Extra Initialization)" + automatic_shard: true dind: false device: mi355_dpx timeout_in_minutes: 60 - parallelism: 10 working_dir: "/vllm-workspace/tests" depends_on: - image-build-amd @@ -70,22 +72,40 @@ steps: - tests/models/transformers/test_backend.py - tests/models/test_registry.py - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v4_vl.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.py - tests/models/test_deepseek_v41_pcp.py - tests/models/test_deepseek_v41_pipeline.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py - - tests/models/test_deepseek_v4_vl_stream_weights.py - tests/models/test_muse_spark_rocm_loading.py - tests/models/test_muse_spark_logits_local.py - tests/models/test_muse_spark_rocm_shared_rope.py + - tests/models/test_muse_spark_rocm_mono.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py + - >- + pytest -v -s + models/test_terratorch.py + models/transformers/test_backend.py + models/test_registry.py + models/test_deepseek_v4_vl_rocm.py + models/test_deepseek_v4_vl.py + models/transformers/test_layer_registry_e2e.py + models/test_deepseek_v41_replay_start.py + models/test_deepseek_v41_decoder_replay_layers.py + models/test_deepseek_v41_replay_batch.py + models/test_deepseek_v41_pcp.py + models/test_deepseek_v41_pipeline.py + models/test_muse_spark_rocm_loading.py + models/test_muse_spark_logits_local.py + models/test_muse_spark_rocm_shared_rope.py + models/test_muse_spark_rocm_mono.py mirror: amd: label: ":amd: (MI355) Basic Models (Other)" dind: false + automatic_shard: true device: mi355_1 timeout_in_minutes: 65 depends_on: @@ -94,6 +114,7 @@ steps: - vllm/models/deepseek_v4/ - vllm/models/deepseek_v41/ - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v4_vl.py - tests/models/test_hyv4_rocm.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py @@ -105,18 +126,40 @@ steps: - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_mono_decode.py - tests/models/test_muse_spark_rocm_loading.py - - tests/models/test_deepseek_v4_vl_stream_weights.py + - tests/models/test_dsv4_mono_decode.py - tests/models/hy_v4/test_triton_ihc.py - vllm/models/muse_spark/ - tests/models/test_muse_spark_logits_local.py - tests/models/test_muse_spark_rocm_shared_rope.py + - tests/models/test_muse_spark_rocm_mono.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_logits_local.py models/test_muse_spark_rocm_shared_rope.py + - >- + pytest -v -s + models/test_terratorch.py + models/transformers/test_backend.py + models/test_registry.py + models/test_deepseek_v4_vl_rocm.py + models/test_deepseek_v4_vl.py + models/transformers/test_layer_registry_e2e.py + models/test_deepseek_v41_replay_start.py + models/test_deepseek_v41_decoder_replay_layers.py + models/test_deepseek_v41_replay_batch.py + models/test_deepseek_v41_amd_replay_forward.py + models/test_deepseek_v4_dspark_rocm.py + models/test_deepseek_v4_rocm_wo_a.py + models/test_deepseek_v4_rocm_compressor_gemm_fusion.py + models/test_dsv41_mono_decode.py + models/test_muse_spark_rocm_loading.py + models/test_muse_spark_logits_local.py + models/test_muse_spark_rocm_shared_rope.py + models/test_muse_spark_rocm_mono.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/hy_v4/test_triton_ihc.py + - pytest -v -s models/test_dsv4_mono_decode.py - label: ":amd: (MI355) MiniMax M3 MonoKernel" key: minimax-m3-monokernel-mi355 + automatic_shard: true dind: false timeout_in_minutes: 45 device: mi355_4 @@ -135,6 +178,7 @@ steps: - label: ":nvidia: (B200) Inkling" key: inkling-unit-tests-b200 + automatic_shard: true timeout_in_minutes: 40 device: b200-k8s source_file_dependencies: @@ -147,6 +191,7 @@ steps: - pytest -v -s models/inkling mirror: amd: + automatic_shard: true label: ":amd: (MI355) Inkling" dind: false device: mi355_1 @@ -163,19 +208,21 @@ steps: commands: - pytest -v -s models/inkling/rocm -- label: ":nvidia: (B200) Flint" - key: flint-unit-tests-b200 - timeout_in_minutes: 30 - device: b200-k8s - source_file_dependencies: - - vllm/models/flint/ - - tests/models/flint/ - commands: - - uv pip install --system 'gguf>=0.19.0' - - pytest -v -s models/flint +# Disabled until the gguf installation is locked. +# - label: ":nvidia: (B200) Flint" +# key: flint-unit-tests-b200 +# timeout_in_minutes: 30 +# device: b200-k8s +# source_file_dependencies: +# - vllm/models/flint/ +# - tests/models/flint/ +# commands: +# - uv pip install --system 'gguf>=0.19.0' +# - pytest -v -s models/flint - label: ":nvidia: (B200) Kimi K3" key: kimi-k3-unit-tests-b200 + automatic_shard: true timeout_in_minutes: 40 device: b200-k8s source_file_dependencies: @@ -185,9 +232,14 @@ steps: - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py commands: # The native NVIDIA Kimi K3 kernels require the SM100 family. - - pytest -v -s models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py kernels/attention/test_kimi_k3_mla_fused_epilogue.py + - >- + pytest -v -s + models/kimi_k3 + --ignore=models/kimi_k3/test_prefix_cache.py + kernels/attention/test_kimi_k3_mla_fused_epilogue.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) Kimi K3" dind: false device: mi355_1 @@ -202,44 +254,91 @@ steps: # The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355. # The kernels/ file the NVIDIA job adds is CUDA-only (it imports # cuda.bindings.driver), so it is deliberately left out here. - - pytest -v -s models/kimi_k3 + - pytest -v -s models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py -- label: ":nvidia: (4xB200) [Nightly] Kimi K3 Prefix Cache" - key: kimi-k3-prefix-cache-4xb200 - # Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs. - optional: true - automatic_shard: true - timeout_in_minutes: 75 - device: b200-k8s - num_devices: 4 - source_file_dependencies: - - vllm/models/kimi_k3/ - - vllm/v1/core/ - - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py - - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/ - - vllm/distributed/mooncake_store.py - - vllm/v1/simple_kv_offload/ - - vllm/v1/spec_decode/ - - vllm/model_executor/layers/mamba/ - - tests/models/kimi_k3/test_prefix_cache.py - - .buildkite/scripts/install-kv-connectors.sh - commands: - - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - pytest -v -s models/kimi_k3/test_prefix_cache.py +# Disabled until the KV connectors installation is locked. +# - label: ":nvidia: (4xB200) [Nightly] Kimi K3 Prefix Cache" +# key: kimi-k3-prefix-cache-4xb200 +# # Nightly only: 4-GPU B200 pods wait hours behind smaller jobs on per-PR runs. +# optional: true +# automatic_shard: true +# timeout_in_minutes: 75 +# device: b200-k8s +# num_devices: 4 +# source_file_dependencies: +# - vllm/models/kimi_k3/ +# - vllm/v1/core/ +# - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ +# - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +# - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/ +# - vllm/distributed/mooncake_store.py +# - vllm/v1/simple_kv_offload/ +# - vllm/v1/spec_decode/ +# - vllm/model_executor/layers/mamba/ +# - tests/models/kimi_k3/test_prefix_cache.py +# - .buildkite/scripts/install-kv-connectors.sh +# commands: +# - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh +# - pytest -v -s models/kimi_k3/test_prefix_cache.py +# mirror: +# amd: +# label: ":amd: (MI355) Kimi K3 Prefix Cache" +# dind: false +# device: mi355_4 +# timeout_in_minutes: 60 +# working_dir: "/vllm-workspace/tests" +# env: +# VLLM_ROCM_USE_AITER: "1" +# MOONCAKE_PROTOCOL: "tcp" +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/platforms/rocm.py +# - requirements/kv_connectors_rocm.txt - label: ":nvidia: (H200) Step5" key: step5-unit-tests + automatic_shard: true timeout_in_minutes: 20 device: h200_35gb source_file_dependencies: - vllm/models/step5/ - tests/models/step5/ + - vllm/model_executor/warmup/step5_dsa_warmup.py commands: - pytest -v -s models/step5 + mirror: + amd: + label: ":amd: (MI355 DPX) Step5" + dind: false + device: mi355_dpx + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/models/step5/ + - tests/models/step5/ + - vllm/model_executor/models/step3p7.py + - vllm/model_executor/models/step3_vl.py + - vllm/model_executor/models/step_vl.py + - vllm/config/compilation.py + - vllm/model_executor/layers/indexer_topk.py + - vllm/transformers_utils/configs/step5.py + - vllm/model_executor/warmup/step5_dsa_warmup.py + - vllm/v1/attention/backends/rocm_aiter_unified_attn.py + - vllm/triton_utils/ + - vllm/platforms/rocm.py + - tests/v1/attention/utils.py + - tests/models/test_initialization.py + - tests/models/registry.py + commands: + - pytest -v -s models/step5 + - pytest -v -s models/test_initialization.py -k Step5ForConditionalGeneration - label: ":nvidia: (H200) GLM5Next" key: glm5next-unit-tests + automatic_shard: true timeout_in_minutes: 20 device: h200_35gb source_file_dependencies: @@ -250,6 +349,7 @@ steps: - pytest -v -s models/glm5next mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) GLM5Next" dind: false device: mi355_dpx @@ -262,6 +362,7 @@ steps: - label: ":nvidia: (H200) Qwen4 Exp" key: qwen4-exp-unit-tests + automatic_shard: true timeout_in_minutes: 40 device: h200_35gb source_file_dependencies: @@ -271,20 +372,48 @@ steps: - vllm/model_executor/layers/mamba/checkpoint.py - tests/models/qwen4_exp/ commands: - - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_prepare.py models/qwen4_exp/test_qsa_reference.py + - >- + pytest -v -s + models/qwen4_exp/test_hc_ops.py + models/qwen4_exp/test_qsa_prepare.py + models/qwen4_exp/test_qsa_reference.py - pytest -v -s models/qwen4_exp/test_ple.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Qwen4 Exp" dind: false device: mi355_dpx depends_on: - image-build-amd commands: - - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py + - >- + pytest -v -s + models/qwen4_exp/test_qsa_amd.py + models/qwen4_exp/test_qsa_prepare_amd.py # PLE and QSA import conflicting HC registrations in a shared process. - pytest -v -s models/qwen4_exp/test_ple.py +- label: ":amd: (MI300) Qwen4 Exp" + key: qwen4-exp-unit-tests-mi300 + automatic_shard: true + timeout_in_minutes: 40 + dind: false + device: mi300_1 + working_dir: "/vllm-workspace/tests" + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/models/qwen4_exp/ + - vllm/config/compilation.py + - vllm/v1/attention/backends/short_conv_attn.py + - vllm/model_executor/layers/mamba/checkpoint.py + - tests/models/qwen4_exp/ + commands: + - pytest -v -s models/qwen4_exp/test_qsa_amd.py models/qwen4_exp/test_qsa_prepare_amd.py + # PLE and QSA import conflicting HC registrations in a shared process. + - pytest -v -s models/qwen4_exp/test_ple.py + - label: ":computer: (CPU) Qwen4 Exp" key: qwen4-exp-unit-tests-cpu depends_on: @@ -297,7 +426,24 @@ steps: - vllm/v1/attention/backends/short_conv_attn.py - tests/models/qwen4_exp/ commands: - - pytest -v -s models/qwen4_exp/test_config.py models/qwen4_exp/test_ple.py + - >- + pytest -v -s + models/qwen4_exp/test_config.py + models/qwen4_exp/test_ple.py + +- label: ":nvidia: (B200) Qwen3.5" + key: qwen3-5-unit-tests-b200 + timeout_in_minutes: 20 + device: b200-k8s + source_file_dependencies: + - vllm/model_executor/layers/linear.py + - vllm/model_executor/layers/mamba/gdn + - vllm/model_executor/layers/utils.py + - vllm/models/qwen3_5/ + - vllm/model_executor/models/qwen3_next.py + - tests/models/qwen3_5/ + commands: + - pytest -v -s models/qwen3_5/ - label: ":computer: (CPU) Basic Models Other" key: basic-models-test-other-cpu @@ -311,6 +457,8 @@ steps: - tests/models/test_vision.py - tests/models/test_adapters.py - tests/models/test_qwen3_5_mtp_config.py + - tests/models/test_qwen3_5.py + - tests/models/test_dory.py - tests/models/test_alice_ai.py - tests/models/test_alice_ai_config.py - tests/models/test_beam_weight_mapping.py @@ -320,16 +468,39 @@ steps: - tests/model_executor/layers/test_activation.py - tests/models/test_language_model_cache_is_weak.py - tests/models/deepseek_v32/test_sequence_parallel.py + - tests/models/deepseek_v32/test_pp_topk.py - tests/models/test_deepseek_v4_fi_moe_ep.py - tests/models/test_dspark_mla.py - tests/models/transformers/test_layer_registry.py - tests/models/test_cohere2_moe_parallel_block.py device: cpu-small commands: - - pytest -v -s models/test_utils.py models/test_vision.py models/test_adapters.py models/test_qwen3_5_mtp_config.py models/transformers/fusers/ models/transformers/test_dynamic_rope.py model_executor/layers/test_activation.py models/test_language_model_cache_is_weak.py models/deepseek_v32/test_sequence_parallel.py models/test_deepseek_v4_fi_moe_ep.py models/test_dspark_mla.py models/transformers/test_layer_registry.py models/test_alice_ai.py models/test_alice_ai_config.py models/test_beam_weight_mapping.py models/test_beam_numerics.py models/test_cohere2_moe_parallel_block.py + - >- + pytest -v -s + models/test_utils.py + models/test_vision.py + models/test_adapters.py + models/test_qwen3_5_mtp_config.py + models/test_qwen3_5.py + models/transformers/fusers/ + models/transformers/test_dynamic_rope.py + model_executor/layers/test_activation.py + models/test_language_model_cache_is_weak.py + models/deepseek_v32/test_sequence_parallel.py + models/deepseek_v32/test_pp_topk.py + models/test_deepseek_v4_fi_moe_ep.py + models/test_dspark_mla.py + models/transformers/test_layer_registry.py + models/test_dory.py + models/test_alice_ai.py + models/test_alice_ai_config.py + models/test_beam_weight_mapping.py + models/test_beam_numerics.py + models/test_cohere2_moe_parallel_block.py - label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics' key: dsv41-mono-decode-numerics-amd + automatic_shard: true dind: false device: mi355_4 num_devices: 4 @@ -347,8 +518,52 @@ steps: - pytest -v -s tests/models/test_dsv41_mono_numerics.py - pytest -v -s tests/models/test_dsv41_mono_layers.py +- label: ':amd: (MI355) DeepSeek-V4.1 Mono Decode Numerics TP8' + key: dsv41-mono-decode-numerics-tp8-amd + automatic_shard: true + dind: false + device: mi355_8 + num_devices: 8 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - docker/Dockerfile.rocm_base + - vllm/models/deepseek_v41/amd/ + - vllm/models/deepseek_v41/nvidia/model_state.py + - tests/models/test_dsv41_mono_numerics.py + - tests/models/test_dsv41_mono_layers.py + commands: + # only the TP8 cases; the 4-GPU step above runs TP2/TP4 (TP8 skips there) + - pytest -v -s "tests/models/test_dsv41_mono_numerics.py::test_mono_decode_layer_matches_vllm[8]" + - >- + pytest -v -s + "tests/models/test_dsv41_mono_layers.py::test_mono_decode_matches_the_model[8-8]" + "tests/models/test_dsv41_mono_layers.py::test_mono_decode_matches_the_model[8-28]" + +- label: ':amd: (MI355) DeepSeek-V4-Pro Mono Decode Numerics' + key: dsv4-mono-decode-numerics-amd + automatic_shard: true + dind: false + device: mi355_8 + num_devices: 8 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - docker/Dockerfile.rocm_base + - vllm/models/deepseek_v4/amd/ + - tests/models/test_dsv4_mono_numerics.py + - tests/models/test_dsv4_mono_layers.py + commands: + - pytest -v -s tests/models/test_dsv4_mono_numerics.py + - pytest -v -s tests/models/test_dsv4_mono_layers.py + - label: ':amd: (MI300) DeepSeek-V4.1 gfx942 Mono Decode Layers' key: dsv41-mono-decode-layers-gfx942-amd + automatic_shard: true dind: false device: mi300_4 num_devices: 4 @@ -393,6 +608,7 @@ steps: - label: ':amd: (MI355) Kimi-K3 Specialized Decode' key: kimi-k3-specialized-decode-amd + automatic_shard: true dind: false device: mi355_8 num_devices: 8 diff --git a/.buildkite/test_areas/models_distributed.yaml b/.buildkite/test_areas/models_distributed.yaml index 924707369f42..22e38ab1aa24 100644 --- a/.buildkite/test_areas/models_distributed.yaml +++ b/.buildkite/test_areas/models_distributed.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (L4) [Nightly] Qwen4 Exp SP" key: qwen4-exp-sp-tests-4-gpus + automatic_shard: true # Nightly only: 4-GPU L4 pods wait behind smaller jobs on per-PR runs. optional: true timeout_in_minutes: 30 @@ -22,6 +23,7 @@ steps: - pytest -v -s models/qwen4_exp/test_tpsp.py mirror: amd: + automatic_shard: true # Run on PRs even when the L4 step is optional (nightly only). optional: false label: ":amd: (MI355) Qwen4 Exp SP" @@ -39,8 +41,9 @@ steps: - tests/models/qwen4_exp/test_ple.py - tests/models/qwen4_exp/test_tpsp.py -- label: ":nvidia: (L4) Distributed Models Shard %N" +- label: ":nvidia: (L4) Distributed Models" key: distributed-model-tests-2-gpus + automatic_shard: true timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" device: l4 @@ -52,26 +55,16 @@ steps: - tests/model_executor/model_loader/test_sharded_state_loader.py - tests/models/ commands: - - | - case "$$BUILDKITE_PARALLEL_JOB" in - 0) - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - ;; - 1) - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' - ;; - 2) - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - ;; - *) exit 2 ;; - esac - parallelism: 3 + - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' mirror: amd: - label: ":amd: (MI300) Distributed Models Shard %N" + automatic_shard: true + label: ":amd: (MI300) Distributed Models" dind: false device: mi300_2 timeout_in_minutes: 80 @@ -96,20 +89,10 @@ steps: - tests/model_executor/model_loader/test_sharded_state_loader.py - tests/models/ commands: - - | - case "$$BUILDKITE_PARALLEL_JOB" in - 0) - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' && - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - ;; - 1) - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' && - pytest models/language -v -s -m 'distributed(num_gpus=2)' && - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' - ;; - 2) - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py && - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' - ;; - *) exit 2 ;; - esac + - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - HIP_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' + - pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)' + - pytest models/language -v -s -m 'distributed(num_gpus=2)' + - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'distributed(num_gpus=2)' + - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py + - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest models/multimodal/generation/test_whisper.py -v -s -m 'distributed(num_gpus=2)' diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 909bb30e5e0f..a157d0ae76f8 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Language Models (Standard)" key: language-models-tests-standard + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -11,12 +12,14 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/language commands: - # Test standard language models, excluding a subset of slow tests + # Test standard language models, excluding a subset of slow tests. + # The exclusion sits in -k so the shard planner reads -m as picking tests. - pip freeze | grep -E 'torch' - - pytest -v -s models/language -m 'core_model and (not slow_test)' + - pytest -v -s models/language -m core_model -k 'not slow_test' mirror: amd: label: ":amd: (MI355 DPX) Language Models (Standard)" + automatic_shard: true dind: false device: mi355_dpx timeout_in_minutes: 60 @@ -90,8 +93,6 @@ steps: depends_on: - image-build-amd commands: - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB # Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD @@ -127,11 +128,6 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/language/generation commands: - - export GIT_CONFIG_COUNT=1 - - export GIT_CONFIG_KEY_0=http.version - - export GIT_CONFIG_VALUE_0=HTTP/1.1 - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview' - label: ":nvidia: (H200 MIG 35GB) [Nightly] Language Models (Extended Generation)" @@ -159,12 +155,11 @@ steps: depends_on: - image-build-amd commands: - - MAMBA_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - CAUSAL_CONV1D_FORCE_BUILD=TRUE uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - label: ":nvidia: (H200 MIG 18GB) [Nightly] Language Models (PPL)" key: language-models-test-ppl + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb optional: true @@ -178,11 +173,13 @@ steps: amd: label: ":amd: (MI250) Language Models (PPL)" dind: false + automatic_shard: true device: mi250_1 timeout_in_minutes: 65 depends_on: - image-build-amd source_file_dependencies: + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/qwen2.py @@ -217,10 +214,10 @@ steps: - pytest -v -s models/language/pooling -m 'not core_model' --timeout=900 --timeout-method=thread mirror: amd: - parallelism: 4 + automatic_shard: true commands: - - "pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" - label: ":amd: (MI355 DPX) Language Models (Extended Pooling) Shard %N" + - "pytest -v -s models/language/pooling -m 'not core_model'" + label: ":amd: (MI355 DPX) Language Models (Extended Pooling)" dind: false device: mi355_dpx timeout_in_minutes: 95 diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 52847b61e986..c47651641055 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Multimodal Models (Standard) 1: qwen2" key: multi-modal-models-standard-1-qwen2 + automatic_shard: true timeout_in_minutes: 68 device: h200_18gb source_file_dependencies: @@ -17,6 +18,7 @@ steps: amd: label: ":amd: (MI355 DPX) Multimodal Models (Standard) 1: qwen2" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 65 depends_on: @@ -38,6 +40,7 @@ steps: mirror: amd: label: ":amd: (MI355 DPX) Multimodal Models (Standard) 2: qwen3 + gemma" + automatic_shard: true dind: false device: mi355_dpx timeout_in_minutes: 55 @@ -47,6 +50,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Multimodal Models (Standard) 3: llava + qwen2_vl" device: h200_35gb key: multi-modal-models-standard-3-llava-qwen2-vl + automatic_shard: true timeout_in_minutes: 40 source_file_dependencies: - vllm/ @@ -57,6 +61,7 @@ steps: - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model mirror: amd: + automatic_shard: true label: ":amd: (MI250) Multimodal Models (Standard) 3: llava + qwen2_vl" dind: false device: mi250_1 @@ -74,7 +79,19 @@ steps: - "!vllm/distributed/kv_transfer/" - tests/models/multimodal commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - >- + pytest -v -s + models/multimodal + -m core_model + --ignore models/multimodal/generation/test_common.py + --ignore models/multimodal/generation/test_ultravox.py + --ignore models/multimodal/generation/test_qwen2_5_vl.py + --ignore models/multimodal/generation/test_qwen2_vl.py + --ignore models/multimodal/generation/test_whisper.py + --ignore models/multimodal/generation/test_mm_prefix_lm.py + --ignore models/multimodal/generation/test_memory_leak.py + --ignore models/multimodal/generation/test_vit_cudagraph.py + --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest models/multimodal/generation/test_memory_leak.py -m core_model - pytest -v -s models/multimodal/generation/test_whisper.py -m core_model @@ -100,7 +117,11 @@ steps: - tests/models/registry.py device: cpu-medium commands: - - pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB + - >- + pytest -v -s + models/multimodal/processing + --ignore models/multimodal/processing/test_tensor_schema.py + --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 4 - label: ":nvidia: (H200 MIG 18GB) Multimodal Processor" @@ -128,6 +149,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Multimodal Accuracy Eval (Small Models)" device: h200_35gb key: multi-modal-accuracy-eval-small-models + automatic_shard: true timeout_in_minutes: 30 working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" source_file_dependencies: @@ -138,6 +160,7 @@ steps: - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1 mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Multimodal Accuracy Eval (Small Models)" dind: false device: mi355_dpx @@ -161,7 +184,11 @@ steps: - tests/models/multimodal/generation - tests/models/multimodal/test_mapping.py commands: - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py + - >- + pytest -v -s + models/multimodal/generation + -m 'not core_model' + --ignore models/multimodal/generation/test_common.py - pytest -v -s models/multimodal/test_mapping.py mirror: amd: @@ -215,6 +242,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) [Nightly] Multimodal Models (Extended Generation 3)" device: h200_35gb key: multi-modal-models-extended-generation-3 + automatic_shard: true optional: true source_file_dependencies: - vllm/ @@ -226,6 +254,7 @@ steps: amd: label: ":amd: (MI355 DPX) Multimodal Models (Extended Generation 3)" dind: false + automatic_shard: true device: mi355_dpx timeout_in_minutes: 90 depends_on: diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index 22189078fa4c..c84a1776d450 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -11,6 +11,13 @@ steps: source_file_dependencies: - vllm/plugins/ - tests/plugins/ + # Import-time platform resolution (see test_import_vllm_does_not_resolve_platform) + - vllm/__init__.py + - vllm/env_override.py + - vllm/platforms/__init__.py + - vllm/utils/torch_utils.py + - vllm/utils/torch_version.py + - tests/plugins_tests/test_platform_plugins.py commands: # begin platform plugin and general plugin tests, all the code in-between runs on dummy platform - pip install -e ./plugins/vllm_add_dummy_platform @@ -62,65 +69,74 @@ steps: - vllm/plugins/ - tests/plugins/ - vllm/platforms/rocm.py + - vllm/__init__.py + - vllm/env_override.py + - vllm/platforms/__init__.py + - vllm/utils/torch_utils.py + - vllm/utils/torch_version.py + - tests/plugins_tests/test_platform_plugins.py -- label: ":nvidia: (H200 MIG 18GB) [Nightly] GGUF Plugin" - key: gguf-plugin - device: h200_18gb - timeout_in_minutes: 30 - soft_fail: true - optional: true - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/plugins_tests/gguf/ - commands: - - pip install "vllm-gguf-plugin >= 0.0.2" - - pytest -v -s plugins_tests/gguf - mirror: - amd: - label: ":amd: (MI355 DPX) GGUF Plugin" - dind: false - device: mi355_dpx - timeout_in_minutes: 45 - working_dir: "/vllm-workspace/tests" - depends_on: - - image-build-amd - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/plugins_tests/gguf - - vllm/platforms/rocm.py +# Disabled until the GGUF plugin installation is locked. +# - label: ":nvidia: (H200 MIG 18GB) [Nightly] GGUF Plugin" +# key: gguf-plugin +# device: h200_18gb +# timeout_in_minutes: 30 +# soft_fail: true +# optional: true +# source_file_dependencies: +# - vllm/model_executor/layers/quantization +# - tests/plugins_tests/gguf/ +# commands: +# - pip install "vllm-gguf-plugin >= 0.0.2" +# - pytest -v -s plugins_tests/gguf +# mirror: +# amd: +# label: ":amd: (MI355 DPX) GGUF Plugin" +# dind: false +# device: mi355_dpx +# timeout_in_minutes: 45 +# working_dir: "/vllm-workspace/tests" +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - vllm/model_executor/layers/quantization +# - tests/plugins_tests/gguf +# - vllm/platforms/rocm.py -- label: ":nvidia: (H200 MIG 18GB) [Nightly] BitsAndBytes Plugin" - key: bitsandbytes-plugin - device: h200_18gb - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/tests" - soft_fail: true - optional: true - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/plugins_tests/bitsandbytes - commands: - # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so - # there is nothing to wrap on the vLLM side. - - unset VLLM_GPU_SYNC_CHECK - - pip install "vllm-bnb-plugin >= 0.0.1" - - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed' +# Disabled until the BitsAndBytes plugin installation is locked. +# - label: ":nvidia: (H200 MIG 18GB) [Nightly] BitsAndBytes Plugin" +# key: bitsandbytes-plugin +# device: h200_18gb +# timeout_in_minutes: 30 +# working_dir: "/vllm-workspace/tests" +# soft_fail: true +# optional: true +# source_file_dependencies: +# - vllm/model_executor/layers/quantization +# - tests/plugins_tests/bitsandbytes +# commands: +# # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so +# # there is nothing to wrap on the vLLM side. +# - unset VLLM_GPU_SYNC_CHECK +# - pip install "vllm-bnb-plugin >= 0.0.1" +# - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed' -- label: ":nvidia: (L4) [Nightly] BitsAndBytes Plugin" - key: bitsandbytes-plugin-2-gpus - timeout_in_minutes: 15 - working_dir: "/vllm-workspace/tests" - device: l4 - num_devices: 2 - soft_fail: true - optional: true - source_file_dependencies: - - vllm/model_executor/layers/quantization - - tests/plugins_tests/bitsandbytes - commands: - # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so - # there is nothing to wrap on the vLLM side. - - unset VLLM_GPU_SYNC_CHECK - - pip install "vllm-bnb-plugin >= 0.0.1" - - pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)' +# Disabled until the BitsAndBytes plugin installation is locked. +# - label: ":nvidia: (L4) [Nightly] BitsAndBytes Plugin" +# key: bitsandbytes-plugin-2-gpus +# timeout_in_minutes: 15 +# working_dir: "/vllm-workspace/tests" +# device: l4 +# num_devices: 2 +# soft_fail: true +# optional: true +# source_file_dependencies: +# - vllm/model_executor/layers/quantization +# - tests/plugins_tests/bitsandbytes +# commands: +# # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so +# # there is nothing to wrap on the vLLM side. +# - unset VLLM_GPU_SYNC_CHECK +# - pip install "vllm-bnb-plugin >= 0.0.1" +# - pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)' diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 1dc6659e8e99..33638dd04f27 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -57,12 +57,14 @@ steps: - label: ":nvidia: (H200 MIG 35GB) PyTorch Compilation / Dynamic Shapes" device: h200_35gb key: pytorch-compilation-dynamic-shapes + automatic_shard: true timeout_in_minutes: 30 source_file_dependencies: *pytorch-compilation-deps commands: - pytest -s -v compile/dynamic_shapes/ mirror: amd: + automatic_shard: true label: ":amd: (MI300) PyTorch Compilation / Dynamic Shapes" dind: false device: mi300_1 @@ -75,6 +77,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) PyTorch Compilation H100 Cases" key: pytorch-compilation-unit-tests-h100 + automatic_shard: true timeout_in_minutes: 30 device: h200_18gb source_file_dependencies: @@ -108,6 +111,7 @@ steps: - "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) PyTorch Compilation H100 Cases" dind: false device: mi355_dpx @@ -155,7 +159,10 @@ steps: - vllm/v1/ - tests/compile/passes commands: - - pytest -s -v compile/passes --ignore compile/passes/distributed + - >- + pytest -s -v + compile/passes + --ignore compile/passes/distributed mirror: amd: label: ':amd: (MI300) PyTorch Compilation Passes' @@ -171,11 +178,15 @@ steps: - tests/compile/passes - vllm/platforms/rocm.py commands: - - pytest -s -v compile/passes --ignore compile/passes/distributed + - >- + pytest -s -v + compile/passes + --ignore compile/passes/distributed - label: ":nvidia: (H200 MIG 35GB) PyTorch Fullgraph" device: h200_35gb key: pytorch-fullgraph-test + automatic_shard: true timeout_in_minutes: 90 source_file_dependencies: - vllm/__init__.py @@ -206,13 +217,11 @@ steps: - tests/compile - csrc/ commands: - # Run smoke tests under fullgraph directory, except test_full_cudagraph.py - # as it is a heavy test that is covered in other steps. - # Use `find` to launch multiple instances of pytest so that - # they do not suffer from https://github.com/vllm-project/vllm/issues/28965 - - "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'" + # The heavy cudagraph test is covered in other steps. + - pytest -v -s compile/fullgraph --ignore=compile/fullgraph/test_full_cudagraph.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) PyTorch Fullgraph" dind: false device: mi355_dpx @@ -229,6 +238,7 @@ steps: # those SM90 cases skip while the architecture-compatible cases still run. - label: ":nvidia: (L4) PyTorch Fullgraph CUDAGraph Compatibility" key: pytorch-fullgraph-cudagraph-l4-compatibility + automatic_shard: true device: l4 num_devices: 1 timeout_in_minutes: 60 @@ -263,6 +273,7 @@ steps: - pytest -s -v compile/fullgraph/test_full_cudagraph.py mirror: amd: + automatic_shard: true label: ':amd: (MI250) PyTorch Fullgraph CUDAGraph Compatibility' dind: false device: mi250_1 @@ -283,6 +294,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) PyTorch Nightly Dependency Override Check" key: pytorch-nightly-dependency-override-check + automatic_shard: true # if this test fails, it means the nightly torch version is not compatible with some # of the dependencies. Please check the error message and add the package to whitelist # in /vllm/tools/pre_commit/generate_nightly_torch_test.py @@ -294,6 +306,7 @@ steps: - bash standalone_tests/pytorch_nightly_dependency.sh mirror: amd: + automatic_shard: true label: ":amd: (MI300) PyTorch Nightly Dependency Override Check" dind: false device: mi300_1 diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 50cfc2ee86fc..9349c767be9f 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -2,54 +2,61 @@ group: Quantization depends_on: - image-build steps: -- label: ":nvidia: (H200 MIG 35GB) Quantization" - device: h200_35gb - key: quantization - automatic_shard: true - timeout_in_minutes: 40 - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/models/deepseek_v41/quant_config.py - - vllm/utils/humming.py - - tests/quantization - commands: - # Pin torchao to the version compatible with the CI PyTorch and CUDA stack. - - uv pip install --system torchao==0.18.0 --index-url https://download.pytorch.org/whl/cu130 - - uv pip install --system conch-triton-kernels - # The SM90-only checkpoint currently contains a removed weight_chan_scale - # parameter. It was not exercised by the previous L4 job. - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py -k 'not test_compressed_tensors_w4a8_fp8' - mirror: - amd: - label: ":amd: (MI300) Quantization" - automatic_shard: true - dind: false - device: mi300_1 - timeout_in_minutes: 115 - working_dir: "/vllm-workspace/tests" - depends_on: - - image-build-amd - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - tests/quantization - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - vllm/model_executor/layers/fused_moe/oracle/unquantized.py - - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - - tests/rocm/test_moe_weight_replay.py - - tests/rocm/test_moe_padding.py - commands: - - uv pip install --system torchao==0.17.0 - - uv pip install --system conch-triton-kernels - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py - - pytest -v -s rocm/test_moe_weight_replay.py - - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py +# Disabled until torchao transitive dependencies are locked. +# - label: ":nvidia: (H200 MIG 35GB) Quantization" +# device: h200_35gb +# key: quantization +# automatic_shard: true +# timeout_in_minutes: 40 +# source_file_dependencies: +# - csrc/ +# - vllm/model_executor/layers/quantization +# - vllm/model_executor/models/qwen3_dflash.py +# - vllm/model_executor/models/utils.py +# - vllm/transformers_utils/model_arch_config_convertor.py +# - vllm/models/deepseek_v41/quant_config.py +# - vllm/utils/humming.py +# - tests/quantization +# commands: +# # Pin torchao to the version compatible with the CI PyTorch and CUDA stack. +# - uv pip install --system torchao==0.18.0 --index-url https://download.pytorch.org/whl/cu130 +# - uv pip install --system conch-triton-kernels +# # The SM90-only checkpoint currently contains a removed weight_chan_scale +# # parameter. It was not exercised by the previous L4 job. +# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py -k 'not test_compressed_tensors_w4a8_fp8' +# mirror: +# amd: +# label: ":amd: (MI300) Quantization" +# automatic_shard: true +# dind: false +# device: mi300_1 +# timeout_in_minutes: 115 +# working_dir: "/vllm-workspace/tests" +# depends_on: +# - image-build-amd +# source_file_dependencies: +# - csrc/ +# - vllm/model_executor/layers/quantization +# - vllm/model_executor/models/qwen3_dflash.py +# - vllm/model_executor/models/utils.py +# - tests/quantization +# - vllm/_aiter_ops.py +# - vllm/platforms/rocm.py +# - vllm/model_executor/layers/fused_moe/oracle/unquantized.py +# - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +# - tests/rocm/test_moe_weight_replay.py +# - tests/rocm/test_moe_padding.py +# commands: +# - uv pip install --system torchao==0.17.0 +# - uv pip install --system conch-triton-kernels +# - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py +# - pytest -v -s rocm/test_moe_weight_replay.py +# - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py - label: ":nvidia: (H200 MIG 35GB) Quantized Fusions" device: h200_35gb key: quantized-fusions + automatic_shard: true timeout_in_minutes: 20 source_file_dependencies: - tests/fusion @@ -61,6 +68,7 @@ steps: - pytest -v -s fusion/ mirror: amd: + automatic_shard: true label: ":amd: (MI300) Quantized Fusions" dind: false device: mi300_1 @@ -84,6 +92,9 @@ steps: device: b200-k8s source_file_dependencies: - tests/quantization/test_blackwell_moe.py + - tests/benchmarks/test_forward_cli.py + - vllm/benchmarks/forward/ + - vllm/models/muse_spark/ - vllm/model_executor/models/deepseek_v2.py - vllm/model_executor/models/gpt_oss.py - vllm/model_executor/models/llama4.py @@ -93,9 +104,13 @@ steps: - vllm/model_executor/layers/quantization/mxfp4.py - vllm/v1/attention/backends/flashinfer.py commands: - - pytest -s -v tests/quantization/test_blackwell_moe.py + - >- + pytest -v -s + tests/quantization/test_blackwell_moe.py + tests/benchmarks/test_forward_cli.py::test_bench_forward_muse_balanced_routing mirror: amd: + automatic_shard: true label: ":amd: (MI355) Quantized MoE" dind: false device: mi355_1 diff --git a/.buildkite/test_areas/rl.yaml b/.buildkite/test_areas/rl.yaml index 7141a81bde1a..6ef92a35c754 100644 --- a/.buildkite/test_areas/rl.yaml +++ b/.buildkite/test_areas/rl.yaml @@ -14,7 +14,10 @@ steps: commands: - pytest -v -s entrypoints/rl/weight_transfer - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s entrypoints/rl --ignore=entrypoints/rl/weight_transfer + - >- + pytest -v -s + entrypoints/rl + --ignore=entrypoints/rl/weight_transfer mirror: amd: label: ":amd: (MI355 DPX) RL Entrypoints" diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index dbc6af87f104..043429033932 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (H200 MIG 18GB) Rust Frontend OpenAI Coverage" key: rust-frontend-openai-coverage + automatic_shard: true timeout_in_minutes: 50 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -62,6 +63,7 @@ steps: mirror: amd: label: ":amd: (MI355 DPX) Rust Frontend OpenAI Coverage" + automatic_shard: true dind: false device: mi355_dpx timeout_in_minutes: 55 @@ -154,6 +156,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Rust Frontend Core Correctness" key: rust-frontend-core-correctness + automatic_shard: true timeout_in_minutes: 20 device: h200_18gb working_dir: "/vllm-workspace/tests" @@ -168,6 +171,7 @@ steps: - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Rust Frontend Core Correctness" dind: false device: mi355_dpx @@ -183,6 +187,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Rust Frontend Tool Use" key: rust-frontend-tool-use + automatic_shard: true device: h200_35gb timeout_in_minutes: 35 working_dir: "/vllm-workspace/tests" @@ -199,7 +204,11 @@ steps: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn # The Rust frontend has no `mistral` tokenizer mode or renderer. - - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 + - >- + pytest -v -s + tool_use + --ignore=tool_use/mistral + --models llama3.2 # The Rust renderer rejects empty tool-call arguments in the history (Python treats them as {}). - pytest -v -s entrypoints/openai/chat_completion/test_completion_with_function_calling.py -k "not test_no_args_tool_call" - pytest -v -s entrypoints/tool_parsers/test_granite4_tool_parser.py @@ -207,6 +216,7 @@ steps: - pytest -v -s entrypoints/tool_parsers/test_hermes_tool_parser.py -k granite4 mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Rust Frontend Tool Use" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index 9d66e785ea8a..68055d6a0fd7 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.buildkite/test_areas/samplers.yaml @@ -5,6 +5,7 @@ steps: - label: ":nvidia: (H200 MIG 18GB) Samplers" device: h200_18gb key: samplers-test + automatic_shard: true timeout_in_minutes: 60 source_file_dependencies: - vllm/model_executor/layers @@ -19,6 +20,7 @@ steps: - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data" mirror: amd: + automatic_shard: true label: ":amd: (MI250) Samplers" dind: false device: mi250_1 @@ -41,6 +43,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search" device: h200_35gb key: samplers-multimodal-beam-search + automatic_shard: true timeout_in_minutes: 30 source_file_dependencies: - vllm/model_executor/layers @@ -52,6 +55,7 @@ steps: - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Samplers Multimodal Beam Search" dind: false device: mi355_dpx diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index d451ec07ea9b..32d377ca5efa 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -36,6 +36,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Spec Decode Eagle 1: DeepSeek + Qwen" key: spec-decode-eagle-1-deepseek-qwen + automatic_shard: true timeout_in_minutes: 25 device: h200_35gb source_file_dependencies: @@ -46,6 +47,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3" mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode Eagle 1: DeepSeek + Qwen" dind: false device: mi355_dpx @@ -63,6 +65,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" key: spec-decode-eagle-2-llama3-qwen-vl-other + automatic_shard: true timeout_in_minutes: 25 device: h200_35gb source_file_dependencies: @@ -73,6 +76,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3" mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" dind: false device: mi355_dpx @@ -127,6 +131,7 @@ steps: # each job below runs whole directories. - label: ":nvidia: (H200 MIG 35GB) Spec Decode Speculators" key: spec-decode-speculators + automatic_shard: true timeout_in_minutes: 30 device: h200_35gb source_file_dependencies: @@ -139,6 +144,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/speculators/ mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode Speculators" dind: false device: mi355_dpx @@ -157,6 +163,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP Gemma4" key: spec-decode-mtp-gemma4 + automatic_shard: true timeout_in_minutes: 30 device: h200_35gb source_file_dependencies: @@ -169,6 +176,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/mtp/gemma4/ mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode MTP Gemma4" dind: false device: mi355_dpx @@ -187,6 +195,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP DeepSeek + MiMo" key: spec-decode-mtp-deepseek-mimo + automatic_shard: true timeout_in_minutes: 30 device: h200_35gb source_file_dependencies: @@ -200,6 +209,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/mtp/mimo/ mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode MTP DeepSeek + MiMo" dind: false device: mi355_dpx @@ -218,6 +228,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Spec Decode MTP Qwen3.5" key: spec-decode-mtp-qwen3-5 + automatic_shard: true timeout_in_minutes: 30 device: h200_35gb source_file_dependencies: @@ -225,11 +236,16 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - vllm/v1/attention/backends/ - vllm/transformers_utils/configs/speculators/ + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/models/qwen3_5/nvidia/mtp.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/config/speculative.py - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/mtp/qwen3_5/ mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode MTP Qwen3.5" dind: false device: mi355_dpx @@ -243,6 +259,9 @@ steps: - vllm/v1/sample/ - vllm/model_executor/layers/ - vllm/transformers_utils/configs/speculators/ + - vllm/model_executor/models/qwen3_5_mtp.py + - vllm/model_executor/models/qwen3_next_mtp.py + - vllm/config/speculative.py - tests/v1/e2e/spec_decode/ - vllm/platforms/rocm.py @@ -286,6 +305,7 @@ steps: - label: ":nvidia: (H200 MIG 35GB) Spec Decode N-Gram + Suffix" key: spec-decode-ngram-suffix + automatic_shard: true timeout_in_minutes: 20 device: h200_35gb source_file_dependencies: @@ -298,6 +318,7 @@ steps: - python3 spec_decode/test_custom_proposer.py mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Spec Decode N-Gram + Suffix" dind: false device: mi355_dpx @@ -344,6 +365,7 @@ steps: - label: ":nvidia: (B200) [Nightly] Spec Decode Draft Model" key: spec-decode-draft-model-nightly-b200 + automatic_shard: true timeout_in_minutes: 40 device: b200-k8s optional: true @@ -351,8 +373,11 @@ steps: - vllm/v1/spec_decode/ - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - tests/models/test_qwen3_dflash_mxfp8_context_kv.py commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ + - pytest -v -s models/test_qwen3_dflash_mxfp8_context_kv.py mirror: amd: label: ":amd: (MI355) Spec Decode Draft Model Nightly" @@ -377,6 +402,7 @@ steps: - label: ":nvidia: (H100) [Nightly] Speculators Correctness" key: speculators-correctness + automatic_shard: true timeout_in_minutes: 30 device: h100 optional: true @@ -391,6 +417,7 @@ steps: - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test mirror: amd: + automatic_shard: true label: ":amd: (MI355 DPX) Speculators Correctness Nightly" dind: false device: mi355_dpx @@ -403,6 +430,7 @@ steps: - label: ":nvidia: (B200) [Nightly] Spec Decode DeepSeek MTP Parallel Load" key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300 + automatic_shard: true timeout_in_minutes: 30 device: b200-k8s optional: true @@ -418,6 +446,7 @@ steps: - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py mirror: amd: + automatic_shard: true label: ":amd: (MI355) Spec Decode DeepSeek MTP Parallel Load" dind: false device: mi355_2 @@ -449,12 +478,15 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/models/laguna_dflash.py - tests/v1/e2e/spec_decode/ commands: # DFlash2 NVFP4 (27B) does not fit a 35GB MIG slice; runs on B200 below. - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" + - pytest -v -s v1/e2e/spec_decode/test_distribution_shift.py mirror: amd: label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly" @@ -464,6 +496,7 @@ steps: timeout_in_minutes: 120 commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" + - pytest -v -s v1/e2e/spec_decode/test_distribution_shift.py depends_on: - image-build-amd source_file_dependencies: @@ -483,6 +516,7 @@ steps: - label: ":nvidia: (B200) [Nightly] Spec Decode AL DFlash2" key: spec-decode-dflash2-nightly-b200 + automatic_shard: true timeout_in_minutes: 60 device: b200-k8s optional: true @@ -491,12 +525,15 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "dflash2" mirror: amd: + automatic_shard: true label: ":amd: (MI355) Spec Decode AL DFlash2 Nightly" dind: false device: mi355_1 @@ -516,7 +553,9 @@ steps: - vllm/v1/attention/ - vllm/model_executor/models/qwen3_dflash.py - vllm/model_executor/models/qwen3_dflash2.py + - vllm/models/qwen3_5/ - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/qwen3_5_mtp.py - vllm/model_executor/model_loader/ - vllm/model_executor/layers/quantization/ - vllm/model_executor/kernels/linear/ @@ -524,6 +563,7 @@ steps: - label: ":nvidia: (H100) [Nightly] Spec Decode AL XPress" key: spec-decode-xpress-nightly + automatic_shard: true timeout_in_minutes: 30 device: h100 optional: true @@ -590,7 +630,13 @@ steps: # Each Gemma4 parametrization is an indivisible ~23m floor. Keep them in # separate shards and put the remaining tests in shard 0, where new tests # land by default. - - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/ --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi + - >- + if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then + pytest -v -s + v1/e2e/spec_decode/acceptance_rates/mtp_other/ + --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]" + --deselect "tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; + fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi mirror: diff --git a/.buildkite/test_areas/torch_abi.yaml b/.buildkite/test_areas/torch_abi.yaml index 576f9b8c9e10..4c38ba0ff2e3 100644 --- a/.buildkite/test_areas/torch_abi.yaml +++ b/.buildkite/test_areas/torch_abi.yaml @@ -4,6 +4,7 @@ depends_on: steps: - label: ":nvidia: (L4) Torch Stable ABI Audit" key: torch-stable-abi-audit + automatic_shard: true timeout_in_minutes: 10 source_file_dependencies: - .buildkite/check-torch-abi.py @@ -12,22 +13,23 @@ steps: - setup.py commands: - python3 /vllm-workspace/.buildkite/check-torch-abi.py - mirror: - amd: - label: ":amd: (MI300) Torch Stable ABI Audit" - dind: false - device: mi300_1 - timeout_in_minutes: 35 - depends_on: - - image-build-amd - no_gpu: true - working_dir: /vllm-workspace/ - source_file_dependencies: - - .buildkite/check-torch-abi.py - - csrc/ - - cmake/ - - setup.py - - vllm/platforms/rocm.py - commands: - - uv pip install --system torch-abi-audit==0.0.1 - - python3 /vllm-workspace/.buildkite/check-torch-abi.py + # Disabled until torch-abi-audit transitive dependencies are locked. + # mirror: + # amd: + # label: ":amd: (MI300) Torch Stable ABI Audit" + # dind: false + # device: mi300_1 + # timeout_in_minutes: 35 + # depends_on: + # - image-build-amd + # no_gpu: true + # working_dir: /vllm-workspace/ + # source_file_dependencies: + # - .buildkite/check-torch-abi.py + # - csrc/ + # - cmake/ + # - setup.py + # - vllm/platforms/rocm.py + # commands: + # - uv pip install --system torch-abi-audit==0.0.1 + # - python3 /vllm-workspace/.buildkite/check-torch-abi.py diff --git a/.dockerignore b/.dockerignore index e0284129eefa..7ed1895898df 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,8 +1,9 @@ /.venv /build -/artifacts/vllm-rocm-install/ -/bake-config-build-*.json -/wheel-export/ +/.rocm-version/ +/.rocm-wheel/ +/.rocm-kernel-map/ +/.rocm-src/ dist vllm/*.so vllm/vllm-rs diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a50febef8c49..2b72eafe2222 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -173,6 +173,7 @@ mkdocs.yaml @hmellor /tools/pre_commit/.markdownlint.yaml @hmellor .pre-commit-config.yaml @hmellor /tools/pre_commit @hmellor +/tools/pre_commit/*_allowlist.txt # CPU /vllm/v1/worker/cpu* @bigPYJ1151 @xuechendi @@ -219,7 +220,8 @@ mkdocs.yaml @hmellor /vllm/v1/attention/ops/triton_unified_attention.py @tdoublep /vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy /vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl @gau-nernst -/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py @gau-nernst +/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_sm100.py @gau-nernst +/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_sm90.py @gau-nernst /vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py @gau-nernst /vllm/cute_utils @gau-nernst diff --git a/.github/mergify.yml b/.github/mergify.yml index bd364fd4e4ac..7ece42cdecfe 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -14,6 +14,19 @@ pull_request_rules: comment: message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/" +- name: label-recipe + description: Automatically apply recipe label + conditions: + - label != stale + - -closed + - files~=^recipes/ + actions: + label: + add: + - recipe + comment: + message: "Recipe preview: https://vllm-recipes-git-preview-vllm-pr-{{number}}-inferact-inc.vercel.app" + - name: comment-pre-commit-failure description: Comment on PR when pre-commit check fails conditions: @@ -739,6 +752,28 @@ pull_request_rules: - "stmatengss" - "LCAIZJ" +- name: label-nixl + description: Automatically apply nixl label + conditions: + - label != stale + - or: + # Any NIXL-named file: KV connector and shared utils, NIXL EP + # dispatch, EC-transfer NIXL data, the offload p2p data layer, + # tests, docs, examples. + - files~=(?i)nixl + # KV offloading paths (multi-tiered KVCR tiering and simple), plus + # their tests and docs. + - files~=kv_offload + - title~=(?i)nixl + actions: + label: + add: + - nixl + assign: + users: + - "ovidiusm" + - "iyastreb" + - name: label-scheduler description: Automatically apply scheduler label conditions: diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml index dd695d28d529..15c7ce0e1179 100644 --- a/.github/workflows/buf.yml +++ b/.github/workflows/buf.yml @@ -17,8 +17,8 @@ jobs: proto: runs-on: ubuntu-latest steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: bufbuild/buf-action@85aebf73123b5c15fd5528aaecbf9129cddf7fa7 # v1.6.0 with: version: 1.72.0 input: rust/proto diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 43a458df7678..03eeb71a5bda 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -15,7 +15,7 @@ jobs: steps: - name: Label issues based on keywords id: label-step - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | // Configuration: Add new labels and keywords here @@ -633,7 +633,7 @@ jobs: - name: CC users for labeled issues if: steps.label-step.outputs.labels_added != '[]' - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | // Configuration: Map labels to GitHub users to CC @@ -714,7 +714,7 @@ jobs: - name: Request missing ROCm info from issue author if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug') - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 + uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 with: script: | const body = (context.payload.issue.body || '').toLowerCase(); diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index 2e63c5676dd8..8de27d848e8b 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -28,9 +28,9 @@ jobs: timeout-minutes: 30 steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: enable-cache: true cache-dependency-glob: | diff --git a/.github/workflows/notify-ci-authorized.yml b/.github/workflows/notify-ci-authorized.yml index c260cc67dba0..1396032d51cb 100644 --- a/.github/workflows/notify-ci-authorized.yml +++ b/.github/workflows/notify-ci-authorized.yml @@ -34,11 +34,11 @@ jobs: github.event.workflow_run.conclusion == 'success')) runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.event.repository.default_branch }} persist-credentials: false - - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: python-version: "3.12" - name: Notify PR author that CI is available diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 861d4a5e498b..1410a8944828 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -49,14 +49,18 @@ jobs: if: github.repository_owner == 'vllm-project' && always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: python-version: "3.12" - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json" - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json" - run: echo "::add-matcher::.github/workflows/matchers/mypy.json" - run: echo "::add-matcher::.github/workflows/matchers/shellcheck.json" - - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 + # Replaces pre-commit/action, whose latest release (v3.0.1) uses the node20 actions/cache@v4 + - uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 with: - extra_args: --all-files --hook-stage manual + path: ~/.cache/pre-commit + key: pre-commit-uv|${{ runner.os }}|${{ env.UV_PYTHON }}|${{ hashFiles('.pre-commit-config.yaml') }} + - run: >- + uvx pre-commit run --show-diff-on-failure --color=always --all-files --hook-stage manual diff --git a/.github/workflows/proto-crate.yml b/.github/workflows/proto-crate.yml index 43be5c8c9214..ad20db39fa66 100644 --- a/.github/workflows/proto-crate.yml +++ b/.github/workflows/proto-crate.yml @@ -24,8 +24,8 @@ jobs: semver: runs-on: ubuntu-latest steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: obi1kenobi/cargo-semver-checks-action@6b69fcf40e9b5fb17adeb57e4b6ecd020649a239 # v2.9 with: manifest-path: rust/proto/Cargo.toml package: vllm-proto @@ -33,7 +33,7 @@ jobs: check: runs-on: ubuntu-latest steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - run: rustup show active-toolchain - name: Check release version if: github.event_name == 'push' @@ -56,13 +56,13 @@ jobs: contents: read id-token: write steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.sha }} - run: rustup show active-toolchain - name: Verify package before authentication run: cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run - - uses: rust-lang/crates-io-auth-action@c6f97d42243bad5fab37ca0427f495c86d5b1a18 # v1.0.5 + - uses: rust-lang/crates-io-auth-action@3268eb413133e2960fa4ae0e963ca76380c54cd9 # v1.0.6 id: auth - name: Publish env: diff --git a/.github/workflows/recipes-preview.yml b/.github/workflows/recipes-preview.yml new file mode 100644 index 000000000000..efaca4681a67 --- /dev/null +++ b/.github/workflows/recipes-preview.yml @@ -0,0 +1,69 @@ +name: Recipes Preview + +# Tells vllm-project/recipes when recipes/ changes, instead of waiting for its +# scheduled polling, which GitHub runs hours late: +# - a recipe PR is opened, pushed or reopened: rebuild its preview branch +# `preview/vllm-pr-`, which Vercel deploys at the environment URL below +# (recipes debounces bursts of pushes); +# - a recipe PR is closed: delete that branch; +# - recipe changes reach main: sync recipes.vllm.ai (`recipes-updated`). +# +# Runs on pull_request_target so fork PRs can reach the dispatch token. This +# is safe only because no step checks out or runs PR code, and only the PR +# number leaves this workflow. Keep it that way. The token comes from a GitHub +# App installed only on vllm-project/recipes with contents: write; its key +# lives in the `recipes-preview` environment, which only main may use. + +on: + pull_request_target: + types: [opened, synchronize, reopened, closed] + paths: + - 'recipes/**' + push: + branches: [main] + paths: + - 'recipes/**' + +permissions: + contents: read + +concurrency: + group: recipes-preview-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + dispatch: + if: github.repository == 'vllm-project/vllm' + runs-on: ubuntu-latest + timeout-minutes: 5 + environment: + name: recipes-preview + url: >- + ${{ github.event_name == 'push' && 'https://recipes.vllm.ai' || + (github.event.action != 'closed' && + format('https://vllm-recipes-git-preview-vllm-pr-{0}-inferact-inc.vercel.app', github.event.pull_request.number) || + '') }} + steps: + - name: Create a token for vllm-project/recipes + id: token + uses: actions/create-github-app-token@64ab729a6294000f5aa4a4bc0b73d1f32bc45510 # v3.3.0 + with: + app-id: ${{ vars.RECIPES_APP_ID }} + private-key: ${{ secrets.RECIPES_APP_PRIVATE_KEY }} + owner: vllm-project + repositories: recipes + permission-contents: write + + - name: Dispatch to vllm-project/recipes + env: + GH_TOKEN: ${{ steps.token.outputs.token }} + EVENT: ${{ github.event_name }} + PR: ${{ github.event.pull_request.number }} + run: | + if [ "$EVENT" = push ]; then + gh api repos/vllm-project/recipes/dispatches -f event_type=recipes-updated + exit 0 + fi + [[ "$PR" =~ ^[0-9]+$ ]] || { echo "Invalid PR number: $PR"; exit 1; } + gh api repos/vllm-project/recipes/dispatches \ + -f event_type=preview-pr -F "client_payload[pr]=${PR}" diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index d3336772717f..877c2d91b94b 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -22,11 +22,11 @@ jobs: startsWith(github.event.comment.body, '/amd-ci')) runs-on: [self-hosted, linux, x64, vllm-runners] steps: - - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.event.repository.default_branch }} persist-credentials: false - - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 + - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3 with: python-version: "3.12" - name: Authorize and run CI command diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index 64ec8cb7623b..f249c08ce5c2 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -76,11 +76,9 @@ } RETRY_STATES = "failed,timed_out,expired" CANCELABLE_BUILD_STATES = ("scheduled", "running", "failing") -SETUP_STEP_KEYS = { - "ensure-ci-base-amd", - "pre-commit", - "refresh-rocm-base-amd", -} +# ensure-ci-base-amd: the old ROCm build chain, which ci-infra still renders +# for branches that predate .buildkite/scripts/rocm/build.sh. +SETUP_STEP_KEYS = {"ensure-ci-base-amd", "refresh-rocm-base-amd", "pre-commit"} class ApiError(RuntimeError): diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index ba807fab7c3a..d2c5533d31f6 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -15,7 +15,7 @@ jobs: actions: write runs-on: ubuntu-latest steps: - - uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0 + - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11.0.0 with: # Increasing this value ensures that changes to this workflow # propagate to all issues and PRs in days rather than months diff --git a/.gitignore b/.gitignore index 57d0d54fc87b..533f9cfa4921 100644 --- a/.gitignore +++ b/.gitignore @@ -251,3 +251,9 @@ vllm/grpc/vllm_engine_pb2.pyi # Ignore generated cpu headers csrc/cpu/cpu_attn_dispatch_generated.h + +# Written by .buildkite/scripts/rocm/build.sh +/.rocm-src/ +/.rocm-version/ +/.rocm-wheel/ +/.rocm-kernel-map/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c940db4ae844..b42d38d93a31 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,42 +5,43 @@ default_stages: - pre-commit # Run locally - manual # Run in CI exclude: 'vllm/third_party/.*|vllm/models/kimi_k3/nvidia/ops/third_party/.*|vllm/models/kimi_k3/amd/ops/third_party/.*' +# Revs are frozen to commit SHAs, bump them with `pre-commit autoupdate --freeze [--repo ]` repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.14.0 + rev: f9351c924055bf6c7b4a4670237d3ce141e9f57c # frozen: v0.14.0 hooks: - id: ruff-check args: [--output-format, github, --fix] - id: ruff-format - repo: https://github.com/crate-ci/typos - rev: v1.43.5 + rev: 9023c1f779400877cafdd98052113fc777e2cfdc # frozen: v1.51.1 hooks: - id: typos args: [--force-exclude] - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v21.1.2 + rev: 98ad4ddb9c057de2f5b9e93e44e6db959807d68c # frozen: v23.1.3 hooks: - id: clang-format exclude: 'csrc/moe/topk_softmax_kernels.cu|vllm/third_party/.*' types_or: [c++, cuda] args: [--style=file, --verbose] - repo: https://github.com/DavidAnson/markdownlint-cli2 - rev: v0.21.0 + rev: 341e2a331ba578029e66c2e01183e7a6b5ca9ab4 # frozen: v0.24.0 hooks: - id: markdownlint-cli2 - language_version: lts + language_version: 24.22.1 # Node 26 (current lts) needs libatomic, which the CI runners lack args: [--config, tools/pre_commit/.markdownlint.yaml, --fix] - repo: https://github.com/rhysd/actionlint - rev: v1.7.7 + rev: 914e7df21a07ef503a81201c76d2b11c789d3fca # frozen: v1.7.12 hooks: - id: actionlint - repo: https://github.com/shellcheck-py/shellcheck-py - rev: 745eface02aef23e168a8afb6b5737818efbea95 # v0.11.0.1 + rev: a6564f5d57444e123433a5c9af733e3975c77229 # frozen: v0.11.0.1-1 hooks: - id: shellcheck args: ["-s", "bash", "--format=gcc", "--rcfile=tools/pre_commit/.shellcheckrc"] - repo: https://github.com/astral-sh/uv-pre-commit - rev: 0.11.1 + rev: 46787482b9fb2003c56e0f7592e312ef381829d9 # frozen: 0.13.6 hooks: - id: pip-compile args: [ @@ -74,6 +75,7 @@ repos: --no-emit-package, cupy-cuda12x, # nvidia packages (unsuffixed / unified naming) --no-emit-package, nvidia-cublas, + --no-emit-package, nvidia-cuda-cccl, --no-emit-package, nvidia-cuda-cupti, --no-emit-package, nvidia-cuda-nvrtc, --no-emit-package, nvidia-cuda-runtime, @@ -158,7 +160,7 @@ repos: ] files: ^requirements/docs\.(in|txt)$ - repo: https://github.com/pre-commit/pre-commit-hooks - rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c + rev: 3e8a8703264a2f4a69428a0aa4dcb512790b2c8c # frozen: v6.0.0 hooks: - id: check-json files: ^\.buildkite/.*\.json$ diff --git a/CMakeLists.txt b/CMakeLists.txt index e1c6169d92be..93d0818b6cc8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -31,6 +31,9 @@ endif() # CUDA by default, can be overridden by using -DVLLM_TARGET_DEVICE=... (used by setup.py) set(VLLM_TARGET_DEVICE "cuda" CACHE STRING "Target device backend for vLLM") +if (VLLM_TARGET_DEVICE STREQUAL "rocm" OR "$ENV{VLLM_KERNEL_SYMBOL_MAP}" STREQUAL "1") + set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +endif() message(STATUS "Build type: ${CMAKE_BUILD_TYPE}") message(STATUS "Target device: ${VLLM_TARGET_DEVICE}") @@ -365,6 +368,20 @@ if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_EXT_SRC "csrc/rocm/quickreduce/torch_bindings.cpp" "csrc/rocm/quickreduce/custom_quickreduce.cu") + # Octave KV cache (store, decode, rotation, prefill): only built for the + # validated architectures, gfx1100 (RDNA3) and gfx942 (MI300). + set(VLLM_OCTAVE_KV OFF) + if(VLLM_GPU_ARCHES MATCHES "gfx942|gfx1100") + set(VLLM_OCTAVE_KV ON) + list(APPEND VLLM_EXT_SRC + "csrc/attention/octave_torch_bindings.cpp" + "csrc/attention/octave_attn.cu" + "csrc/attention/octave_decode_wmma_rdna3.cu" + "csrc/attention/octave_prefill_rdna3.cu" + "csrc/attention/octave_prefill_cdna.cu" + "csrc/attention/octave_decode_mfma_cdna.cu" + ) + endif() message(STATUS "Enabling C extension.") define_extension_target( @@ -376,6 +393,13 @@ if(VLLM_GPU_LANG STREQUAL "HIP") ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + if(VLLM_OCTAVE_KV) + # MFMA accumulators in VGPRs: in AGPRs every online-softmax rescale of the + # output copies it through 64 extra VGPRs and halves the occupancy. + set_source_files_properties( + "${CMAKE_CURRENT_BINARY_DIR}/csrc/attention/octave_decode_mfma_cdna.hip" + PROPERTIES COMPILE_FLAGS "-mllvm -amdgpu-mfma-vgpr-form") + endif() endif() # _C HIP endif if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") @@ -406,6 +430,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") WORKING_DIRECTORY "${TORCH_SHADOW_INCLUDE_DIR}" COMMAND_ERROR_IS_FATAL ANY) + execute_process( + COMMAND ${Python_EXECUTABLE} -c + "import json,sys; from pathlib import Path; Path(sys.argv[1]).write_text(json.dumps({sys.argv[2]: sys.argv[3:]}))" + "${CMAKE_CURRENT_BINARY_DIR}/generated-source-deps.json" + "${TORCH_SHADOW_INCLUDE_DIR}/${_stable_hdr_rel}" + "${TORCH_INCLUDE_DIR}/${_stable_hdr_rel}" + "${CMAKE_CURRENT_LIST_DIR}/cmake/patches/pytorch_stable_string.patch" + COMMAND_ERROR_IS_FATAL ANY) + include_directories(BEFORE "${TORCH_SHADOW_INCLUDE_DIR}") endif() endif() @@ -1199,6 +1232,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") message(STATUS "Building hadacore") endif() + # Fused MQA logits + top-k prefill kernel (SM100 family). It compiles against + # DeepGEMM's headers, so DeepGEMM is fetched here. + include(cmake/external_projects/deepgemm.cmake) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + cuda_archs_loose_intersection(FUSED_MQA_TOPK_ARCHS "10.0f" "${CUDA_ARCHS}") + endif() + if(FUSED_MQA_TOPK_ARCHS) + set(FUSED_MQA_TOPK_SRC "csrc/fused_mqa_topk/fused_mqa_topk.cu") + set_gencode_flags_for_srcs( + SRCS "${FUSED_MQA_TOPK_SRC}" + CUDA_ARCHS "${FUSED_MQA_TOPK_ARCHS}") + set_property(SOURCE ${FUSED_MQA_TOPK_SRC} APPEND PROPERTY INCLUDE_DIRECTORIES + "${deepgemm_SOURCE_DIR}/deep_gemm/include" + "${deepgemm_SOURCE_DIR}/third-party/cutlass/include") + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_MQA_TOPK_SRC}") + message(STATUS + "Building fused MQA top-k for archs: ${FUSED_MQA_TOPK_ARCHS}") + endif() + # if CUDA endif endif() @@ -1263,6 +1315,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_KDA_DECODE=1) endif() + if(FUSED_MQA_TOPK_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_FUSED_MQA_TOPK=1) + endif() if(FUSED_GDN_DECODE_ARCHS) target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_GDN_DECODE=1) @@ -1339,6 +1395,19 @@ set_gencode_flags_for_srcs( SRCS "${VLLM_MOE_EXT_SRC}" CUDA_ARCHS "${CUDA_ARCHS}") +if(VLLM_GPU_LANG STREQUAL "CUDA") + set(MUSE_SPARK_ROUTING_ARCHS "${CUDA_ARCHS}") + if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0) + # Native FP32 warp-max reduction requires the Blackwell family target. + list(TRANSFORM MUSE_SPARK_ROUTING_ARCHS + REPLACE "^10[.]0([+]PTX)?$" "10.0f") + endif() + set_gencode_flags_for_srcs( + SRCS "csrc/moe/muse_spark_routing.cu" + CUDA_ARCHS "${MUSE_SPARK_ROUTING_ARCHS}") + list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/muse_spark_routing.cu") +endif() + if(VLLM_GPU_LANG STREQUAL "CUDA") set(VLLM_MOE_WNA16_SRC "csrc/moe/moe_wna16.cu") @@ -1618,7 +1687,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP") list(APPEND VLLM_ROCM_EXT_SRC "csrc/rocm/q_gemm_rdna3.cu" "csrc/rocm/q_gemm_rdna3_wmma.cu" - "csrc/rocm/moe_q_gemm_rdna3.cu") + "csrc/rocm/moe_q_gemm_rdna3.cu" + "csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu") endif() define_extension_target( @@ -1667,7 +1737,6 @@ endif() # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") - include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/deepselect.cmake) include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 943133252be0..c42ebe895b3c 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -32,6 +32,7 @@ import sys from dataclasses import replace from pathlib import Path +from typing import get_args import yaml from rich.console import Console @@ -50,6 +51,7 @@ is_mla_backend, ) +from vllm.config.cache import CacheDType from vllm.v1.worker.workspace import init_workspace_manager @@ -548,6 +550,11 @@ def main(): # Benchmark settings parser.add_argument("--device", default="cuda:0", help="Device") + parser.add_argument( + "--model", + default="meta-llama/Meta-Llama-3-8B", + help="Model whose config the backend is built with", + ) parser.add_argument( "--warmup-ms", type=int, @@ -561,8 +568,8 @@ def main(): parser.add_argument( "--kv-cache-dtype", default="auto", - choices=["auto", "fp8"], - help="KV cache dtype: auto or fp8", + choices=get_args(CacheDType), + help="KV cache dtype", ) parser.add_argument( "--cuda-graphs", @@ -894,6 +901,7 @@ def main(): device=args.device, profile_memory=args.profile_memory, kv_cache_dtype=args.kv_cache_dtype, + model=args.model, use_cuda_graphs=args.cuda_graphs, prefill_backend="fa4", ) @@ -964,6 +972,7 @@ def main(): device=args.device, profile_memory=args.profile_memory, kv_cache_dtype=args.kv_cache_dtype, + model=args.model, use_cuda_graphs=args.cuda_graphs, ncu_profile=args.ncu_profile, warmup_ms=args.warmup_ms, @@ -1036,6 +1045,8 @@ def main(): pbar.update(1) if args.ncu_profile: + if any(not result.success for result in all_results): + raise SystemExit(1) return # Display decode vs prefill results @@ -1208,6 +1219,7 @@ def main(): device=args.device, max_model_len=getattr(args, "max_model_len", None), kv_cache_dtype=args.kv_cache_dtype, + model=args.model, profile_memory=args.profile_memory, use_cuda_graphs=args.cuda_graphs, ncu_profile=args.ncu_profile, @@ -1353,6 +1365,7 @@ def main(): device=args.device, profile_memory=args.profile_memory, kv_cache_dtype=args.kv_cache_dtype, + model=args.model, use_cuda_graphs=args.cuda_graphs, ncu_profile=args.ncu_profile, warmup_ms=args.warmup_ms, @@ -1435,6 +1448,9 @@ def main(): if args.output_json: formatter.save_json(all_results, args.output_json) + if any(not result.success for result in all_results): + raise SystemExit(1) + if __name__ == "__main__": main() diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index b6fd26c80344..0af316d59a53 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -284,8 +284,10 @@ class BenchmarkConfig: torch_profile_iters: int = 3 warmup_ms: int | None = None - # "auto" or "fp8" + # Any CacheDType: "auto", "fp8", "int8_per_token_head", ... kv_cache_dtype: str = "auto" + # Model whose config the backend sees (dtype, rotary dims, ...) + model: str = "meta-llama/Meta-Llama-3-8B" # MLA-specific prefill_backend: str | None = None diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index a1b7d045050e..abb8858abe6d 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -36,6 +36,7 @@ set_current_vllm_config, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.attention.backends.utils import ( CommonAttentionMetadata, get_supported_kv_cache_layouts, @@ -48,6 +49,7 @@ compute_layer_kv_cache_shape_bytes, compute_layout_strides, create_kv_cache_views, + get_kv_quant_mode, ) # ============================================================================ @@ -142,15 +144,16 @@ def _build_common_attn_metadata( def _create_vllm_config( config: BenchmarkConfig, max_num_blocks: int, + max_model_len: int = 0, ) -> VllmConfig: """Create a VllmConfig for benchmarking with mock model methods.""" model_config = ModelConfig( - model="meta-llama/Meta-Llama-3-8B", - tokenizer="meta-llama/Meta-Llama-3-8B", + model=config.model, + tokenizer=config.model, trust_remote_code=False, dtype="auto", # Use model's native dtype seed=0, - max_model_len=1024, + max_model_len=max(1024, max_model_len), ) cache_config = CacheConfig( @@ -164,7 +167,7 @@ def _create_vllm_config( scheduler_config = SchedulerConfig( max_num_seqs=256, max_num_batched_tokens=8192, - max_model_len=8192, + max_model_len=max(8192, max_model_len), is_encoder_decoder=False, enable_chunked_prefill=True, ) @@ -342,6 +345,22 @@ def _create_input_tensors( return q_list, k_list, v_list +def _kv_cache_spec(config: BenchmarkConfig, dtype: torch.dtype) -> FullAttentionSpec: + if config.kv_cache_dtype == "auto": + cache_dtype = dtype + elif config.kv_cache_dtype.startswith("fp8"): + cache_dtype = current_platform.fp8_dtype() + else: + cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[config.kv_cache_dtype] + return FullAttentionSpec( + block_size=config.block_size, + num_kv_heads=config.num_kv_heads, + head_size=config.head_dim, + dtype=cache_dtype, + kv_quant_mode=get_kv_quant_mode(config.kv_cache_dtype), + ) + + def _create_kv_cache( config: BenchmarkConfig, max_num_blocks: int, @@ -351,16 +370,7 @@ def _create_kv_cache( layout: KVCacheLayout, ) -> list: """Create KV cache tensors for all layers using the standard allocator.""" - if config.kv_cache_dtype.startswith("fp8"): - cache_dtype = current_platform.fp8_dtype() - else: - cache_dtype = dtype - spec = FullAttentionSpec( - block_size=config.block_size, - num_kv_heads=config.num_kv_heads, - head_size=config.head_dim, - dtype=cache_dtype, - ) + spec = _kv_cache_spec(config, dtype) # Apply the backend's page customization, as the worker does for the real spec. spec = backend_class.customize_spec(spec) total_bytes = ( @@ -492,7 +502,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: # Suppress vLLM logs during setup to reduce spam with log_warnings_and_errors_only(): # Create vllm_config first - uses model's native dtype via "auto" - vllm_config = _create_vllm_config(config, max_num_blocks) + vllm_config = _create_vllm_config(config, max_num_blocks, max_kv) dtype = vllm_config.model_config.dtype # Wrap everything in set_current_vllm_config context @@ -510,12 +520,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: q_lens, kv_lens, config.block_size, device ) - kv_cache_spec = FullAttentionSpec( - block_size=config.block_size, - num_kv_heads=config.num_kv_heads, - head_size=config.head_dim, - dtype=dtype, - ) + kv_cache_spec = backend_class.customize_spec(_kv_cache_spec(config, dtype)) builder = _create_metadata_builder( backend_class, kv_cache_spec, vllm_config, device, config.backend diff --git a/benchmarks/kernels/benchmark_dsv41_mega_attn.py b/benchmarks/kernels/benchmark_dsv41_mega_attn.py index 1b3b5d9c3848..ede9021e0a36 100644 --- a/benchmarks/kernels/benchmark_dsv41_mega_attn.py +++ b/benchmarks/kernels/benchmark_dsv41_mega_attn.py @@ -118,7 +118,6 @@ def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads): n_wv_group = padded_heads // 8 cos_sin = make_cos_sin_cache(65536, device) positions = torch.randint(0, 65536, (s_q,), device=device, dtype=torch.int64) - pos32 = positions.to(torch.int32) # Per-step inputs: the wq_b output (local heads) and this step's KV row. q_local = torch.randn( @@ -177,7 +176,7 @@ def mega(): ex_len, False, 0.0, - pos32, + positions, False, 64, cos_sin, diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py index 3e47f61c74f0..9c06e3ba21f3 100644 --- a/benchmarks/kernels/benchmark_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py @@ -18,6 +18,11 @@ torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \ --model deepseek_v41 --backend flashinfer_cutedsl + +or the MiniMax-M3 MXFP8 attention O-proj and dense down-proj at TP4 with: + + torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \ + --model minimax_m3 --backend flashinfer_cutedsl """ import argparse @@ -45,10 +50,12 @@ # Per model: (hidden size N, global input widths K). Kimi-K3 lists the # shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the -# ``wo_b`` output projection (o_groups * o_lora_rank). +# ``wo_b`` output projection (o_groups * o_lora_rank). MiniMax-M3 lists the +# attention O-proj and the dense-layer down-proj. _MODEL_PROJECTIONS = { "kimi_k3": (7168, (6144, 12288)), "deepseek_v41": (5120, (8192,)), + "minimax_m3": (6144, (8192, 12288)), } diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index cc96ab291b02..3d6f52d4f184 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -123,6 +123,7 @@ def benchmark_config( num_iters: int = 100, block_quant_shape: list[int] = None, use_deep_gemm: bool = False, + per_act_token_quant: bool = False, ) -> float: init_dtype = torch.float16 if use_fp8_w8a8 else dtype x = torch.randn(num_tokens, hidden_size, dtype=dtype) @@ -224,12 +225,19 @@ def benchmark_config( torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32) * factor_for_scale ) + elif per_act_token_quant: + # Per-channel weights; activations are quantized per token at run time. + w1_scale = torch.rand( + (num_experts, shard_intermediate_size, 1), dtype=torch.float32 + ) + w2_scale = torch.rand((num_experts, hidden_size, 1), dtype=torch.float32) else: w1_scale = torch.randn(num_experts, dtype=torch.float32) w2_scale = torch.randn(num_experts, dtype=torch.float32) - a1_scale = torch.randn(1, dtype=torch.float32) - a2_scale = torch.randn(1, dtype=torch.float32) + if not per_act_token_quant: + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) w1 = w1.to(FP8_DTYPE) w2 = w2.to(FP8_DTYPE) @@ -255,6 +263,8 @@ def run(): w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, + per_act_token_quant=per_act_token_quant, + per_out_ch_quant=per_act_token_quant, block_shape=block_quant_shape, weight_dtype="int4" if use_int4_w4a16 else None, ) @@ -606,6 +616,7 @@ def benchmark( use_int4_w4a16: bool = False, block_quant_shape: list[int] = None, use_deep_gemm: bool = False, + per_act_token_quant: bool = False, ) -> tuple[dict[str, int], float]: # local import to allow serialization by ray @@ -649,6 +660,7 @@ def benchmark( num_iters=100, block_quant_shape=block_quant_shape, use_deep_gemm=use_deep_gemm, + per_act_token_quant=per_act_token_quant, ) return config, kernel_time @@ -666,6 +678,7 @@ def tune( search_space: list[dict[str, int]], block_quant_shape: list[int], use_deep_gemm: bool, + per_act_token_quant: bool = False, ) -> dict[str, int]: # local import to allow serialization by ray from vllm.platforms import current_platform @@ -713,6 +726,7 @@ def tune( num_iters=20, block_quant_shape=block_quant_shape, use_deep_gemm=use_deep_gemm, + per_act_token_quant=per_act_token_quant, ) except triton.runtime.autotuner.OutOfResources: # Some configurations may be invalid and fail to compile. @@ -819,6 +833,24 @@ def get_weight_block_size_safety(config, default_value=None): return default_value +def uses_per_token_act_quant(config) -> bool: + """Whether a compressed-tensors checkpoint quantizes activations per token. + + vLLM then runs the fused MoE with per-token activation and per-channel + weight scales, which is a different kernel specialization from per-tensor + FP8 but reads the same config file. + """ + quantization_config = getattr(config, "quantization_config", {}) + if not isinstance(quantization_config, dict): + return False + config_groups = quantization_config.get("config_groups", {}) + if len(config_groups) != 1: + return False + group = next(iter(config_groups.values())) + input_activations = group.get("input_activations") or {} + return input_activations.get("strategy") == "token" + + def get_model_params(config): architectures = getattr(config, "architectures", None) or [type(config).__name__] architecture = architectures[0] @@ -850,6 +882,7 @@ def get_model_params(config): hidden_size = config.hidden_size elif architecture in ( "BailingMoeV3ForCausalLM", + "BeamForCausalLM", "Qwen2MoeForCausalLM", "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", @@ -990,6 +1023,11 @@ def main(args: argparse.Namespace): use_int8_w8a16 = args.dtype == "int8_w8a16" use_int4_w4a16 = args.dtype == "int4_w4a16" block_quant_shape = get_weight_block_size_safety(config) + per_act_token_quant = ( + use_fp8_w8a8 and block_quant_shape is None and uses_per_token_act_quant(config) + ) + if per_act_token_quant: + print("Benchmarking per-token activation and per-channel weight FP8.") if use_int4_w4a16: group_size = get_quantization_group_size(config) if group_size is None: @@ -1092,6 +1130,7 @@ def _distribute(method: str, inputs: list[Any]) -> list[Any]: search_space, block_quant_shape, use_deep_gemm, + per_act_token_quant, ) for batch_size in batch_sizes ], @@ -1130,6 +1169,7 @@ def _distribute(method: str, inputs: list[Any]) -> list[Any]: use_int4_w4a16, block_quant_shape, use_deep_gemm, + per_act_token_quant, ) for batch_size in batch_sizes ], diff --git a/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py b/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py index 115459d4150e..2ec5788b8811 100644 --- a/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py +++ b/benchmarks/kernels/benchmark_muse_spark_qknorm_rope_kv_insert.py @@ -16,6 +16,16 @@ HEAD_DIM = 64 DEFAULT_TOKENS = [1, 8, 64, 256, 1024, 4096, 8192, 16384, 32768, 65536] +try: + from vllm.models.muse_spark.nvidia.ops.cute_dsl.qk_norm import ( + qk_norm, + supports_qk_norm, + ) + + HAS_CUTEDSL = supports_qk_norm(torch.bfloat16, HEAD_DIM) +except ImportError: + HAS_CUTEDSL = False + def make_cos_sin_cache(max_pos, rotary_dim, base, device): inv_freq = 1.0 / ( @@ -135,6 +145,12 @@ def unfused_rope(self, token_heads_per_warp: int = -1): ) self._insert() + def unfused_nope(self): + q, k, _ = self.views() + qk_norm(q, HEAD_DIM, self.eps) + qk_norm(k, HEAD_DIM, self.eps) + self._insert() + def fused_bytes(self): el = self.qkv.element_size() cache_el = self.kv_cache.element_size() @@ -260,7 +276,7 @@ def main(): f"dtype={args.dtype} kv_cache={args.kv_cache_dtype}\n" ) - modes = ["rope"] + modes = ["rope", "nope"] l2 = _l2_bytes(device) print( f"rows marked * have a working set inside the {l2 / 2**20:.0f} MiB L2, " @@ -274,6 +290,9 @@ def main(): for mode, use_graph in itertools.product(modes, [False, True]): if use_graph and args.no_graph: continue + if mode == "nope" and not HAS_CUTEDSL: + print("skipping the NoPE baseline: CuTe-DSL qk_norm unavailable") + continue label = "graph" if use_graph else "eager" print(f"── {mode} layer, {label} ──") print(header) @@ -282,7 +301,7 @@ def main(): case = Case(nt, args.nq, args.nkv, dtype, args.kv_cache_dtype, device) use_rope = mode == "rope" fused_us = timer(lambda c=case, r=use_rope: c.fused(r)) - baseline = case.unfused_rope + baseline = case.unfused_rope if use_rope else case.unfused_nope unfused_us = timer(baseline) tuned_us = ( min(timer(lambda c=case, t=t: c.unfused_rope(t)) for t in (4, 8)) diff --git a/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py b/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py new file mode 100644 index 000000000000..12f7571e93c8 --- /dev/null +++ b/benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py @@ -0,0 +1,538 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Benchmark MXFP4 MoE activation quant at GLM-5.3-shaped sizes. + +Sweeps expert count at a fixed hidden width (default 6144). Two timed regions +(``--timed-region``): + +* ``quant``: ``mxfp4_experts_quant`` / ``silu_and_mul_mxfp4_experts_quant`` only, + with synthetic routing offsets (kernel isolation). +* ``metadata_and_quant``: ``get_cutlass_moe_mm_data`` → ``shuffle_rows`` → quant. +* ``both``: run ``quant`` then ``metadata_and_quant``. + +Routing uses 32 rows (4 tokens, topk 8, two picks per expert 0..15). GEMMs are +outside the timed region. + +Run from a CUDA-built vLLM checkout, for example: + + .venv/bin/python benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py + .venv/bin/python benchmarks/kernels/benchmark_mxfp4_experts_activation_quant.py \\ + --timed-region both --experts 16,32,64,128,256 +""" + +from __future__ import annotations + +import argparse +import sys +from dataclasses import dataclass +from enum import Enum +from pathlib import Path + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +import vllm._custom_ops # noqa: E402, F401 -- registers torch.ops._C kernels +from vllm import _custom_ops as ops # noqa: E402 +from vllm.triton_utils import triton # noqa: E402 + +ROUTED_ROWS = 32 +TOPK = 8 +DEFAULT_HIDDEN_SIZE = 6144 +DEFAULT_MOE_INTERMEDIATE_SIZE = 2048 +DEFAULT_TP_SIZE = 4 +DEFAULT_EXPERT_COUNTS = (16, 32, 64, 128, 256) + + +class TimedRegion(str, Enum): + QUANT = "quant" + METADATA_AND_QUANT = "metadata_and_quant" + BOTH = "both" + + +@dataclass(frozen=True) +class Case: + kernel: str + output_width: int + fused_silu_mul: bool + experts: int + mm_n: int + mm_k: int + is_gated: bool + + +def routing_offsets( + num_experts: int, + rows: int = ROUTED_ROWS, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """32 rows, two assignments on experts 0..15 (production-like small decode).""" + if rows != ROUTED_ROWS: + raise ValueError(f"expected {ROUTED_ROWS} routed rows, got {rows}") + if num_experts < 16: + raise ValueError("expert sweep must use at least 16 experts") + + counts = [0] * num_experts + counts[:16] = [2] * 16 + + row_offsets = [0] + scale_offsets = [0] + expert_ids: list[int] = [] + for expert, count in enumerate(counts): + expert_ids.extend([expert] * count) + row_offsets.append(row_offsets[-1] + count) + scale_offsets.append(scale_offsets[-1] + ((count + 127) // 128) * 128) + assert row_offsets[-1] == rows + return ( + torch.tensor(row_offsets, dtype=torch.int32, device="cuda"), + torch.tensor(scale_offsets, dtype=torch.int32, device="cuda"), + torch.tensor(expert_ids, dtype=torch.int32, device="cuda"), + scale_offsets[-1], + ) + + +def decode_topk_ids(rows: int = ROUTED_ROWS) -> torch.Tensor: + """Two picks per expert 0..15 across 4 tokens (topk=8); shuffled (fixed seed).""" + if rows % TOPK != 0: + raise ValueError(f"routed rows {rows} must be divisible by topk {TOPK}") + expert_ids: list[int] = [] + for expert in range(16): + expert_ids.extend([expert, expert]) + assert len(expert_ids) == rows + gen = torch.Generator(device="cuda") + gen.manual_seed(0) + perm = torch.randperm(rows, generator=gen, device="cuda") + ordered = torch.tensor(expert_ids, device="cuda", dtype=torch.int32) + return ordered[perm].view(rows // TOPK, TOPK) + + +def benchmark_invoke(invoke, rep: int) -> float: + """Median latency in microseconds (CUDA graph replay via Triton).""" + quantiles = [0.5, 0.2, 0.8] + ms, _, _ = triton.testing.do_bench_cudagraph(invoke, rep=rep, quantiles=quantiles) + return ms * 1_000 + + +def _quant_tensors( + case: Case, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + scale_rows = int(blockscale_offsets[-1].item()) + scale_cols_int32 = (case.output_width // 32 + 3) // 4 + output = torch.empty( + (ROUTED_ROWS, case.output_width // 2), dtype=torch.uint8, device="cuda" + ) + output_scales = torch.zeros( + (scale_rows, scale_cols_int32), dtype=torch.int32, device="cuda" + ) + return output, output_scales + + +def benchmark_mxfp4_quant_only( + case: Case, + inputs: torch.Tensor, + rep: int, + sorted_first: bool, +) -> tuple[float, float]: + expert_offsets, blockscale_offsets, expert_ids, _ = routing_offsets(case.experts) + output, output_scales = _quant_tensors(case, expert_offsets, blockscale_offsets) + op = ( + torch.ops._C.silu_and_mul_mxfp4_experts_quant + if case.fused_silu_mul + else torch.ops._C.mxfp4_experts_quant + ) + + def invoke_offset() -> None: + op( + output, + output_scales, + inputs, + expert_offsets, + blockscale_offsets, + case.experts, + None, + ) + + def invoke_sorted() -> None: + op( + output, + output_scales, + inputs, + expert_offsets, + blockscale_offsets, + case.experts, + expert_ids, + ) + + if sorted_first: + direct_id_us = benchmark_invoke(invoke_sorted, rep) + baseline_us = benchmark_invoke(invoke_offset, rep) + else: + baseline_us = benchmark_invoke(invoke_offset, rep) + baseline_output = output.clone() + baseline_scales = output_scales.clone() + direct_id_us = benchmark_invoke(invoke_sorted, rep) + if not ( + torch.equal(output, baseline_output) + and torch.equal(output_scales, baseline_scales) + ): + raise AssertionError( + f"MXFP4 output mismatch for {case.kernel} E={case.experts}" + ) + return baseline_us, direct_id_us + + +def benchmark_mxfp4_metadata_and_quant( + case: Case, + inputs: torch.Tensor, + rep: int, + sorted_first: bool, +) -> tuple[float, float]: + num_experts = case.experts + topk_ids = decode_topk_ids() + expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") + blockscale_offsets = torch.empty_like(expert_offsets) + problem_sizes1 = torch.empty(num_experts, 3, dtype=torch.int32, device="cuda") + problem_sizes2 = torch.empty_like(problem_sizes1) + input_permutation = torch.empty(ROUTED_ROWS, dtype=torch.int32, device="cuda") + output_permutation = torch.empty_like(input_permutation) + sorted_expert_ids = torch.empty_like(input_permutation) + + torch.ops._C.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + case.mm_n, + case.mm_k, + blockscale_offsets, + case.is_gated, + sorted_expert_ids, + ) + output, output_scales = _quant_tensors(case, expert_offsets, blockscale_offsets) + op = ( + torch.ops._C.silu_and_mul_mxfp4_experts_quant + if case.fused_silu_mul + else torch.ops._C.mxfp4_experts_quant + ) + + def invoke_offset() -> None: + torch.ops._C.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + case.mm_n, + case.mm_k, + blockscale_offsets, + case.is_gated, + None, + ) + sorted_inputs = ops.shuffle_rows(inputs, input_permutation) + op( + output, + output_scales, + sorted_inputs, + expert_offsets, + blockscale_offsets, + num_experts, + None, + ) + + def invoke_sorted() -> None: + torch.ops._C.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + case.mm_n, + case.mm_k, + blockscale_offsets, + case.is_gated, + sorted_expert_ids, + ) + sorted_inputs = ops.shuffle_rows(inputs, input_permutation) + op( + output, + output_scales, + sorted_inputs, + expert_offsets, + blockscale_offsets, + num_experts, + sorted_expert_ids, + ) + + if sorted_first: + direct_id_us = benchmark_invoke(invoke_sorted, rep) + baseline_us = benchmark_invoke(invoke_offset, rep) + else: + baseline_us = benchmark_invoke(invoke_offset, rep) + baseline_output = output.clone() + baseline_scales = output_scales.clone() + direct_id_us = benchmark_invoke(invoke_sorted, rep) + if not ( + torch.equal(output, baseline_output) + and torch.equal(output_scales, baseline_scales) + ): + raise AssertionError( + f"MXFP4 output mismatch for {case.kernel} E={case.experts}" + ) + return baseline_us, direct_id_us + + +def verify_sorted_id_writer(num_experts: int) -> None: + """Check the routing sort writes the expected sorted expert ID at each row.""" + topk_ids = torch.arange(ROUTED_ROWS, device="cuda", dtype=torch.int32) % num_experts + topk_ids = topk_ids.view(4, 8) + expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") + blockscale_offsets = torch.empty_like(expert_offsets) + problem_sizes1 = torch.empty(num_experts, 3, dtype=torch.int32, device="cuda") + problem_sizes2 = torch.empty_like(problem_sizes1) + input_permutation = torch.empty(ROUTED_ROWS, dtype=torch.int32, device="cuda") + output_permutation = torch.empty_like(input_permutation) + sorted_expert_ids = torch.empty_like(input_permutation) + torch.ops._C.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + 64, + 64, + blockscale_offsets, + True, + sorted_expert_ids, + ) + expected = torch.sort(topk_ids.flatten()).values + torch.accelerator.synchronize() + if not torch.equal(sorted_expert_ids, expected): + raise AssertionError("routing sort produced incorrect sorted expert IDs") + + +def parse_expert_counts(raw: str) -> tuple[int, ...]: + counts = tuple(int(x.strip()) for x in raw.split(",") if x.strip()) + if not counts or any(c < 16 for c in counts): + raise ValueError( + "expert counts must be a comma-separated list of integers >= 16" + ) + return counts + + +def parse_timed_region(raw: str) -> TimedRegion: + try: + return TimedRegion(raw) + except ValueError: + choices = ", ".join(r.value for r in TimedRegion) + raise ValueError(f"timed-region must be one of: {choices}") from None + + +def intermediate_dims(moe_intermediate_size: int, tp_size: int) -> tuple[int, int]: + if moe_intermediate_size % tp_size != 0: + raise ValueError( + f"moe intermediate size {moe_intermediate_size} must be divisible " + f"by tp size {tp_size}" + ) + n_per_rank = moe_intermediate_size // tp_size + return 2 * n_per_rank, n_per_rank + + +def build_cases( + experts: tuple[int, ...], + hidden_size: int, + moe_intermediate_size: int, + tp_size: int, +) -> list[tuple[Case, torch.Tensor]]: + intermediate_in, intermediate_out = intermediate_dims( + moe_intermediate_size, tp_size + ) + cases: list[tuple[Case, torch.Tensor]] = [] + hidden_input = torch.randn( + (ROUTED_ROWS, hidden_size), device="cuda", dtype=torch.bfloat16 + ) + intermediate_input = torch.randn( + (ROUTED_ROWS, intermediate_in), device="cuda", dtype=torch.bfloat16 + ) + w1_n = 2 * intermediate_out + for expert_count in experts: + cases.append( + ( + Case( + "mxfp4_experts_quant", + hidden_size, + False, + expert_count, + mm_n=w1_n, + mm_k=hidden_size, + is_gated=True, + ), + hidden_input, + ) + ) + cases.append( + ( + Case( + "silu_and_mul_mxfp4_experts_quant", + intermediate_out, + True, + expert_count, + mm_n=hidden_size, + mm_k=intermediate_out, + is_gated=True, + ), + intermediate_input, + ) + ) + return cases + + +def run_region( + region: TimedRegion, + cases: list[tuple[Case, torch.Tensor]], + rep: int, + sorted_first: bool, +) -> None: + print(f"timed_region={region.value}") + print("kernel,experts,variant,median_us,speedup_vs_offset_search") + bench_fn = ( + benchmark_mxfp4_metadata_and_quant + if region == TimedRegion.METADATA_AND_QUANT + else benchmark_mxfp4_quant_only + ) + for case, inputs in cases: + baseline_us, direct_id_us = bench_fn(case, inputs, rep, sorted_first) + print(f"{case.kernel},{case.experts},offset_search,{baseline_us:.4f},1.0000") + print( + f"{case.kernel},{case.experts},sorted_expert_id," + f"{direct_id_us:.4f},{baseline_us / direct_id_us:.4f}" + ) + print() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--rep", + type=int, + default=100, + help="Repeat count (default: 100)", + ) + parser.add_argument( + "--hidden-size", + type=int, + default=DEFAULT_HIDDEN_SIZE, + help=f"hidden activation width K (default: {DEFAULT_HIDDEN_SIZE})", + ) + parser.add_argument( + "--tp-size", + type=int, + default=DEFAULT_TP_SIZE, + help=( + "tensor-parallel size; shards moe intermediate per rank " + f"(default: {DEFAULT_TP_SIZE})" + ), + ) + parser.add_argument( + "--moe-intermediate-size", + type=int, + default=DEFAULT_MOE_INTERMEDIATE_SIZE, + help=( + "total MoE intermediate size before TP; per-rank N is this divided " + f"by --tp-size (default: {DEFAULT_MOE_INTERMEDIATE_SIZE})" + ), + ) + parser.add_argument( + "--experts", + type=parse_expert_counts, + default=DEFAULT_EXPERT_COUNTS, + help=( + "comma-separated expert counts (default: " + f"{','.join(map(str, DEFAULT_EXPERT_COUNTS))})" + ), + ) + parser.add_argument( + "--timed-region", + type=parse_timed_region, + default=TimedRegion.QUANT, + help=( + "quant: activation quant only; metadata_and_quant: mm_data + shuffle + " + "quant; both: run each region" + ), + ) + parser.add_argument( + "--sorted-id-first", + action="store_true", + help="time sorted-expert-ID path before offset search (default: offset first)", + ) + parser.add_argument("--skip-metadata-check", action="store_true") + args = parser.parse_args() + if args.rep < 1: + parser.error("--rep must be positive") + if args.hidden_size < 128 or args.hidden_size % 128 != 0: + parser.error("--hidden-size must be a multiple of 128 and >= 128") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + if args.moe_intermediate_size < args.tp_size: + parser.error("--moe-intermediate-size must be >= --tp-size") + if args.moe_intermediate_size % args.tp_size != 0: + parser.error("--moe-intermediate-size must be divisible by --tp-size") + try: + intermediate_in, intermediate_out = intermediate_dims( + args.moe_intermediate_size, args.tp_size + ) + except ValueError as exc: + parser.error(str(exc)) + if intermediate_out < 32 or intermediate_out % 32 != 0: + parser.error( + "per-rank MoE intermediate (moe-intermediate-size / tp-size) " + "must be a multiple of 32 and >= 32" + ) + if not torch.cuda.is_available(): + parser.error("requires a CUDA GPU") + capability = torch.cuda.get_device_capability() + if capability[0] not in (10, 11): + parser.error(f"requires an SM10x/11x GPU, got SM{capability[0]}{capability[1]}") + + torch.manual_seed(0) + if not args.skip_metadata_check: + verify_sorted_id_writer(num_experts=256) + + cases = build_cases( + args.experts, + args.hidden_size, + args.moe_intermediate_size, + args.tp_size, + ) + print(f"GPU: {torch.cuda.get_device_name()} (SM{capability[0]}{capability[1]})") + print( + f"hidden={ROUTED_ROWS}x{args.hidden_size}, " + f"intermediate={ROUTED_ROWS}x{intermediate_in}->{intermediate_out} " + f"(moe_intermediate={args.moe_intermediate_size}, tp={args.tp_size}), " + f"topk={TOPK}" + ) + + regions: list[TimedRegion] + if args.timed_region == TimedRegion.BOTH: + regions = [TimedRegion.QUANT, TimedRegion.METADATA_AND_QUANT] + else: + regions = [args.timed_region] + + for region in regions: + run_region(region, cases, args.rep, args.sorted_id_first) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_mxfp8_tp4.py b/benchmarks/kernels/benchmark_mxfp8_tp4.py new file mode 100644 index 000000000000..e9ee03c8a431 --- /dev/null +++ b/benchmarks/kernels/benchmark_mxfp8_tp4.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate MXFP8 TP4 linears and time HIP graph replay over rotating weights. + +Example (record the baseline and candidate SHAs): + baseline=$(git rev-parse origin/main) + git rev-parse "$baseline" HEAD + git show "$baseline":vllm/model_executor/kernels/linear/mxfp8/rocm_native.py \ + > /tmp/baseline_rocm_native.py + .venv/bin/python benchmarks/kernels/benchmark_mxfp8_tp4.py \ + --baseline /tmp/baseline_rocm_native.py + +Reports cold (>400 MiB rotating weights) and warm timings separately. Each timing +is the median of 30 device-event samples; Python wall time is never ranked. +""" + +import argparse +import importlib.util +import json +import math +import statistics +from functools import partial +from pathlib import Path + +import torch + +from vllm.model_executor.kernels.linear.mxfp8 import rocm_native +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, +) +from vllm.triton_utils import triton + +SHAPES = [ + (m, n, k) + for n, k in [(2560, 6144), (6144, 2048), (1536, 6144), (6144, 768)] + for m in [1, 8, 16] +] + + +def time_graph(calls): + for call in calls: + call() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(2): + for call in calls: + call() + for _ in range(5): + graph.replay() + events = [ + (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(30) + ] + for start, end in events: + start.record() + for _ in range(5): + graph.replay() + end.record() + # Bound outstanding HIP graph nodes on short-K cases with many slots. + end.synchronize() + torch.accelerator.synchronize() + return statistics.median( + start.elapsed_time(end) / (10 * len(calls)) for start, end in events + ) + + +def launch(module, xq, xs, w, ws, out, config): + bm, bn, bk, nw, ns = config + m, k = xq.shape + n = w.shape[0] + return module._mxfp8_linear_kernel[(triton.cdiv(m, bm), triton.cdiv(n, bn))]( + xq, + xs, + w, + ws, + out, + m, + n, + k, + *xq.stride(), + *xs.stride(), + *w.stride(), + *ws.stride(), + *out.stride(), + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + num_warps=nw, + num_stages=ns, + ) + + +def emit(**data): + print(json.dumps(data), flush=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", type=Path) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--shape", type=int, choices=range(len(SHAPES))) + parser.add_argument("--tune", action="store_true") + parser.add_argument("--profile", action="store_true") + parser.add_argument("--asm-dir", type=Path) + args = parser.parse_args() + torch.accelerator.set_device_index(args.device) + torch.manual_seed(42) + variants = {"candidate": rocm_native} + if args.baseline: + spec = importlib.util.spec_from_file_location( + "vllm.model_executor.kernels.linear.mxfp8.baseline_rocm_native", + args.baseline, + ) + baseline = importlib.util.module_from_spec(spec) + spec.loader.exec_module(baseline) + variants = {"baseline": baseline, **variants} + emit( + torch=torch.__version__, + triton=triton.__version__, + device=str(torch.cuda.get_device_properties(args.device)), + ) + for index, (m, n, k) in enumerate(SHAPES): + if args.shape is not None and index != args.shape: + continue + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.5 + w, ws = mxfp8_e4m3_quantize( + torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.05 + ) + xq, xs = mxfp8_e4m3_quantize(x) + reference = ( + dequant_mxfp8_to_bf16(xq, xs).float() + @ dequant_mxfp8_to_bf16(w, ws).float().T + ) + out = torch.empty(m, n, device="cuda", dtype=x.dtype) + slots = math.ceil(400 * 1024**2 / (n * k)) + weights = [(w.clone(), ws.clone()) for _ in range(slots)] + calls = {} + for name, module in variants.items(): + configs = [module._select_cfg(m, n, k)] + if args.tune: + configs += [ + (16, 16, 1024, 2, 2), + (16, 16, 1024, 2, 3), + (16, 16, 1024, 2, 4), + (16, 16, 512, 2, 2), + (16, 32, 256, 4, 3), + (16, 16, 256, 2, 3), + (128, 128, 128, 8, 2), + ] + for config in dict.fromkeys(configs): + if k % config[2]: + continue + fn = partial(launch, module, xq, xs, config=config, out=out) + obj = fn(w=w, ws=ws) + rel = ((out.float() - reference).norm() / reference.norm()).item() + snr = -20 * math.log10(rel) if rel else float("inf") + assert rel < 0.005 and snr >= 30, (name, config, rel, snr) + emit( + shape=[m, n, k], + variant=name, + config=config, + rel=rel, + snr_db=snr, + registers=obj.n_regs, + spills=obj.n_spills, + lds=obj.metadata.shared, + ) + if args.asm_dir: + args.asm_dir.mkdir(parents=True, exist_ok=True) + (args.asm_dir / f"{name}-{m}-{n}-{k}-{config}.amdgcn").write_text( + obj.asm["amdgcn"] + ) + if args.profile: + for _ in range(10): + fn(w=w, ws=ws) + torch.accelerator.synchronize() + continue + calls[f"{name}-kernel-{config}"] = [ + partial(fn, w=wi, ws=si) for wi, si in weights + ] + if not args.profile and not args.tune: + wrapper = partial(module._mxfp8_dot_scaled_linear, x) + actual = wrapper(w, ws) + assert ((actual.float() - reference).norm() / reference.norm()) < 0.005 + calls[f"{name}-wrapper"] = [ + partial(wrapper, wi, si) for wi, si in weights + ] + for repeat in range(2): + order = list(calls.items()) + if repeat: + order.reverse() + for name, functions in order: + cold = time_graph(functions) + warm = time_graph([functions[0]] * 32) + emit( + shape=[m, n, k], + variant=name, + repeat=repeat, + cold_ms=cold, + warm_ms=warm, + slots=slots, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py index 84d9cc707c40..229f5385a4c2 100644 --- a/benchmarks/kernels/benchmark_reshape_and_cache_flash.py +++ b/benchmarks/kernels/benchmark_reshape_and_cache_flash.py @@ -22,10 +22,45 @@ @torch.inference_mode() +def _make_packed_kv_caches( + num_blocks: int, + block_size: int, + num_heads: int, + head_size: int, + v_head_size: int, + kv_cache_dtype: str, + dtype: torch.dtype, + kv_cache_layout: str, + device: str, +): + """K and V as views of one [.., head_size + v_head_size] buffer. + + The asymmetric kernel needs both sides to share a head stride, which is what + a packed K|V buffer gives (and what backends such as ROCm AITER unified + attention produce). + """ + if kv_cache_dtype not in ("auto", "bfloat16", "float16", "fp8"): + raise ValueError(f"asymmetric benchmark does not support {kv_cache_dtype}.") + content = head_size + v_head_size + if kv_cache_layout == "NHD": + physical = (num_blocks, block_size, num_heads, content) + permute = None + else: + physical = (num_blocks, num_heads, block_size, content) + permute = (0, 2, 1, 3) + + cache_dtype = torch.uint8 if kv_cache_dtype == "fp8" else dtype + packed = torch.zeros(physical, dtype=cache_dtype, device=device) + view = packed if permute is None else packed.permute(permute) + key_cache, value_cache = view.split([head_size, v_head_size], dim=-1) + return key_cache, value_cache + + def run_benchmark( num_tokens: int, num_heads: int, head_size: int, + v_head_size: int | None, block_size: int, num_blocks: int, dtype: torch.dtype, @@ -37,8 +72,14 @@ def run_benchmark( device: str = "cuda", ) -> float: """Return latency (seconds) for given num_tokens.""" - if kv_cache_dtype == "fp8" and head_size % 16: - raise ValueError("fp8 kv-cache requires head_size to be a multiple of 16.") + if v_head_size is None: + v_head_size = head_size + asymmetric = v_head_size != head_size + + if kv_cache_dtype == "fp8" and (head_size % 16 or v_head_size % 16): + raise ValueError( + "fp8 kv-cache requires head_size and v_head_size to be multiples of 16." + ) if implementation not in ("cuda", "triton"): raise ValueError( @@ -47,13 +88,23 @@ def run_benchmark( ) if implementation == "triton" and kv_cache_layout == "HND": return float("nan") # Triton does not support HND layout yet. + if asymmetric and implementation != "cuda": + raise ValueError( + "asymmetric head sizes (v_head_size != head_size) are only supported " + "by the 'cuda' implementation." + ) set_random_seed(42) torch.set_default_device(device) # create random key / value tensors [T, H, D]. key = torch.randn(num_tokens, num_heads, head_size, dtype=dtype, device=device) - value = torch.randn_like(key) + if asymmetric: + value = torch.randn( + num_tokens, num_heads, v_head_size, dtype=dtype, device=device + ) + else: + value = torch.randn_like(key) # prepare the slot mapping. # each token is assigned a unique slot in the KV-cache. @@ -63,20 +114,36 @@ def run_benchmark( slot_mapping_lst = random.sample(range(num_slots), num_tokens) slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device) - key_caches, value_caches = create_kv_caches_with_random_flash( - num_blocks, - block_size, - 1, # num_layers - num_heads, - head_size, - kv_cache_dtype, - dtype, - device=device, - cache_layout=kv_cache_layout, - ) - key_cache, value_cache = key_caches[0], value_caches[0] - # to free unused memory - del key_caches, value_caches + if asymmetric: + # K and V are views of one buffer, so both sides share the head stride; + # that is the only layout the asymmetric path accepts (see + # reshape_and_cache_flash). + key_cache, value_cache = _make_packed_kv_caches( + num_blocks, + block_size, + num_heads, + head_size, + v_head_size, + kv_cache_dtype, + dtype, + kv_cache_layout, + device, + ) + else: + key_caches, value_caches = create_kv_caches_with_random_flash( + num_blocks, + block_size, + 1, # num_layers + num_heads, + head_size, + kv_cache_dtype, + dtype, + device=device, + cache_layout=kv_cache_layout, + ) + key_cache, value_cache = key_caches[0], value_caches[0] + # to free unused memory + del key_caches, value_caches # compute per-kernel scaling factors for fp8 conversion (if used). k_scale = (key.amax() / 64.0).to(torch.float32) @@ -142,6 +209,7 @@ def main(args): num_tokens=n_tok, num_heads=args.num_heads, head_size=args.head_size, + v_head_size=args.v_head_size, block_size=args.block_size, num_blocks=args.num_blocks, dtype=STR_DTYPE_TO_TORCH_DTYPE[args.dtype], @@ -154,9 +222,14 @@ def main(args): ) rows.append([n_tok, layout, f"{lat * 1e6:.3f}"]) + v_note = ( + "" + if args.v_head_size is None + else f", head_size={args.head_size} v_head_size={args.v_head_size}" + ) print( f"Benchmark results for implementation {args.implementation}" - f" (measuring with {args.mode}):" + f" (measuring with {args.mode}{v_note}):" ) print(tabulate(rows, headers=["num_tokens", "layout", "latency (µs)"])) @@ -171,6 +244,13 @@ def main(args): choices=[64, 80, 96, 112, 120, 128, 192, 256], default=128, ) + parser.add_argument( + "--v-head-size", + type=int, + default=None, + help="V head dim; defaults to --head-size. Setting it differently " + "benchmarks the asymmetric K/V path (cuda implementation only).", + ) parser.add_argument("--block-size", type=int, choices=[16, 32], default=16) parser.add_argument("--num-blocks", type=int, default=128 * 512) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index d92ad58a9cf3..39b327cb76b7 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -492,6 +492,8 @@ set(VLLM_EXT_SRC "csrc/cpu/activation.cpp" "csrc/cpu/sgl-kernels/fla.cpp" "csrc/cpu/utils.cpp" + "csrc/cpu/block_table.cpp" + "csrc/cpu/input_batch.cpp" "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" @@ -693,6 +695,7 @@ else() USE_SABI 3 WITH_SOABI ) + target_compile_definitions(_C PRIVATE VLLM_STANDALONE_CPU_MRV2) if (VLLM_OPENBLAS_LIB) target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS) endif() diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index cde756127af7..e9d0bbbd2803 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,9 +32,12 @@ else() # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, # configurable SwiGLU alpha/beta, SiTU for FP8/FP4 Mega MoE, and NVFP4 # Mega MoE, plus swiglustep for BF16/FP8/NVFP4 Mega MoE and F8_128x4 - # SFs for the FP8 einsum output (vllm-project/DeepGEMM#40). + # SFs for the FP8 einsum output (vllm-project/DeepGEMM#40), plus + # fp4_fp4_mega_moe and fused post-expert RMSNorm for Muse Spark Open + # (vllm-project/DeepGEMM#39), and the MX MegaMoE alignment fix + # (vllm-project/DeepGEMM#43). # Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API. - set(_DEEPGEMM_UPSTREAM_TAG "727f847082d2e7946f70dc02b94949db206310d3") + set(_DEEPGEMM_UPSTREAM_TAG "5cdb145704ec4672ecf4b61441599521ce635df6") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index fadee58efbbd..0675e75e1d3f 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -13,7 +13,7 @@ else() FetchContent_Declare( flashkda GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git - GIT_TAG 17a037d98da546deb4591e967cf961a43c034d8b + GIT_TAG 47fc9fec6661d04aa9a42e9e5e069a51ad3b649b GIT_PROGRESS TRUE GIT_SUBMODULES cutlass ) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index d60fd84e914b..9a4c3ef895e7 100644 --- a/cmake/external_projects/flashmla.cmake +++ b/cmake/external_projects/flashmla.cmake @@ -19,7 +19,7 @@ else() FetchContent_Declare( flashmla GIT_REPOSITORY https://github.com/vllm-project/FlashMLA - GIT_TAG e2afa5fd8c220fbd05fe03fc24c9d5e64c24641b + GIT_TAG d1a1c92eff2cf558ddf49293828cd2bd21ca48c4 GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 6b84b35c9049..2336d66d23dc 100644 --- a/cmake/external_projects/fmha_sm100.cmake +++ b/cmake/external_projects/fmha_sm100.cmake @@ -17,7 +17,7 @@ else() FetchContent_Declare( fmha_sm100 GIT_REPOSITORY https://github.com/vllm-project/MSA.git - GIT_TAG 968061ab16cb14991a5c7d4c0e8eabf531be0dcf + GIT_TAG 25c144700f727b638067e98066a60507081b63ad GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index b6df33c170b8..222e3d83faa2 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -51,6 +51,13 @@ endif() if(NOT qutlass_SOURCE_DIR) message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.") endif() + +if(NOT EXISTS "${qutlass_SOURCE_DIR}/qutlass/csrc/include/registration.h") + message(FATAL_ERROR + "[QUTLASS] Stable ABI support was not found in the Qutlass source. " + "Please update your Qutlass package.") +endif() + message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) @@ -142,9 +149,9 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS) target_compile_definitions(_qutlass_C PRIVATE QUTLASS_MINIMAL_BUILD=1 - QUTLASS_DISABLE_PYBIND=1 TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC} CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1 + TORCH_TARGET_VERSION=0x020D000000000000ULL USE_CUDA) set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS diff --git a/cmake/hipify.py b/cmake/hipify.py index 2cf58a7b7d33..2135a625dfae 100755 --- a/cmake/hipify.py +++ b/cmake/hipify.py @@ -11,8 +11,10 @@ # import argparse +import json import os import shutil +from pathlib import Path from torch.utils.hipify.hipify_python import get_hip_file_path, hipify @@ -48,14 +50,29 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: parser.add_argument( "sources", help="Source files to hipify.", nargs="*", default=[] ) + parser.add_argument( + "--manifest", help="Write generated-to-original source/header provenance." + ) args = parser.parse_args() # Limit include scope to project_dir only includes = [os.path.join(args.project_dir, "*")] - # Get absolute path for all source files. - extra_files = [os.path.abspath(s) for s in args.sources] + project_dir = Path(args.project_dir).resolve() + output_dir = Path(args.output_dir).resolve() + provenance = { + str(output_dir / source.relative_to(project_dir)): str(source) + for source in project_dir.rglob("*") + if source.is_file() + } + source_copies = { + os.path.abspath(source): str( + output_dir / Path(source).resolve().relative_to(project_dir) + ) + for source in args.sources + } + extra_files = list(source_copies.values()) # Copy sources from project directory to output directory. # The directory might already exist to hold object files so we ignore that. @@ -75,16 +92,24 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: is_pytorch_extension=True, hipify_extra_files_only=True, ) + copied_sources = dict(provenance) + for source, result in hipify_result.items(): + if result.hipified_path is not None: + original = copied_sources.get( + os.path.abspath(source), os.path.abspath(source) + ) + provenance[os.path.abspath(result.hipified_path)] = original hipified_sources = [] for source in args.sources: s_abs = os.path.abspath(source) - if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None: - path = hipify_result[s_abs].hipified_path + copied = source_copies[s_abs] + if copied in hipify_result and hipify_result[copied].hipified_path is not None: + path = hipify_result[copied].hipified_path # PyTorch skips writing when is_pytorch_extension and text unchanged; # hipified_path then stays *.cu. CMake expects *.hip under output_dir. if s_abs.endswith(".cu") and path.endswith(".cu"): - dest = _expected_hip_build_path(s_abs, args.output_dir) + dest = _expected_hip_build_path(copied, args.output_dir) if os.path.normpath(path) != os.path.normpath(dest): os.makedirs(os.path.dirname(dest), exist_ok=True) shutil.copy2(path, dest) @@ -92,10 +117,19 @@ def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: else: hipified_s_abs = path else: - hipified_s_abs = s_abs + hipified_s_abs = copied hipified_sources.append(hipified_s_abs) + provenance[os.path.abspath(hipified_s_abs)] = s_abs assert len(hipified_sources) == len(args.sources) + manifest = ( + Path(args.manifest) if args.manifest else output_dir / "hipify-source-map.json" + ) + manifest.parent.mkdir(parents=True, exist_ok=True) + temporary = manifest.with_suffix(manifest.suffix + ".tmp") + temporary.write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n") + temporary.replace(manifest) + # Print hipified source files. print("\n".join(hipified_sources)) diff --git a/cmake/utils.cmake b/cmake/utils.cmake index 83eed340eb7d..7d6987122f79 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -113,11 +113,14 @@ function (vllm_finalize_hipify_target) list(REMOVE_DUPLICATES ALL_BYPRODUCTS) set(CSRC_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/csrc) + set(HIPIFY_MANIFEST ${CMAKE_CURRENT_BINARY_DIR}/hipify-source-map.json) add_custom_target( hipify_all - COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} ${ALL_SRCS} + COMMAND ${Python_EXECUTABLE} ${CMAKE_SOURCE_DIR}/cmake/hipify.py -p ${CMAKE_SOURCE_DIR}/csrc -o ${CSRC_BUILD_DIR} --manifest ${HIPIFY_MANIFEST} ${ALL_SRCS} DEPENDS ${CMAKE_SOURCE_DIR}/cmake/hipify.py ${ALL_SRCS} - BYPRODUCTS ${ALL_BYPRODUCTS} + BYPRODUCTS ${ALL_BYPRODUCTS} ${HIPIFY_MANIFEST} + WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} + VERBATIM COMMENT "Running hipify on all extension source files.") endfunction() diff --git a/csrc/attention/octave_arch.cuh b/csrc/attention/octave_arch.cuh new file mode 100644 index 000000000000..536f36b25943 --- /dev/null +++ b/csrc/attention/octave_arch.cuh @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Per-architecture primitives of the Octave kernels. Every AMD target vLLM +// builds compiles the portable tier; faster tiers are selected at runtime by +// the host code (see docs/design/rocm_octave.md, "Architecture contract"). +// +// int8 dot4: v_dot4_i32_iu8 on gfx11/gfx12, v_dot4_i32_i8 on gfx9/gfx10.3 +// decode, fast tier: octave_decode_wmma_rdna3.cu (WMMA, gfx11) +// prefill: octave_prefill_rdna3.cu (WMMA, gfx11) +// decode, portable: octave_attn.cu decode_kernel (dot4, every target) +// +// The portable kernels are written for 32-lane logical waves: shuffles never +// cross 32 lanes, so they are correct on wave64 (CDNA) as well. + +#pragma once + +#include + +namespace octave { + +// sum_i a.i8[i] * b.i8[i] + c, signed bytes on both sides. +__device__ __forceinline__ int dot4_i8(int a, int b, int c) { +#if defined(__GFX11__) || defined(__GFX12__) + return __builtin_amdgcn_sudot4(true, a, true, b, c, false); +#elif defined(__gfx942__) + return __builtin_amdgcn_sdot4(a, b, c, false); +#else + #pragma unroll + for (int i = 0; i < 4; ++i) + c += (int)(int8_t)(a >> (8 * i)) * (int)(int8_t)(b >> (8 * i)); + return c; +#endif +} + +} // namespace octave diff --git a/csrc/attention/octave_attn.cu b/csrc/attention/octave_attn.cu new file mode 100644 index 000000000000..b4056a1662a0 --- /dev/null +++ b/csrc/attention/octave_attn.cu @@ -0,0 +1,1180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave KV cache kernels, head_size 256 with 64 RoPE dims. Portable across +// the AMD targets vLLM builds (octave_arch.cuh); the WMMA decode tier for +// gfx11 lives in octave_decode_wmma_rdna3.cu. +// Format: vllm/v1/attention/ops/rocm_octave.py (the PyTorch reference there +// is the source of truth for the byte layout). +// +// octave_cache_store quantize + pack K/V into slots (one wave per slot). +// octave_decode split-KV attention per query group; K dot products +// with int8 dot4 on the codebook values of the packed +// codes, the query quantized to int8 once per block. +// Consecutive query tokens of one request (MTP verify) +// share the KV read. +// octave_rotate the store-time rotations, for the prefill path. +// +// Everything is in the rotated space: K and V are sign-flipped and +// Hadamard-rotated at store time; the query is rotated like K and the +// attention output is rotated back in the reduce kernel. + +#include +#include +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #include "octave_arch.cuh" + #include "octave_format.cuh" + +namespace octave { + +constexpr int DPL = D / 32; // dims per lane when a wave owns one slot + +template +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(half x) { + return __half2float(x); +} +template +__device__ __forceinline__ T from_f(float x) { + return (T)x; +} +template <> +__device__ __forceinline__ half from_f(float x) { + return __float2half(x); +} + +__device__ __forceinline__ float half_bits_to_f(uint32_t bits16) { + return __half2float(__ushort_as_half((unsigned short)bits16)); +} +__device__ __forceinline__ uint16_t f_to_half_bits(float x) { + return __half_as_ushort(__float2half(x)); +} + +template +__device__ __forceinline__ float xor_sum(float v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v += __shfl_xor(v, m); + return v; +} +template +__device__ __forceinline__ float xor_max(float v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v = fmaxf(v, __shfl_xor(v, m)); + return v; +} +template +__device__ __forceinline__ uint32_t xor_or(uint32_t v) { + #pragma unroll + for (int m = 1; m <= MASK; m <<= 1) v |= (uint32_t)__shfl_xor((int)v, m); + return v; +} + +// Walsh-Hadamard butterflies over DPL consecutive dims per lane, then across +// lanes with xor distance 1..LANE_MASK. Unnormalized. +template +__device__ __forceinline__ void fwht_wave(float (&x)[DPL], int lane) { + #pragma unroll + for (int h = 1; h < DPL; h <<= 1) { + #pragma unroll + for (int i = 0; i < DPL; ++i) { + if ((i & h) == 0) { + float a = x[i], b = x[i + h]; + x[i] = a + b; + x[i + h] = a - b; + } + } + } + #pragma unroll + for (int m = 1; m <= LANE_MASK; m <<= 1) { + const bool hi = lane & m; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + float o = __shfl_xor(x[i], m); + x[i] = hi ? o - x[i] : x[i] + o; + } + } +} + +__device__ __forceinline__ float sign_of(const int* bits, int i) { + return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f; +} + +// Rotate a lane's 8 dims like K (64-dim blocks) or like V (one 256-dim +// block). +__device__ __forceinline__ void rotate_k(float (&x)[DPL], int lane, + const int* signs) { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= sign_of(signs, lane * DPL + i); + fwht_wave<4>(x, lane); + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= 0.125f; +} + +__device__ __forceinline__ void rotate_full(float (&x)[DPL], int lane, + const int* signs) { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= sign_of(signs, lane * DPL + i); + fwht_wave<16>(x, lane); + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] *= 0.0625f; +} + +// Midpoints of the int8 codebooks (rocm_octave.LUT), in codebook units. +__device__ constexpr float kThr3[7] = {-103.f, -62.f, -29.5f, 0.f, + 29.5f, 62.f, 103.f}; +__device__ constexpr float kThr4[15] = {-110.5f, -84.5f, -66.f, -50.5f, -36.5f, + -24.f, -12.f, 0.f, 12.f, 24.f, + 36.5f, 50.5f, 66.f, 84.5f, 110.5f}; + +// Nearest codebook entry on the group RMS (4- or 3-bit per lane), codes in +// [0, 2^bits). Returns the stored scale, times LUT_ONE: the one that makes +// x . x_hat = |x|^2, so quantization does not shrink scores or outputs. +template +__device__ __forceinline__ float quantize_group(const float (&x)[DPL], + int (&codes)[DPL], + float group_dims, bool four) { + float ss = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) ss += x[i] * x[i]; + ss = xor_sum(ss); + const float to_lut = + (four ? 46.0f : 59.0f) / fmaxf(sqrtf(ss / group_dims), 1e-12f); + float num = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + const float y = x[i] * to_lut; + int c = 0; + #pragma unroll + for (int k = 0; k < 15; ++k) + c += (four ? y > kThr4[k] : (k < 7 && y > kThr3[k])) ? 1 : 0; + codes[i] = c; + num += x[i] * (float)(int8_t)(four ? lut4((uint32_t)c) : lut3((uint32_t)c)); + } + num = xor_sum(num); + return num > 0.f ? ss / num * LUT_ONE : 0.f; +} + +// 4-bit: one word per 8 dims (lane-local). 3-bit: this lane's contribution to +// the 2-bit plane word (16 dims, 2 lanes) and the 1-bit plane word (32 dims, +// 4 lanes); callers OR across lanes. +__device__ __forceinline__ uint32_t pack4(const int (&c)[DPL]) { + uint32_t w = 0; + #pragma unroll + for (int i = 0; i < DPL; ++i) + w |= (uint32_t)c[i] << (8 * (i & 3) + 4 * (i >> 2)); + return w; +} +__device__ __forceinline__ void pack3(const int (&c)[DPL], int idx, + uint32_t& lo, uint32_t& hi) { + lo = 0; + hi = 0; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + const int e16 = (idx & 1) * DPL + i; // dim within the 16-dim lo word + const int e32 = (idx & 3) * DPL + i; // dim within the 32-dim hi word + lo |= (uint32_t)(c[i] & 3) << (8 * (e16 & 3) + 2 * (e16 >> 2)); + hi |= (uint32_t)(c[i] >> 2) << (8 * (e32 & 3) + (e32 >> 2)); + } +} + +// --------------------------------------------------------------------------- +// Store: grid (num_tokens, num_kv_heads), block 32. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(32) + store_kernel(const KT* __restrict__ key, const KT* __restrict__ value, + uint8_t* __restrict__ cache, + const int64_t* __restrict__ slot_mapping, + const int* __restrict__ k_signs, + const int* __restrict__ v_signs, int block_size, int64_t sk0, + int64_t sk1, int64_t sv0, int64_t sv1, int64_t scb, + int64_t sch, int64_t sct) { + using F = Format; + const int t = blockIdx.x; + const int h = blockIdx.y; + const int lane = threadIdx.x; + const int64_t slot = slot_mapping[t]; + if (slot < 0) return; + uint8_t* dst = + cache + (slot / block_size) * scb + h * sch + (slot % block_size) * sct; + + float k[DPL], v[DPL]; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + k[i] = to_f(key[t * sk0 + h * sk1 + lane * DPL + i]); + v[i] = to_f(value[t * sv0 + h * sv1 + lane * DPL + i]); + } + + // K, block layout: rotate, then quantize per 64-dim block (8 lanes): the + // RoPE block at 4 bits, the NoPE blocks at 3. Compact layout: rotate over + // all dims like V, 3 bits, one scale. + const bool rope_lane = !KC && lane < R / DPL; + int kc[DPL]; + float k_scale; + if constexpr (KC) { + rotate_full(k, lane, k_signs); + k_scale = quantize_group<16>(k, kc, (float)D, false); + } else { + rotate_k(k, lane, k_signs); + k_scale = quantize_group<4>(k, kc, 64.f, rope_lane); + } + const int nl = lane - F::KR / DPL; // 3-bit lane index (valid when >= 0) + { + uint32_t lo, hi; + pack3(kc, nl & 3, lo, hi); + lo = xor_or<1>(lo); + hi = xor_or<2>(hi); + if (rope_lane) { + *reinterpret_cast(dst + 4 * lane) = pack4(kc); + } else { + if ((nl & 1) == 0) + *reinterpret_cast(dst + F::OFF_KN + 4 * (nl >> 1)) = lo; + if ((nl & 3) == 0) + *reinterpret_cast(dst + F::OFF_KN + F::KN / 4 + + 4 * (nl >> 2)) = hi; + } + } + + // V: rotate all 256 dims, one scale. + rotate_full(v, lane, v_signs); + int vc[DPL]; + const float v_scale = quantize_group<16>(v, vc, (float)D, VB == 4); + if constexpr (VB == 4) { + *reinterpret_cast(dst + F::OFF_V + 4 * lane) = pack4(vc); + } else { + uint32_t lo, hi; + pack3(vc, lane, lo, hi); + lo = xor_or<1>(lo); + hi = xor_or<2>(hi); + if ((lane & 1) == 0) + *reinterpret_cast(dst + F::OFF_V + 4 * (lane >> 1)) = lo; + if ((lane & 3) == 0) + *reinterpret_cast(dst + F::OFF_V + D / 4 + 4 * (lane >> 2)) = + hi; + } + + uint16_t* sc = reinterpret_cast(dst + F::OFF_SC); + if (KC ? lane == 0 : (lane & 7) == 0) sc[lane >> 3] = f_to_half_bits(k_scale); + if (lane == 0) sc[F::NUM_KSC] = f_to_half_bits(v_scale); +} + +// --------------------------------------------------------------------------- +// Decode stage 1: grid (num_q, num_kv_heads, num_splits), block 128. +// A block serves up to QG consecutive query tokens of one request (only the +// first token of each group of a run launches work) and all G query heads of +// its KV head: ROWS = QG * G attention rows share every K/V load. +// --------------------------------------------------------------------------- +constexpr int WAVES = 4; +constexpr int CHUNK = 32 * WAVES; + +template +__global__ void __launch_bounds__(128) + decode_kernel(const QT* __restrict__ Q, const uint8_t* __restrict__ cache, + const int* __restrict__ block_table, + const int* __restrict__ q_to_req, + const int* __restrict__ q_to_klen, float* __restrict__ mid_o, + const int* __restrict__ k_signs, float sm_scale_log2, + int num_q, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int num_splits, int64_t sq0, int64_t sq1, + int64_t scb, int64_t sch, int64_t sct, int64_t smo, + int64_t smh, int64_t sms) { + using F = Format; + constexpr int ROWS = QG * G; + constexpr int QW = D / 4; // int8 query words per row + + __shared__ int q8[ROWS][QW]; + __shared__ float q_bs[ROWS][NKB]; // query scale per K block + __shared__ float S[ROWS][CHUNK]; + __shared__ float m_s[ROWS], l_s[ROWS], alpha_s[ROWS]; + __shared__ const uint8_t* tok_ptr[CHUNK]; + __shared__ float tok_vs[CHUNK]; + + const int qi = blockIdx.x; + const int kvh = blockIdx.y; + const int si = blockIdx.z; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int wave = tid >> 5; + + int req = q_to_req[qi]; + const bool req_ok = req >= 0 && req < num_reqs; + int nq = 1; + if constexpr (QG > 1) { + if (req_ok) { + int s = qi; + while (s > 0 && qi - s < 256 && q_to_req[s - 1] == req) --s; + if ((qi - s) % QG != 0) return; // served by the group leader + while (nq < QG && qi + nq < num_q && q_to_req[qi + nq] == req) ++nq; + } + } + const int max_kv = max_blocks * block_size; + int klen[QG]; + int kv_max = 0; + #pragma unroll + for (int j = 0; j < QG; ++j) { + int kl = (req_ok && j < nq) ? q_to_klen[qi + j] : 0; + kl = min(max(kl, 0), max_kv); + klen[j] = kl; + kv_max = max(kv_max, kl); + } + if (!req_ok) req = 0; + + const int tps = (kv_max + num_splits - 1) / num_splits; + const int start = si * tps; + const int end = min(start + tps, kv_max); + + auto write_empty = [&]() { + for (int r = wave; r < ROWS; r += WAVES) { + const int j = r / G, g = r % G; + if (j >= nq) continue; + float* op = mid_o + (qi + j) * smo + (kvh * G + g) * smh + si * sms; + for (int d = lane; d < D; d += 32) op[d] = 0.f; + if (lane == 0) { + op[D] = -INFINITY; + op[D + 1] = 0.f; + } + } + }; + if (start >= end) { + write_empty(); + return; + } + + // Query: rotate like K, quantize to int8 with one scale per 64-dim block. + for (int r = wave; r < ROWS; r += WAVES) { + const int j = r / G, g = r % G; + float x[DPL]; + if (j < nq) { + const QT* qp = Q + (qi + j) * sq0 + (kvh * G + g) * sq1 + lane * DPL; + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] = to_f(qp[i]); + } else { + #pragma unroll + for (int i = 0; i < DPL; ++i) x[i] = 0.f; + } + if constexpr (KC) + rotate_full(x, lane, k_signs); + else + rotate_k(x, lane, k_signs); + float amax = 0.f; + #pragma unroll + for (int i = 0; i < DPL; ++i) amax = fmaxf(amax, fabsf(x[i])); + const float s = fmaxf(xor_max<4>(amax), 1e-20f) / 127.0f; + uint32_t w[2] = {0, 0}; + #pragma unroll + for (int i = 0; i < DPL; ++i) { + int q = __float2int_rn(x[i] / s); + q = min(max(q, -127), 127); + w[i >> 2] |= (uint32_t)(q & 0xFF) << (8 * (i & 3)); + } + q8[r][2 * lane] = (int)w[0]; + q8[r][2 * lane + 1] = (int)w[1]; + if ((lane & 7) == 0) q_bs[r][lane >> 3] = s; + } + __syncthreads(); + + if (tid < ROWS) { + m_s[tid] = -INFINITY; + l_s[tid] = 0.f; + } + float o0[ROWS], o1[ROWS]; // this lane's two V dims (rotated space) + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + o0[r] = 0.f; + o1[r] = 0.f; + } + const int my_d = wave * 64 + 2 * lane; // first of the two V dims owned + + for (int base = start; base < end; base += CHUNK) { + // ---- scores: lane = token ---- + const int tok = base + wave * 32 + lane; + const bool valid = tok < end; + if (valid) { + const int lb = tok / block_size; + int pb = block_table[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) pb = 0; + const uint8_t* p = + cache + pb * scb + kvh * sch + (int64_t)(tok - lb * block_size) * sct; + tok_ptr[wave * 32 + lane] = p; + + constexpr int KW = F::OFF_V / 4; // K code words + uint32_t kw[KW]; + const uint32_t* p4 = reinterpret_cast(p); + #pragma unroll + for (int w = 0; w < KW; ++w) kw[w] = p4[w]; + const uint32_t* kn = kw + F::OFF_KN / 4; + const uint32_t* ps = reinterpret_cast(p + F::OFF_SC); + float ks[NKB]; + if constexpr (KC) { + const uint32_t u = ps[0]; + #pragma unroll + for (int b = 0; b < NKB; ++b) ks[b] = half_bits_to_f(u & 0xFFFF); + tok_vs[wave * 32 + lane] = half_bits_to_f(u >> 16) * (1.0f / LUT_ONE); + } else { + const uint32_t u0 = ps[0], u1 = ps[1]; + ks[0] = half_bits_to_f(u0 & 0xFFFF); + ks[1] = half_bits_to_f(u0 >> 16); + ks[2] = half_bits_to_f(u1 & 0xFFFF); + ks[3] = half_bits_to_f(u1 >> 16); + tok_vs[wave * 32 + lane] = + half_bits_to_f(ps[2] & 0xFFFF) * (1.0f / LUT_ONE); + } + + // Rows two at a time: unrolling all of them kept too much live and + // spilled to scratch on gfx11. + #pragma unroll 2 + for (int r = 0; r < ROWS; ++r) { + const int* qr = q8[r]; + float sk = 0.f; + if constexpr (F::KR > 0) { + int dr = 0; + #pragma unroll + for (int u = 0; u < F::KR / 8; ++u) { + const uint32_t w = kw[u]; + dr = dot4_i8(qr[2 * u], (int)lut4(w & 0x0F0F0F0Fu), dr); + dr = dot4_i8(qr[2 * u + 1], (int)lut4((w >> 4) & 0x0F0F0F0Fu), dr); + } + sk = q_bs[r][0] * ks[0] * (float)dr; + } + #pragma unroll + for (int b = 0; b < F::KN / 64; ++b) { + constexpr int B0 = F::KR / 64; // first 3-bit block + int dn = 0; + #pragma unroll + for (int v = 0; v < 4; ++v) { // 2-bit plane word: 16 dims + const uint32_t lo = kn[4 * b + v]; + const uint32_t hi = kn[F::KN / 16 + 2 * b + (v >> 1)]; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t c = + ((lo >> (2 * i)) & 0x03030303u) | + (((hi >> (4 * (v & 1) + i)) & 0x01010101u) << 2); + dn = + dot4_i8(qr[F::KR / 4 + 16 * b + 4 * v + i], (int)lut3(c), dn); + } + } + sk += q_bs[r][B0 + b] * ks[B0 + b] * (float)dn; + } + const float sc = sk * (sm_scale_log2 / LUT_ONE); + int kl = klen[0]; // klen by constant index only (no scratch) + #pragma unroll + for (int j = 1; j < QG; ++j) kl = r / G == j ? klen[j] : kl; + S[r][wave * 32 + lane] = tok < kl ? sc : -INFINITY; + } + } else { + #pragma unroll + for (int r = 0; r < ROWS; ++r) S[r][wave * 32 + lane] = -INFINITY; + } + __syncthreads(); + + // ---- online softmax: each wave owns rows r = wave, wave + WAVES, ... ---- + for (int r = wave; r < ROWS; r += WAVES) { + float x[CHUNK / 32]; + float mx = -INFINITY; + #pragma unroll + for (int c = 0; c < CHUNK / 32; ++c) { + x[c] = S[r][32 * c + lane]; + mx = fmaxf(mx, x[c]); + } + mx = xor_max<16>(mx); + const float m_old = m_s[r]; + const float m_new = fmaxf(m_old, mx); + float ps = 0.f; + #pragma unroll + for (int c = 0; c < CHUNK / 32; ++c) { + const float p = x[c] == -INFINITY ? 0.f : exp2f(x[c] - m_new); + S[r][32 * c + lane] = p; + ps += p; + } + ps = xor_sum<16>(ps); + if (lane == 0) { + const float a = m_new == -INFINITY ? 1.f : exp2f(m_old - m_new); + alpha_s[r] = a; + l_s[r] = l_s[r] * a + ps; + m_s[r] = m_new; + } + } + __syncthreads(); + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float a = alpha_s[r]; + o0[r] *= a; + o1[r] *= a; + } + + // ---- P @ V: wave owns 64 dims, lane owns 2 ---- + const int n_valid = min(CHUNK, end - base); + for (int t = 0; t < n_valid; ++t) { + const uint8_t* p = tok_ptr[t]; + uint32_t c; // codes of dims my_d, my_d + 1 in bytes 0, 1 + if constexpr (VB == 4) { + const uint32_t w = + *reinterpret_cast(p + F::OFF_V + 4 * (my_d >> 3)); + const int e = my_d & 7; // even + c = ((w >> (8 * (e & 3) + 4 * (e >> 2))) & 15) | + (((w >> (8 * ((e + 1) & 3) + 4 * ((e + 1) >> 2))) & 15) << 8); + } else { + const uint32_t lo = + *reinterpret_cast(p + F::OFF_V + 4 * (my_d >> 4)); + const uint32_t hi = *reinterpret_cast( + p + F::OFF_V + D / 4 + 4 * (my_d >> 5)); + const int e16 = my_d & 15, e32 = my_d & 31; + const int f16 = e16 + 1, f32 = e32 + 1; + c = (((lo >> (8 * (e16 & 3) + 2 * (e16 >> 2))) & 3) | + (((hi >> (8 * (e32 & 3) + (e32 >> 2))) & 1) << 2)) | + ((((lo >> (8 * (f16 & 3) + 2 * (f16 >> 2))) & 3) | + (((hi >> (8 * (f32 & 3) + (f32 >> 2))) & 1) << 2)) + << 8); + } + const uint32_t val = lut(c); + const float vs = tok_vs[t]; + const float c0 = (float)(int8_t)(val & 0xFF) * vs; + const float c1 = (float)(int8_t)((val >> 8) & 0xFF) * vs; + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float pr = S[r][t]; + o0[r] += pr * c0; + o1[r] += pr * c1; + } + } + __syncthreads(); + } + + #pragma unroll + for (int r = 0; r < ROWS; ++r) { + const int j = r / G, g = r % G; + if (j >= nq) continue; + float* op = mid_o + (qi + j) * smo + (kvh * G + g) * smh + si * sms; + op[my_d] = o0[r]; + op[my_d + 1] = o1[r]; + if (tid == 0) { + op[D] = m_s[r]; + op[D + 1] = l_s[r]; + } + } +} + +// --------------------------------------------------------------------------- +// Decode stage 2: combine splits, rotate V back. grid (num_q, num_q_heads), +// block 256. Up to 1024 splits: their weights at once, then 8 waves take the +// splits round-robin (lane = 8 dims) and skip the ones that weigh zero. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(256) + reduce_kernel(const float* __restrict__ mid_o, OutT* __restrict__ out, + const int* __restrict__ v_signs, int num_splits, int64_t smo, + int64_t smh, int64_t sms, int64_t soo, int64_t soh) { + __shared__ float sw[1024], swm[8], swl[8]; + __shared__ float so[8][D]; + const int qi = blockIdx.x, hi = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 5, lane = tid & 31; + const float* base = mid_o + qi * smo + hi * smh; + + float ms[4], ls[4]; + float m = -INFINITY; + #pragma unroll + for (int k = 0; k < 4; ++k) { + const int s = tid + 256 * k; + ms[k] = s < num_splits ? base[s * sms + D] : -INFINITY; + ls[k] = s < num_splits ? base[s * sms + D + 1] : 0.0f; + m = fmaxf(m, ms[k]); + } + m = xor_max<16>(m); + if (lane == 0) swm[w] = m; + __syncthreads(); + float mg = swm[0]; + #pragma unroll + for (int i = 1; i < 8; ++i) mg = fmaxf(mg, swm[i]); + float l = 0.0f; + #pragma unroll + for (int k = 0; k < 4; ++k) { + const float a = ms[k] == -INFINITY ? 0.0f : exp2f(ms[k] - mg); + if (tid + 256 * k < 1024) sw[tid + 256 * k] = a; + l += ls[k] * a; + } + l = xor_sum<16>(l); + if (lane == 0) swl[w] = l; + __syncthreads(); + + float o[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int s = w; s < num_splits; s += 8) { + const float as = sw[s]; + if (as == 0.0f) continue; + const float4* v = (const float4*)(base + s * sms + lane * 8); + const float4 x = v[0], y = v[1]; + o[0] += x.x * as; + o[1] += x.y * as; + o[2] += x.z * as; + o[3] += x.w * as; + o[4] += y.x * as; + o[5] += y.y * as; + o[6] += y.z * as; + o[7] += y.w * as; + } + #pragma unroll + for (int d = 0; d < 8; ++d) so[w][lane * 8 + d] = o[d]; + __syncthreads(); + float lg = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) lg += swl[i]; + float acc = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) acc += so[i][tid]; + __syncthreads(); + // Inverse rotation of V over the 256 dims (so[0] reused as scratch). + float* x = so[0]; + x[tid] = acc; + __syncthreads(); + for (int hh = 1; hh < D; hh <<= 1) { + float a = 0.f, b = 0.f; + const bool lo = (tid & hh) == 0; + if (lo) { + a = x[tid]; + b = x[tid + hh]; + } + __syncthreads(); + if (lo) { + x[tid] = a + b; + x[tid + hh] = a - b; + } + __syncthreads(); + } + const float inv_l = lg > 0.f ? 1.0f / lg : 0.f; + out[qi * soo + hi * soh + tid] = + from_f(x[tid] * 0.0625f * sign_of(v_signs, tid) * inv_l); +} + +// --------------------------------------------------------------------------- +// In-place rotation of [T, H, 256] vectors for the prefill path, like K +// (signs, then 64-dim Hadamard blocks) or like V (one 256-dim block); +// `inverse` applies the transform then the signs. grid (T, H), block 32. +// --------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(32) + rotate_kernel(T* __restrict__ x, const int* __restrict__ signs, + bool k_layout, bool inverse, int64_t s0, int64_t s1) { + const int lane = threadIdx.x; + T* p = x + blockIdx.x * s0 + blockIdx.y * s1 + lane * DPL; + float v[DPL]; + #pragma unroll + for (int i = 0; i < DPL; ++i) v[i] = to_f(p[i]); + if (!inverse) { + #pragma unroll + for (int i = 0; i < DPL; ++i) v[i] *= sign_of(signs, lane * DPL + i); + } + float scale; + if (k_layout) { + fwht_wave<4>(v, lane); + scale = 0.125f; + } else { + fwht_wave<16>(v, lane); + scale = 0.0625f; + } + #pragma unroll + for (int i = 0; i < DPL; ++i) { + float y = v[i] * scale; + if (inverse) y *= sign_of(signs, lane * DPL + i); + p[i] = from_f(y); + } +} + +// --------------------------------------------------------------------------- +// Exact window (sparse decode). Direct-mapped fp16 shadow tables keyed by the +// cache slot: entry slot & mask holds that token's exact K and V ([hkv, 2D]) +// and is valid while tags[entry] == slot. The recent table takes every stored +// token, the sink table the first kSinks positions of a sequence. A lost +// entry only means the token is read from the packed cache instead. +// --------------------------------------------------------------------------- +constexpr int kWindow = 64; +constexpr int kSinks = 4; +constexpr int kMaxWindowHits = 128; + +// One wave per token. The per-entry lock keeps two tokens that share an +// entry in the same launch from leaving a mixed row behind. +template +__global__ void __launch_bounds__(32) + window_store_kernel(const T* __restrict__ key, const T* __restrict__ value, + const int64_t* __restrict__ slot_mapping, + const int* __restrict__ positions, T* __restrict__ wtab, + int* __restrict__ wtags, int* __restrict__ wlocks, + int wmask, T* __restrict__ stab, + int* __restrict__ stags, int* __restrict__ slocks, + int smask, int hkv, int64_t sk0, int64_t sk1, + int64_t sv0, int64_t sv1) { + const int i = blockIdx.x, lane = threadIdx.x; + const int64_t slot = slot_mapping[i]; + if (slot < 0) return; + const int pos = positions[i]; + auto put = [&](T* tab, int* tags, int* locks, int e) { + if (lane == 0) { + while (atomicCAS(&locks[e], 0, 1) != 0) { + } + atomicExch(&tags[e], -1); + } + __threadfence(); + for (int h = 0; h < hkv; ++h) { + T* row = tab + ((int64_t)e * hkv + h) * (2 * D); + #pragma unroll + for (int d = 0; d < DPL; ++d) { + row[lane * DPL + d] = key[i * sk0 + h * sk1 + lane * DPL + d]; + row[D + lane * DPL + d] = value[i * sv0 + h * sv1 + lane * DPL + d]; + } + } + __threadfence(); + if (lane == 0) { + atomicExch(&tags[e], (int)slot); + __threadfence(); + atomicExch(&locks[e], 0); + } + }; + put(wtab, wtags, wlocks, (int)(slot & wmask)); + if (pos >= 0 && pos < kSinks) put(stab, stags, slocks, (int)(slot & smask)); +} + +// Attention of one query row over its selected positions that sit in the +// exact window, written as one more split of mid_o (in the rotated V space, +// like the packed-cache splits). grid (num_q, hkv), block 256. +template +__global__ void __launch_bounds__(256) window_attn_kernel( + const T* __restrict__ query, const int* __restrict__ positions, + const int* __restrict__ q_to_req, const int* __restrict__ block_table, + const int* __restrict__ indices, const T* __restrict__ wtab, + const int* __restrict__ wtags, int wmask, const T* __restrict__ stab, + const int* __restrict__ stags, int smask, float* __restrict__ mid_o, + const int* __restrict__ v_signs, float sm_scale, int split, int topk, + int64_t sidx, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int hq, int hkv, int64_t sq0, int64_t sq1, int64_t smo, + int64_t smh, int64_t sms) { + __shared__ int s_entry[kMaxWindowHits]; + __shared__ float s_score[8][kMaxWindowHits]; + __shared__ int s_count; + const int qi = blockIdx.x, kvh = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 5, lane = tid & 31; + const int hpk = hq / hkv; + if (tid == 0) s_count = 0; + __syncthreads(); + const int req = q_to_req[qi]; + const int p = positions[qi]; + if (req >= 0 && req < num_reqs) { + for (int j = tid; j < topk; j += 256) { + const int t = indices[qi * sidx + j]; + const int d = p - t; + if (t < 0 || d < 0 || (d >= kWindow && t >= kSinks)) continue; + const int lb = t / block_size; + if (lb >= max_blocks) continue; + const int pb = block_table[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) continue; + const int slot = pb * block_size + (t - lb * block_size); + int e = -1; + if (d < kWindow && wtags[slot & wmask] == slot) + e = slot & wmask; + else if (t < kSinks && stags[slot & smask] == slot) + e = (slot & smask) | (1 << 30); + if (e < 0) continue; + const int k = atomicAdd(&s_count, 1); + if (k < kMaxWindowHits) s_entry[k] = e; + } + } + __syncthreads(); + const int n = min(s_count, kMaxWindowHits); + const float sl2 = sm_scale * 1.4426950408889634f; + for (int g = w; g < hpk; g += 8) { + const int head = kvh * hpk + g; + float q[DPL]; + #pragma unroll + for (int d = 0; d < DPL; ++d) + q[d] = to_f(query[qi * sq0 + head * sq1 + lane * DPL + d]); + auto row_of = [&](int e) { + const T* tab = (e >> 30) ? stab : wtab; + return tab + ((int64_t)(e & ((1 << 30) - 1)) * hkv + kvh) * (2 * D); + }; + float m = -INFINITY; + for (int i = 0; i < n; ++i) { + const T* r = row_of(s_entry[i]); + float acc = 0.f; + #pragma unroll + for (int d = 0; d < DPL; ++d) acc += q[d] * to_f(r[lane * DPL + d]); + acc = xor_sum<16>(acc) * sl2; + if (lane == 0) s_score[w][i] = acc; + m = fmaxf(m, acc); + } + __builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront"); + __builtin_amdgcn_wave_barrier(); + __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront"); + float l = 0.f, o[DPL] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int i = 0; i < n; ++i) { + const float pr = exp2f(s_score[w][i] - m); + l += pr; + const T* r = row_of(s_entry[i]) + D; + #pragma unroll + for (int d = 0; d < DPL; ++d) o[d] += pr * to_f(r[lane * DPL + d]); + } + rotate_full(o, lane, v_signs); + float* op = mid_o + qi * smo + head * smh + (int64_t)split * sms; + #pragma unroll + for (int d = 0; d < DPL; ++d) op[lane * DPL + d] = o[d]; + if (lane == 0) { + op[D] = n > 0 ? m : -INFINITY; + op[D + 1] = l; + } + } +} + +} // namespace octave + +// --------------------------------------------------------------------------- +// Host launchers +// --------------------------------------------------------------------------- +using namespace octave; + +// fmt: V fmt (3 or 4), plus kCompactFlag for compact K (3-bit V only). +int octave_slot_bytes(int64_t fmt) { + switch (fmt) { + case 3: + return Format<3>::SLOT; + case 4: + return Format<4>::SLOT; + case 3 + kCompactFlag: + return Format<3, true>::SLOT; + } + TORCH_CHECK(false, "octave: unknown format id ", fmt); + return 0; +} + +static void check_cache(const torch::Tensor& cache, int64_t fmt) { + TORCH_CHECK(cache.dtype() == at::kByte && cache.dim() == 4, + "octave: cache must be uint8 [blocks, heads, block_size, slot]"); + const int slot = octave_slot_bytes(fmt); + TORCH_CHECK(cache.size(3) == slot && cache.stride(3) == 1, + "octave: slot size mismatch, expected ", slot, " got ", + cache.size(3)); +} + +void octave_cache_store(torch::Tensor key, torch::Tensor value, + torch::Tensor cache, torch::Tensor slot_mapping, + torch::Tensor k_signs, torch::Tensor v_signs, + int64_t fmt) { + const int n = slot_mapping.size(0); + if (n == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(key.size(-1) == D && key.stride(-1) == 1 && value.stride(-1) == 1, + "octave: head_size must be 256 with contiguous last dim"); + TORCH_CHECK(slot_mapping.dtype() == at::kLong); + const int hkv = cache.size(1); + const at::cuda::OptionalCUDAGuard guard(device_of(key)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + dim3 grid(n, hkv); + #define SQ_STORE(B, KC, T) \ + store_kernel<<>>( \ + (const T*)key.data_ptr(), (const T*)value.data_ptr(), \ + (uint8_t*)cache.data_ptr(), slot_mapping.data_ptr(), \ + k_signs.data_ptr(), v_signs.data_ptr(), cache.size(2), \ + key.stride(0), key.stride(1), value.stride(0), value.stride(1), \ + cache.stride(0), cache.stride(1), cache.stride(2)) + const bool bf = key.dtype() == at::kBFloat16; + if (fmt == 4) { + if (bf) + SQ_STORE(4, false, __hip_bfloat16); + else + SQ_STORE(4, false, half); + } else if (fmt == 3) { + if (bf) + SQ_STORE(3, false, __hip_bfloat16); + else + SQ_STORE(3, false, half); + } else { + if (bf) + SQ_STORE(3, true, __hip_bfloat16); + else + SQ_STORE(3, true, half); + } + #undef SQ_STORE +} + +template +static void launch_decode(torch::Tensor& query, torch::Tensor& cache, + torch::Tensor& block_table, torch::Tensor& q_to_req, + torch::Tensor& q_to_klen, torch::Tensor& mid_o, + torch::Tensor& k_signs, float sm_scale_log2, int ns, + hipStream_t stream) { + const int num_q = query.size(0); + dim3 grid(num_q, cache.size(1), ns); + decode_kernel<<>>( + (const QT*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), + block_table.data_ptr(), q_to_req.data_ptr(), + q_to_klen.data_ptr(), mid_o.data_ptr(), + k_signs.data_ptr(), sm_scale_log2, num_q, cache.size(2), + block_table.size(1), block_table.size(0), cache.size(0), ns, + query.stride(0), query.stride(1), cache.stride(0), cache.stride(1), + cache.stride(2), mid_o.stride(0), mid_o.stride(1), mid_o.stride(2)); +} + +template +static void dispatch_decode(int g, int qg, torch::Tensor& query, + torch::Tensor& cache, torch::Tensor& block_table, + torch::Tensor& q_to_req, torch::Tensor& q_to_klen, + torch::Tensor& mid_o, torch::Tensor& k_signs, + float sl2, int ns, hipStream_t stream) { + #define SQ_DEC(GG, QQ) \ + if (g == GG && qg == QQ) { \ + launch_decode(query, cache, block_table, q_to_req, \ + q_to_klen, mid_o, k_signs, sl2, ns, \ + stream); \ + return; \ + } + SQ_DEC(1, 1); + SQ_DEC(1, 4); + SQ_DEC(2, 1); + SQ_DEC(2, 4); + SQ_DEC(4, 1); + SQ_DEC(4, 4); + SQ_DEC(6, 1); + SQ_DEC(6, 4); + SQ_DEC(8, 1); + SQ_DEC(8, 4); + #undef SQ_DEC + TORCH_CHECK(false, "octave_decode: unsupported GQA group ", g); +} + +int octave_decode_wmma(torch::Tensor query, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor q_to_req, + torch::Tensor q_to_klen, torch::Tensor mid_o, + torch::Tensor k_signs, double sm_scale, + int64_t num_kv_splits, int64_t fmt, + const std::optional& indices, + const int* qpos, const int* wtags, int wmask, + const int* stags, int smask, int64_t query_group); + +bool octave_reduce_mfma(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream); + +static void launch_reduce(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream) { + static const bool cdna3 = + std::string(at::cuda::getCurrentDeviceProperties()->gcnArchName) + .rfind("gfx942", 0) == 0; + if (cdna3 && octave_reduce_mfma(out, mid_o, v_signs, num_splits, stream)) + return; + dim3 grid(out.size(0), out.size(1)); + if (out.dtype() == at::kBFloat16) + reduce_kernel<__hip_bfloat16><<>>( + mid_o.data_ptr(), (__hip_bfloat16*)out.data_ptr(), + v_signs.data_ptr(), num_splits, mid_o.stride(0), mid_o.stride(1), + mid_o.stride(2), out.stride(0), out.stride(1)); + else + reduce_kernel<<>>( + mid_o.data_ptr(), (half*)out.data_ptr(), v_signs.data_ptr(), + num_splits, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), + out.stride(0), out.stride(1)); +} + +void octave_decode(torch::Tensor out, torch::Tensor query, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor q_to_req, + torch::Tensor q_to_klen, torch::Tensor mid_o, + torch::Tensor k_signs, torch::Tensor v_signs, + double sm_scale, int64_t num_kv_splits, int64_t fmt, + int64_t query_group, bool use_wmma) { + const int num_q = query.size(0); + if (num_q == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(query.size(2) == D && query.stride(2) == 1); + TORCH_CHECK(out.dtype() == query.dtype()); + TORCH_CHECK(mid_o.size(2) >= num_kv_splits && mid_o.size(3) >= D + 2); + TORCH_CHECK(num_kv_splits <= 1024 && mid_o.stride(3) == 1); + const int hq = query.size(1), hkv = cache.size(1); + TORCH_CHECK(hq % hkv == 0); + const int g = hq / hkv; + const int qg = query_group > 1 ? 4 : 1; + int ns = (int)num_kv_splits; + const float sl2 = (float)sm_scale * 1.4426950408889634f; + const at::cuda::OptionalCUDAGuard guard(device_of(query)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const bool bf = query.dtype() == at::kBFloat16; + // Fast tier when the device and shapes allow it, else the portable kernel. + const int wmma_splits = + use_wmma ? octave_decode_wmma(query, cache, block_table, q_to_req, + q_to_klen, mid_o, k_signs, sm_scale, + num_kv_splits, fmt, std::nullopt, nullptr, + nullptr, 0, nullptr, 0, qg) + : 0; + if (wmma_splits > 0) { + ns = wmma_splits; + } else { + #define SQ_DISPATCH(B, KC) \ + if (bf) \ + dispatch_decode(g, qg, query, cache, block_table, \ + q_to_req, q_to_klen, mid_o, \ + k_signs, sl2, ns, stream); \ + else \ + dispatch_decode(g, qg, query, cache, block_table, q_to_req, \ + q_to_klen, mid_o, k_signs, sl2, ns, stream) + if (fmt == 4) { + SQ_DISPATCH(4, false); + } else if (fmt == 3) { + SQ_DISPATCH(3, false); + } else { + SQ_DISPATCH(3, true); + } + #undef SQ_DISPATCH + } + launch_reduce(out, mid_o, v_signs, ns, stream); +} + +// Sparse decode: query row i attends to the logical positions indices[i, :] +// of request q_to_req[i] (negative entries are skipped). WMMA tier only. +void octave_decode_sparse(torch::Tensor out, torch::Tensor query, + torch::Tensor cache, torch::Tensor block_table, + torch::Tensor q_to_req, torch::Tensor indices, + torch::Tensor mid_o, torch::Tensor k_signs, + torch::Tensor v_signs, double sm_scale, + int64_t num_kv_splits, int64_t fmt, + const std::optional& positions, + const std::optional& wtab, + const std::optional& wtags, + const std::optional& stab, + const std::optional& stags) { + const int num_q = query.size(0); + if (num_q == 0) return; + check_cache(cache, fmt); + TORCH_CHECK(query.size(2) == D && query.stride(2) == 1); + TORCH_CHECK(out.dtype() == query.dtype()); + TORCH_CHECK(indices.dim() == 2 && indices.size(0) >= num_q && + indices.stride(1) == 1 && indices.dtype() == at::kInt); + TORCH_CHECK(mid_o.size(2) >= num_kv_splits && mid_o.size(3) >= D + 2); + TORCH_CHECK(num_kv_splits <= 1024 && mid_o.stride(3) == 1); + const at::cuda::OptionalCUDAGuard guard(device_of(query)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + // With the exact window, the packed-cache splits leave the last split of + // mid_o to octave_window_attn. + const bool win = positions.has_value(); + TORCH_CHECK(!win || (wtab && wtags && stab && stags && num_kv_splits >= 2)); + const int wmask = win ? (int)wtags->size(0) - 1 : 0; + const int smask = win ? (int)stags->size(0) - 1 : 0; + TORCH_CHECK( + !win || (((wmask + 1) & wmask) == 0 && ((smask + 1) & smask) == 0), + "octave_decode_sparse: shadow tables must be powers of two"); + const int ns = octave_decode_wmma( + query, cache, block_table, q_to_req, q_to_req, mid_o, k_signs, sm_scale, + win ? num_kv_splits - 1 : num_kv_splits, fmt, indices, + win ? positions->data_ptr() : nullptr, + win ? wtags->data_ptr() : nullptr, wmask, + win ? stags->data_ptr() : nullptr, smask, 4); + TORCH_CHECK(ns > 0, + "octave_decode_sparse: needs the WMMA tier (gfx1100, head size " + "256, GQA group <= 8)"); + if (!win) { + launch_reduce(out, mid_o, v_signs, ns, stream); + return; + } + const int hq = query.size(1), hkv = cache.size(1); + TORCH_CHECK(hq / hkv <= 8 && wtab->size(1) == hkv && wtab->size(2) == 2 * D); + dim3 grid(num_q, hkv); + #define SQ_WIN(T) \ + window_attn_kernel<<>>( \ + (const T*)query.data_ptr(), positions->data_ptr(), \ + q_to_req.data_ptr(), block_table.data_ptr(), \ + indices.data_ptr(), (const T*)wtab->data_ptr(), \ + wtags->data_ptr(), wmask, (const T*)stab->data_ptr(), \ + stags->data_ptr(), smask, mid_o.data_ptr(), \ + v_signs.data_ptr(), (float)sm_scale, ns, (int)indices.size(1), \ + indices.stride(0), cache.size(2), block_table.size(1), \ + block_table.size(0), cache.size(0), hq, hkv, query.stride(0), \ + query.stride(1), mid_o.stride(0), mid_o.stride(1), mid_o.stride(2)) + if (query.dtype() == at::kBFloat16) + SQ_WIN(__hip_bfloat16); + else + SQ_WIN(half); + #undef SQ_WIN + launch_reduce(out, mid_o, v_signs, ns + 1, stream); +} + +void octave_window_store(torch::Tensor key, torch::Tensor value, + torch::Tensor slot_mapping, torch::Tensor positions, + torch::Tensor wtab, torch::Tensor wtags, + torch::Tensor wlocks, torch::Tensor stab, + torch::Tensor stags, torch::Tensor slocks) { + const int n = slot_mapping.size(0); + if (n == 0) return; + TORCH_CHECK(slot_mapping.dtype() == at::kLong && + positions.dtype() == at::kInt); + TORCH_CHECK(wtab.dtype() == key.dtype() && stab.dtype() == key.dtype()); + const int hkv = wtab.size(1); + const int wmask = (int)wtags.size(0) - 1, smask = (int)stags.size(0) - 1; + const at::cuda::OptionalCUDAGuard guard(device_of(key)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + #define SQ_WSTORE(T) \ + window_store_kernel<<>>( \ + (const T*)key.data_ptr(), (const T*)value.data_ptr(), \ + slot_mapping.data_ptr(), positions.data_ptr(), \ + (T*)wtab.data_ptr(), wtags.data_ptr(), wlocks.data_ptr(), \ + wmask, (T*)stab.data_ptr(), stags.data_ptr(), \ + slocks.data_ptr(), smask, hkv, key.stride(0), key.stride(1), \ + value.stride(0), value.stride(1)) + if (key.dtype() == at::kBFloat16) + SQ_WSTORE(__hip_bfloat16); + else + SQ_WSTORE(half); + #undef SQ_WSTORE +} + +void octave_rotate(torch::Tensor x, torch::Tensor signs, bool k_layout, + bool inverse) { + if (x.numel() == 0) return; + TORCH_CHECK(x.dim() == 3 && x.size(2) == D && x.stride(2) == 1, + "octave_rotate: x must be [T, H, 256] with contiguous last dim"); + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + dim3 grid(x.size(0), x.size(1)); + if (x.dtype() == at::kBFloat16) + rotate_kernel<__hip_bfloat16><<>>( + (__hip_bfloat16*)x.data_ptr(), signs.data_ptr(), k_layout, inverse, + x.stride(0), x.stride(1)); + else + rotate_kernel + <<>>((half*)x.data_ptr(), signs.data_ptr(), + k_layout, inverse, x.stride(0), x.stride(1)); +} + +#else +void octave_cache_store(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, int64_t) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_decode(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, double, int64_t, int64_t, int64_t, bool) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_decode_sparse(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, double, + int64_t, int64_t, const std::optional&, + const std::optional&, + const std::optional&, + const std::optional&, + const std::optional&) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_window_store(torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor) { + TORCH_CHECK(false, "octave requires ROCm"); +} +void octave_rotate(torch::Tensor, torch::Tensor, bool, bool) { + TORCH_CHECK(false, "octave requires ROCm"); +} +#endif diff --git a/csrc/attention/octave_decode_mfma_cdna.cu b/csrc/attention/octave_decode_mfma_cdna.cu new file mode 100644 index 000000000000..07f899600c6f --- /dev/null +++ b/csrc/attention/octave_decode_mfma_cdna.cu @@ -0,0 +1,719 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave decode on MFMA for CDNA3 (gfx942). Same scheme as the WMMA decode +// (octave_decode_wmma_rdna3.cu): independent waves, one (split, 16-query row +// tile) per wave, transposed orientation, codebook bytes fed to fp16 as +// x + 1152 and debiased afterwards. With 16x16x16 MFMA on wave64 (lane l +// holds column l % 16 and rows 4 (l / 16) .. + 3 of the result): +// +// QK S^T[token][q] = K x Q^T. A lane owns one token (l % 16) and, in the +// four k-steps of each 64-dim block, the 16 contiguous dims +// 16 (l / 16) .. + 15 of that block, so its codes are one contiguous +// run of the slot. One accumulator per K scale. +// PV O^T[dim][q] = V^T x P^T. Output row r of fragment c is dim 16 r + c, +// so a lane reads the 16 contiguous dims 16 (l % 16) .. + 15 of tokens +// 4 (l / 16) .. + 3. P^T is the S^T accumulator, one query per lane. + +#include +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #include "octave_format.cuh" + + #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx942__) + #define OCTAVE_MFMA_STUB + #endif + +namespace octave_mfma { + +typedef _Float16 h4v __attribute__((ext_vector_type(4))); +typedef _Float16 h8v __attribute__((ext_vector_type(8))); +typedef _Float16 h2v __attribute__((ext_vector_type(2))); +typedef float f4v __attribute__((ext_vector_type(4))); +typedef int i4v __attribute__((ext_vector_type(4))); +typedef const int __attribute__((address_space(4))) cint; + +using octave::D; +using octave::Format; +using octave::LUT_ONE; +[[maybe_unused]] constexpr int QROW = D + 8; +[[maybe_unused]] constexpr float PSCALE = 256.0f; +[[maybe_unused]] constexpr float BIAS = 1152.0f; // byte x fed as x + 1152 +[[maybe_unused]] constexpr uint32_t X = 0x80808080u; +// Bytes of a 16-token tile (slot <= 256). +[[maybe_unused]] constexpr int TILE_MAX = 4096; + +// Bytes b of a dword -> fp16 1024 + b, exact. +__device__ __forceinline__ h4v bytes_to_h4(uint32_t w) { + const uint32_t lo = __builtin_amdgcn_perm(0x64646464u, w, 0x07010700u); + const uint32_t hi = __builtin_amdgcn_perm(0x64646464u, w, 0x07030702u); + return __builtin_bit_cast(h4v, (uint2){lo, hi}); +} + +template +__device__ __forceinline__ uint32_t lut4_fast(uint32_t codes) { + constexpr uint32_t B = BIASED ? X : 0u; + const uint32_t lo = __builtin_amdgcn_perm( + octave::LUT4_1 ^ B, octave::LUT4_0 ^ B, codes & 0x07070707u); + const uint32_t hi = __builtin_amdgcn_perm( + octave::LUT4_3 ^ B, octave::LUT4_2 ^ B, codes & 0x07070707u); + const uint32_t t = codes & 0x08080808u; + const uint32_t m = (t << 5) - (t >> 3); + return (hi & m) | (lo & ~m); +} + +// 4 dwords of 4 tokens each (same 4 dims) -> 4 dwords of 4 dims each. +__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2, + uint32_t a3, uint32_t o[4]) { + const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u); + const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u); + const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u); + const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u); + o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u); + o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u); + o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u); + o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u); +} + +template +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(half x) { + return __half2float(x); +} + +__device__ __forceinline__ float sign_of(const int* bits, int i) { + return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f; +} +__device__ __forceinline__ float lo16(uint32_t x) { + return __half2float(__ushort_as_half((unsigned short)(x & 0xFFFF))); +} +__device__ __forceinline__ float hi16(uint32_t x) { + return __half2float(__ushort_as_half((unsigned short)(x >> 16))); +} +// Over the four lanes of a query (l, l ^ 16, l ^ 32, l ^ 48). +__device__ __forceinline__ float qmax4(float v) { + v = fmaxf(v, __shfl_xor(v, 16)); + return fmaxf(v, __shfl_xor(v, 32)); +} +__device__ __forceinline__ float qsum4(float v) { + v += __shfl_xor(v, 16); + return v + __shfl_xor(v, 32); +} + +// QG: query tokens per block (4 for multi-token verification, else 1). +template +__global__ __launch_bounds__(64 * NW) void decode_mfma( + const QT* __restrict__ Q, const uint8_t* __restrict__ cache, + const int* __restrict__ block_table, const int* __restrict__ q_to_req, + const int* __restrict__ q_to_klen, float* __restrict__ mid_o, + const int* __restrict__ k_signs, float sm_scale, int num_q, int num_q_heads, + int num_kv_heads, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int num_splits, int min_tps, int num_row_tiles, + int64_t sq0, int64_t sq1, int64_t scb, int64_t sch, int64_t sct, + int64_t smo, int64_t smh, int64_t sms) { + #ifndef OCTAVE_MFMA_STUB + using F = Format; + constexpr int NG = KC ? 1 : 4; // QK accumulators (one per K scale) + const int grp = blockIdx.x; + const int kvh = blockIdx.y; + const int tid = threadIdx.x; + const int w = __builtin_amdgcn_readfirstlane(tid >> 6), lane = tid & 63; + const int col = lane & 15, gq = lane >> 4; + const int rt = num_row_tiles == 2 ? (w & 1) : 0; + const int splits_per_block = NW / num_row_tiles; + const int si = + blockIdx.z * splits_per_block + (num_row_tiles == 2 ? w >> 1 : w); + const int hpk = num_q_heads / num_kv_heads; + const int max_kv = max_blocks * block_size; + const float sm_log2 = sm_scale * 1.4426950408889634f / LUT_ONE; + + // The staged queries, then (once they are in registers) one 4 KB tile + // buffer per wave. + constexpr int SMEM = + 32 * QROW * 2 > NW * TILE_MAX ? 32 * QROW * 2 : NW * TILE_MAX; + __shared__ __attribute__((aligned(16))) uint8_t smem[SMEM]; + auto sQ = reinterpret_cast<_Float16 (*)[QROW]>(smem); + + // Stage the group's queries rotated like K: each 32-lane half owns one + // query vector, 8 dims per lane. + { + const int l32 = lane & 31; + for (int v = 2 * w + (lane >> 5); v < 32; v += 2 * NW) { + const int r = v / hpk, g = v % hpk; + const int qi = grp * QG + r; + float x[8]; + const bool live = v < QG * hpk && qi < num_q; + #pragma unroll + for (int i = 0; i < 8; ++i) + x[i] = live ? to_f(Q[(int64_t)qi * sq0 + + (int64_t)(kvh * hpk + g) * sq1 + l32 * 8 + i]) + : 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) x[i] *= sign_of(k_signs, l32 * 8 + i); + #pragma unroll + for (int hh = 1; hh < 8; hh <<= 1) { + #pragma unroll + for (int i = 0; i < 8; ++i) + if ((i & hh) == 0) { + const float a = x[i], b = x[i + hh]; + x[i] = a + b; + x[i + hh] = a - b; + } + } + #pragma unroll + for (int m = 1; m <= (KC ? 16 : 4); m <<= 1) { + const bool hi = l32 & m; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const float o = __shfl_xor(x[i], m); + x[i] = hi ? o - x[i] : x[i] + o; + } + } + #pragma unroll + for (int i = 0; i < 8; ++i) + sQ[v][l32 * 8 + i] = (_Float16)(x[i] * (KC ? 0.0625f : 0.125f)); + } + } + __syncthreads(); + + int rreq[QG], rlen[QG]; + #pragma unroll + for (int r = 0; r < QG; ++r) { + const int qi = grp * QG + r; + int rq = -1, ln = 0; + if (qi < num_q) { + rq = q_to_req[qi]; + ln = q_to_klen[qi]; + if (rq < 0 || rq >= num_reqs) { + rq = 0; + ln = 0; + } + ln = min(max(ln, 0), max_kv); + } + rreq[r] = __builtin_amdgcn_readfirstlane(rq); + rlen[r] = __builtin_amdgcn_readfirstlane(ln); + } + + // This lane's query (column) as int8 B fragments for the i8 MFMA: step + // 2 b + h holds dims 64 b + 16 gq + 8 h .. + 7 (the dot product does not + // care how dims map to k, only that K uses the same map). Two int8 halves, + // q ~ qsc (hi + lo / 256), keep ~16 bits of the query. + const int qv = rt * 16 + col; + long Qh[8], Ql[8]; + float qsc; + { + float x[64]; + float qmax = 1e-20f; + #pragma unroll + for (int i = 0; i < 64; ++i) { + x[i] = (float)sQ[qv][64 * (i >> 4) + 16 * gq + (i & 15)]; + qmax = fmaxf(qmax, fabsf(x[i])); + } + qmax = qmax4(qmax); + const float inv = 127.0f / qmax; + qsc = qmax * (1.0f / 127.0f); + #pragma unroll + for (int st = 0; st < 8; ++st) { + uint64_t wh = 0, wl = 0; + #pragma unroll + for (int e = 0; e < 8; ++e) { + const float y = x[8 * st + e] * inv; + const float h = rintf(y); + const int lo = max(-127, min(127, (int)rintf((y - h) * 256.0f))); + wh |= (uint64_t)(uint8_t)(int8_t)(int)h << (8 * e); + wl |= (uint64_t)(uint8_t)(int8_t)lo << (8 * e); + } + Qh[st] = (long)wh; + Ql[st] = (long)wl; + } + } + const float sm_q = sm_log2 * qsc * (1.0f / 256.0f); + __syncthreads(); + uint8_t* lt = smem + w * TILE_MAX; + + // rreq/rlen are only ever indexed by constants (unrolled selections): a + // dynamic index would put them in scratch memory. + auto pick = [&](const int* a, int i) { + int x = a[0]; + #pragma unroll + for (int r = 1; r < QG; ++r) x = i == r ? a[r] : x; + return x; + }; + int seg0 = 0; + while (seg0 < QG && pick(rreq, seg0) >= 0) { + const int req = pick(rreq, seg0); + int seg1 = seg0 + 1; + while (seg1 < QG && pick(rreq, seg1) == req) ++seg1; + int Lmax = 0; + #pragma unroll + for (int r = 0; r < QG; ++r) + if (r >= seg0 && r < seg1) Lmax = max(Lmax, rlen[r]); + const int qv_lo = seg0 * hpk, qv_hi = seg1 * hpk; + int tps = (Lmax + num_splits - 1) / num_splits; + if (tps < min_tps) tps = min_tps; + tps = (tps + 15) & ~15; + + const bool tile_live = qv_hi > rt * 16 && qv_lo < rt * 16 + 16; + const bool qlive = qv >= qv_lo && qv < qv_hi; + if (blockIdx.z * splits_per_block * tps >= Lmax) { + if (tile_live && si < num_splits && gq == 0 && qlive) { + float* op = mid_o + (int64_t)(grp * QG + qv / hpk) * smo + + (int64_t)(kvh * hpk + qv % hpk) * smh + (int64_t)si * sms; + op[D] = -INFINITY; + op[D + 1] = 0.0f; + } + seg0 = seg1; + continue; + } + + if (tile_live && si < num_splits) { + int qlen = 0; + #pragma unroll + for (int r = 0; r < QG; ++r) + if (qlive && qv / hpk == r) qlen = rlen[r]; + + const int start = si * tps; + const int end = min(start + tps, Lmax); + float m = -INFINITY, l = 0.0f, cp = 0.0f; + f4v O[16]; + #pragma unroll + for (int c = 0; c < 16; ++c) O[c] = (f4v){0, 0, 0, 0}; + + const uint8_t* cache_h = cache + (int64_t)kvh * sch; + auto tile_ptr = [&](int pos) { + const int lb = pos / block_size; + int pb = ((cint*)block_table)[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) pb = 0; + return cache_h + (int64_t)pb * scb + + (int64_t)(pos - lb * block_size) * sct; + }; + + // Raw K of token col: per 64-dim block, the codes of dims + // 16 gq .. + 15. 4-bit: 8 bytes. 3-bit: 2-bit plane word gq and the + // 16 bits of the 1-bit plane. Scales of tokens 4 gq .. + 3. + constexpr int NB3 = KC ? 4 : 3; // 3-bit blocks + uint2 kr = {0, 0}; + uint2 kn[NB3]; + uint32_t sa[4], sb[4], sv[4]; + // A 16-token tile is one contiguous run of a block: each lane fetches + // four 16-byte chunks of the next tile into registers, then stores + // them to its wave's LDS buffer, where the fragments are read. + constexpr int tb = 16 * F::SLOT; + uint4 raw[TILE_MAX / 1024], raw2[TILE_MAX / 1024]; + auto fetch = [&](uint4* dst, const uint8_t* tile) { + #pragma unroll + for (int j = 0; j < TILE_MAX / 1024; ++j) { + const int o = 16 * (lane + 64 * j); + dst[j] = o < tb ? *(const uint4*)(tile + o) : make_uint4(0, 0, 0, 0); + } + }; + auto stash = [&]() { + #pragma unroll + for (int j = 0; j < TILE_MAX / 1024; ++j) { + const int o = 16 * (lane + 64 * j); + if (o < tb) *(uint4*)(lt + o) = raw[j]; + } + }; + auto load_k = [&](const uint8_t* tile) { + const uint8_t* pk = tile + (int64_t)col * F::SLOT; + if constexpr (!KC) kr = *(const uint2*)(pk + 8 * gq); + #pragma unroll + for (int b = 0; b < NB3; ++b) { + kn[b].x = *(const uint32_t*)(pk + F::OFF_KN + 16 * b + 4 * gq); + kn[b].y = *(const uint32_t*)(pk + F::OFF_KN + F::KN / 4 + 8 * b + + 4 * (gq >> 1)) >> + (4 * (gq & 1)); + } + #pragma unroll + for (int e = 0; e < 4; ++e) { + const uint8_t* ps = + tile + (int64_t)(4 * gq + e) * F::SLOT + F::OFF_SC; + if constexpr (KC) { + sa[e] = *(const uint32_t*)ps; + sv[e] = sa[e] >> 16; + } else { + const uint2 u = *(const uint2*)ps; + sa[e] = u.x; + sb[e] = u.y; + sv[e] = *(const uint16_t*)(ps + 8); + } + } + }; + // Raw V of tokens 4 gq .. + 3, dims 16 col .. + 15. + uint2 vr[4]; + auto load_v = [&](const uint8_t* tile) { + #pragma unroll + for (int e = 0; e < 4; ++e) { + const uint8_t* p = tile + (int64_t)(4 * gq + e) * F::SLOT + F::OFF_V; + if constexpr (VB == 4) { + vr[e] = *(const uint2*)(p + 8 * col); + } else { + vr[e].x = *(const uint32_t*)(p + 4 * col); + vr[e].y = *(const uint32_t*)(p + D / 4 + 4 * (col >> 1)) >> + (4 * (col & 1)); + } + } + }; + // Codebook bytes of dims 4 s .. + 3 of a 16-dim run: int8 for K, biased + // (x + 128, the bias folded into the tables) for V. + auto codes4 = [&](uint2 raw, int s) -> uint32_t { + const uint32_t wd = s < 2 ? raw.x : raw.y; + return (wd >> (4 * (s & 1))) & 0x0F0F0F0Fu; + }; + auto codes3 = [&](uint2 raw, int s) -> uint32_t { + return ((raw.x >> (2 * s)) & 0x03030303u) | + (((raw.y >> s) & 0x01010101u) << 2); + }; + auto k_bytes = [&](int b, int s) -> uint32_t { + return (!KC && b == 0) ? lut4_fast(codes4(kr, s)) + : octave::lut3(codes3(kn[KC ? b : b - 1], s)); + }; + auto v_bytes = [&](uint2 raw, int s) -> uint32_t { + if constexpr (VB == 4) + return lut4_fast(codes4(raw, s)); + else + return __builtin_amdgcn_perm(octave::LUT3_HI ^ X, octave::LUT3_LO ^ X, + codes3(raw, s)); + }; + + if (start < end) fetch(raw, tile_ptr(start)); + if (start + 16 < end) fetch(raw2, tile_ptr(start + 16)); + + for (int base = start; base < end; base += 16) { + stash(); + load_k(lt); + #pragma unroll + for (int j = 0; j < TILE_MAX / 1024; ++j) raw[j] = raw2[j]; + if (base + 32 < end) fetch(raw2, tile_ptr(base + 32)); + + i4v Sh[NG], Sl[NG]; + #pragma unroll + for (int g = 0; g < NG; ++g) Sh[g] = Sl[g] = (i4v){0, 0, 0, 0}; + #pragma unroll + for (int b = 0; b < 4; ++b) { + const int g = KC ? 0 : b; + #pragma unroll + for (int h = 0; h < 2; ++h) { + const long a = (long)(((uint64_t)k_bytes(b, 2 * h + 1) << 32) | + k_bytes(b, 2 * h)); + Sh[g] = __builtin_amdgcn_mfma_i32_16x16x32_i8(a, Qh[2 * b + h], + Sh[g], 0, 0, 0); + Sl[g] = __builtin_amdgcn_mfma_i32_16x16x32_i8(a, Ql[2 * b + h], + Sl[g], 0, 0, 0); + } + } + auto dot = [&](int g, int e) { + return (float)(Sh[g][e] * 256 + Sl[g][e]); + }; + + load_v(lt); + + float sc[4]; + float mloc = -INFINITY; + #pragma unroll + for (int e = 0; e < 4; ++e) { + const int t = base + 4 * gq + e; + const bool ok = t < end && t < qlen; + float v; + if constexpr (KC) { + v = lo16(sa[e]) * dot(0, e); + } else { + v = lo16(sa[e]) * dot(0, e) + hi16(sa[e]) * dot(1, e) + + lo16(sb[e]) * dot(2, e) + hi16(sb[e]) * dot(3, e); + } + sc[e] = ok ? v * sm_q : -INFINITY; + mloc = fmaxf(mloc, sc[e]); + } + const float mn = fmaxf(m, qmax4(mloc)); + float alpha = 1.0f; + if (mn > m) + alpha = (m == -INFINITY) ? 0.0f : __builtin_amdgcn_exp2f(m - mn); + m = mn; + h4v pb; + float psum = 0.0f, lsum = 0.0f; + #pragma unroll + for (int e = 0; e < 4; ++e) { + const float p = + sc[e] == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(sc[e] - mn); + lsum += p; + const _Float16 hp = (_Float16)(p * lo16(sv[e]) * PSCALE); + psum += (float)hp; + pb[e] = hp; + } + l = l * alpha + lsum; + cp = cp * alpha + psum; + if (__ballot(alpha != 1.0f)) { + #pragma unroll + for (int c = 0; c < 16; ++c) O[c] *= alpha; + } + + // A of fragment 4 s + d: dim 16 col + 4 s + d of tokens 4 gq .. + 3. + // One perm per token pair transposes the biased bytes and widens + // them to fp16 [b, 0x00] (selector 0x0C), i.e. the exact denormal + // b * 2^-24. + #pragma unroll + for (int s = 0; s < 4; ++s) { + uint32_t a[4]; + #pragma unroll + for (int e = 0; e < 4; ++e) a[e] = v_bytes(vr[e], s); + #pragma unroll + for (int d = 0; d < 4; ++d) { + constexpr uint32_t kSel = 0x0C000C00u; + const uint32_t sel = kSel | ((4u + d) << 16) | (uint32_t)d; + const h4v A = __builtin_bit_cast( + h4v, (uint2){__builtin_amdgcn_perm(a[1], a[0], sel), + __builtin_amdgcn_perm(a[3], a[2], sel)}); + O[4 * s + d] = __builtin_amdgcn_mfma_f32_16x16x16f16( + A, pb, O[4 * s + d], 0, 0, 0); + } + } + } + + l = qsum4(l); + cp = qsum4(cp); + if (qlive) { + float* op = mid_o + (int64_t)(grp * QG + qv / hpk) * smo + + (int64_t)(kvh * hpk + qv % hpk) * smh + (int64_t)si * sms; + const bool any = m != -INFINITY; + // O holds sum hp (x + 128) 2^-24. + const float corr = 128.0f * 0x1p-24f * cp; + constexpr float kOut = 0x1p24f / (PSCALE * LUT_ONE); + #pragma unroll + for (int e = 0; e < 4; ++e) { + float* dp = op + 16 * (4 * gq + e); + #pragma unroll + for (int c = 0; c < 16; c += 4) { + float4 v; + v.x = any ? (O[c][e] - corr) * kOut : 0.0f; + v.y = any ? (O[c + 1][e] - corr) * kOut : 0.0f; + v.z = any ? (O[c + 2][e] - corr) * kOut : 0.0f; + v.w = any ? (O[c + 3][e] - corr) * kOut : 0.0f; + *(float4*)(dp + c) = v; + } + } + if (gq == 0) { + op[D] = m; + op[D + 1] = l; + } + } + } + seg0 = seg1; + } + #endif // OCTAVE_MFMA_STUB +} + +// Split combine for gfx942: grid (num_q, num_q_heads), 16 waves. Thread t +// weighs split t; wave w sums splits w, w + 16, ... four at a time over +// dims 4 lane .. + 3; wave 0 adds the 16 partials and undoes the 256-dim V +// rotation in registers (two in-lane stages, six shuffle stages). +template +__global__ __launch_bounds__(1024) void reduce_mfma( + const float* __restrict__ mid_o, OutT* __restrict__ out, + const int* __restrict__ v_signs, int num_splits, int64_t smo, int64_t smh, + int64_t sms, int64_t soo, int64_t soh) { + #ifndef OCTAVE_MFMA_STUB + __shared__ float sw[1024], swm[16], swl[16]; + __shared__ __attribute__((aligned(16))) float so[16][D]; + const int qi = blockIdx.x, hi = blockIdx.y; + const int tid = threadIdx.x, w = tid >> 6, lane = tid & 63; + const float* base = mid_o + qi * smo + hi * smh; + + const float ms = tid < num_splits ? base[tid * sms + D] : -INFINITY; + const float ls = tid < num_splits ? base[tid * sms + D + 1] : 0.0f; + float m = ms; + #pragma unroll + for (int k = 1; k < 64; k <<= 1) m = fmaxf(m, __shfl_xor(m, k)); + if (lane == 0) swm[w] = m; + __syncthreads(); + float mg = swm[0]; + #pragma unroll + for (int i = 1; i < 16; ++i) mg = fmaxf(mg, swm[i]); + const float a = ms == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(ms - mg); + sw[tid] = a; + float l = ls * a; + #pragma unroll + for (int k = 1; k < 64; k <<= 1) l += __shfl_xor(l, k); + if (lane == 0) swl[w] = l; + __syncthreads(); + + float4 o = {0, 0, 0, 0}; + for (int s0 = w; s0 < num_splits; s0 += 64) { + float4 x[4]; + float as[4]; + #pragma unroll + for (int u = 0; u < 4; ++u) { + const int s = s0 + 16 * u; + as[u] = s < num_splits ? sw[s] : 0.0f; + x[u] = s < num_splits ? *(const float4*)(base + s * sms + 4 * lane) + : make_float4(0, 0, 0, 0); + } + #pragma unroll + for (int u = 0; u < 4; ++u) { + // Splits without tokens weigh zero; their O was never written. + o.x += as[u] == 0.0f ? 0.0f : x[u].x * as[u]; + o.y += as[u] == 0.0f ? 0.0f : x[u].y * as[u]; + o.z += as[u] == 0.0f ? 0.0f : x[u].z * as[u]; + o.w += as[u] == 0.0f ? 0.0f : x[u].w * as[u]; + } + } + *(float4*)&so[w][4 * lane] = o; + __syncthreads(); + if (w != 0) return; + float x[4] = {0, 0, 0, 0}; + #pragma unroll + for (int i = 0; i < 16; ++i) { + const float4 v = *(const float4*)&so[i][4 * lane]; + x[0] += v.x; + x[1] += v.y; + x[2] += v.z; + x[3] += v.w; + } + float lg = 0.0f; + #pragma unroll + for (int i = 0; i < 16; ++i) lg += swl[i]; + { + const float a0 = x[0] + x[1], a1 = x[0] - x[1]; + const float a2 = x[2] + x[3], a3 = x[2] - x[3]; + x[0] = a0 + a2; + x[1] = a1 + a3; + x[2] = a0 - a2; + x[3] = a1 - a3; + } + #pragma unroll + for (int mk = 1; mk < 64; mk <<= 1) { + const bool hi_lane = lane & mk; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const float p = __shfl_xor(x[i], mk); + x[i] = hi_lane ? p - x[i] : x[i] + p; + } + } + const float inv_l = lg > 0.f ? 0.0625f / lg : 0.f; + OutT* op = out + qi * soo + hi * soh + 4 * lane; + #pragma unroll + for (int i = 0; i < 4; ++i) + op[i] = (OutT)(x[i] * sign_of(v_signs, 4 * lane + i) * inv_l); + #endif // OCTAVE_MFMA_STUB +} + +} // namespace octave_mfma + +// Split combine on gfx942; false when the shape is not covered. +bool octave_reduce_mfma(torch::Tensor out, torch::Tensor mid_o, + torch::Tensor v_signs, int num_splits, + cudaStream_t stream) { + using namespace octave_mfma; + // 1024 threads per (query, head) pay off while those blocks alone cannot + // fill the GPU; big batches keep the 256-thread kernel. + static const int max_blocks = + 2 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + if (num_splits > 1024 || out.size(0) * out.size(1) > max_blocks || + out.size(2) != D || out.stride(2) != 1 || mid_o.stride(3) != 1) + return false; + dim3 grid(out.size(0), out.size(1)); + if (out.dtype() == at::kBFloat16) + reduce_mfma<__hip_bfloat16><<>>( + mid_o.data_ptr(), (__hip_bfloat16*)out.data_ptr(), + v_signs.data_ptr(), num_splits, mid_o.stride(0), mid_o.stride(1), + mid_o.stride(2), out.stride(0), out.stride(1)); + else + reduce_mfma<<>>( + mid_o.data_ptr(), (half*)out.data_ptr(), v_signs.data_ptr(), + num_splits, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), + out.stride(0), out.stride(1)); + return true; +} + +constexpr int kMfmaSplits = 256; +constexpr int kMfmaMinTps = 32; +constexpr int kMfmaNw = 4; + +// Launches the MFMA decode when the shape is covered; returns the number of +// splits it wrote, or 0 when the caller must use its own kernel. Dense only. +int octave_decode_mfma(torch::Tensor query, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor q_to_req, + torch::Tensor q_to_klen, torch::Tensor mid_o, + torch::Tensor k_signs, double sm_scale, + int64_t num_kv_splits, int64_t fmt, + int64_t query_group) { + using namespace octave_mfma; + const int num_q = query.size(0); + const int num_q_heads = query.size(1); + const int num_kv_heads = cache.size(1); + const int hpk = num_q_heads / num_kv_heads; + const int QG = query_group > 1 ? 4 : 1; + const int slot = fmt == 4 ? octave::Format<4, false>::SLOT + : fmt == 3 ? octave::Format<3, false>::SLOT + : octave::Format<3, true>::SLOT; + const bool eligible = + query.size(2) == D && query.stride(2) == 1 && hpk * QG <= 32 && + cache.size(2) % 16 == 0 && + cache.stride(2) % (fmt & octave::kCompactFlag ? 4 : 8) == 0 && + cache.stride(2) == slot && cache.stride(1) % 16 == 0 && + cache.stride(0) % 16 == 0 && + reinterpret_cast(cache.data_ptr()) % 16 == 0 && + num_kv_splits >= 1; + if (!eligible) return 0; + const int nrt = std::min(num_q, QG) * hpk > 16 ? 2 : 1; + // Enough splits for ~two resident waves per SIMD, no more: every extra + // split is another partial for the reduction to read. + static const int target_waves = + 8 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int groups = std::max(1, (num_q + QG - 1) / QG) * num_kv_heads * nrt; + const int ns = + std::min({kMfmaSplits, (int)num_kv_splits, + std::max(1, (target_waves + groups - 1) / groups)}); + if (num_q == 0) return ns; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const int spb = kMfmaNw / nrt; + dim3 grid((num_q + QG - 1) / QG, num_kv_heads, (ns + spb - 1) / spb); + #define SQM_G(B, KC, G, T) \ + decode_mfma<<>>( \ + (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \ + block_table.data_ptr(), q_to_req.data_ptr(), \ + q_to_klen.data_ptr(), mid_o.data_ptr(), \ + k_signs.data_ptr(), (float)sm_scale, num_q, num_q_heads, \ + num_kv_heads, cache.size(2), block_table.size(1), block_table.size(0), \ + cache.size(0), ns, kMfmaMinTps, nrt, query.stride(0), query.stride(1), \ + cache.stride(0), cache.stride(1), cache.stride(2), mid_o.stride(0), \ + mid_o.stride(1), mid_o.stride(2)) + #define SQM(B, KC, T) \ + do { \ + if (QG == 4) \ + SQM_G(B, KC, 4, T); \ + else \ + SQM_G(B, KC, 1, T); \ + } while (0) + const bool bf = query.dtype() == at::kBFloat16; + if (fmt == 4) { + if (bf) + SQM(4, false, __hip_bfloat16); + else + SQM(4, false, half); + } else if (fmt == 3) { + if (bf) + SQM(3, false, __hip_bfloat16); + else + SQM(3, false, half); + } else { + if (bf) + SQM(3, true, __hip_bfloat16); + else + SQM(3, true, half); + } + #undef SQM + #undef SQM_G + return ns; +} + +#endif // USE_ROCM diff --git a/csrc/attention/octave_decode_wmma_rdna3.cu b/csrc/attention/octave_decode_wmma_rdna3.cu new file mode 100644 index 000000000000..6f4c6537a08f --- /dev/null +++ b/csrc/attention/octave_decode_wmma_rdna3.cu @@ -0,0 +1,687 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave decode on WMMA for RDNA3 (gfx11). Independent waves, one (split, +// 16-query row tile) per wave, transposed orientation, and lane pairs sharing +// A fragments through a DPP swap. The products: +// +// QK S^T[token][q] = K x Q^T over 16 k-steps of 16 dims: steps 0-3 the +// RoPE block (4-bit codes), 4-15 the NoPE blocks (3-bit codes). Codes +// are mapped to their int8 codebook values and fed as x + 1152. Each +// 64-dim block has its own per-token scale, so the steps accumulate +// into four tiles that are scaled and debiased in the softmax. Compact +// K: 16 steps of 3-bit codes and one scale, so one tile. +// PV O^T[dim][q] = V^T x P^T with the V codebook values fed as x + 1152; +// the per-token V scale is folded into P, the bias removed with sum(P). +// +// The query is rotated like K while it is staged into LDS; the output stays +// rotated and octave_reduce rotates it back. + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + #include + + #include "octave_format.cuh" + + #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1100__) + #define OCTAVE_WMMA_STUB + #endif + +namespace octave_wmma { + +typedef _Float16 h16v __attribute__((ext_vector_type(16))); +typedef _Float16 h8v __attribute__((ext_vector_type(8))); +typedef _Float16 h2v __attribute__((ext_vector_type(2))); +typedef float f8v __attribute__((ext_vector_type(8))); +typedef uint32_t u8v __attribute__((ext_vector_type(8))); +typedef const int __attribute__((address_space(4))) cint; + +using octave::D; +using octave::Format; +using octave::LUT_ONE; +using octave::R; +constexpr int QG = 4; // query tokens per block +// Exact recent and leading positions (sparse window mode). +[[maybe_unused]] constexpr int kWindow = 64; +[[maybe_unused]] constexpr int kSinks = 4; +[[maybe_unused]] constexpr int QROW = D + 8; // padded LDS row (bank spread) +[[maybe_unused]] constexpr float PSCALE = 256.0f; +[[maybe_unused]] constexpr float BIAS = 1152.0f; // int8 byte x fed as x + 1152 +[[maybe_unused]] constexpr uint32_t X = 0x80808080u; + +__device__ __forceinline__ uint32_t swap1(uint32_t v) { + return (uint32_t)__builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xF, 0xF, + false); +} +__device__ __forceinline__ uint32_t xhalf(uint32_t v) { + return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false, + false); +} +__device__ __forceinline__ float xhalf(float v) { + return __uint_as_float(xhalf(__float_as_uint(v))); +} + +// Bytes b -> fp16 1024 + b, exact. +__device__ __forceinline__ h16v bytes_to_h16(uint32_t w0, uint32_t w1, + uint32_t w2, uint32_t w3) { + const uint32_t ws[4] = {w0, w1, w2, w3}; + u8v r; + #pragma unroll + for (int i = 0; i < 4; ++i) { + r[2 * i] = __builtin_amdgcn_perm(0x64646464u, ws[i], 0x07010700u); + r[2 * i + 1] = __builtin_amdgcn_perm(0x64646464u, ws[i], 0x07030702u); + } + return __builtin_bit_cast(h16v, r); +} + +__device__ __forceinline__ h16v ld16(const _Float16* p) { + const h8v lo = *(const h8v*)p; + const h8v hi = *(const h8v*)(p + 8); + return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15); +} + +// 4 dwords of 4 tokens each (same 4 dims) -> 4 dwords of 4 dims each. +__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2, + uint32_t a3, uint32_t o[4]) { + const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u); + const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u); + const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u); + const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u); + o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u); + o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u); + o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u); + o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u); +} + +template +__device__ __forceinline__ float to_f(T x) { + return (float)x; +} +template <> +__device__ __forceinline__ float to_f(half x) { + return __half2float(x); +} + +__device__ __forceinline__ float sign_of(const int* bits, int i) { + return ((bits[i >> 5] >> (i & 31)) & 1) ? -1.0f : 1.0f; +} + +// SP: sparse mode. Row qi attends to the logical positions idx[qi, 0..topk) +// (negative entries are skipped) instead of its first q_to_klen[qi] tokens, +// and every row is its own segment since rows do not share positions. +template +__global__ __launch_bounds__(64 * NSB) void decode_wmma( + const QT* __restrict__ Q, const uint8_t* __restrict__ cache, + const int* __restrict__ block_table, const int* __restrict__ q_to_req, + const int* __restrict__ q_to_klen, float* __restrict__ mid_o, + const int* __restrict__ k_signs, float sm_scale, int num_q, int num_q_heads, + int num_kv_heads, int block_size, int max_blocks, int num_reqs, + int num_phys_blocks, int num_splits, int min_tps, int num_row_tiles, + int64_t sq0, int64_t sq1, int64_t scb, int64_t sch, int64_t sct, + int64_t smo, int64_t smh, int64_t sms, const int* __restrict__ idx, + int64_t sidx, int topk, const int* __restrict__ qpos, + const int* __restrict__ wtags, int wmask, const int* __restrict__ stags, + int smask) { + #ifndef OCTAVE_WMMA_STUB + using F = Format; + constexpr int NG = KC ? 1 : 4; // QK tiles (one per K scale) + const int grp = blockIdx.x; + const int kvh = blockIdx.y; + const int tid = threadIdx.x; + const int w = __builtin_amdgcn_readfirstlane(tid >> 5), lane = tid & 31; + // Two row tiles when the group can hold more than 16 query vectors; + // otherwise every wave takes its own split. + const int rt = num_row_tiles == 2 ? (w & 1) : 0; + const int splits_per_block = 2 * NSB / num_row_tiles; + const int si = + blockIdx.z * splits_per_block + (num_row_tiles == 2 ? w >> 1 : w); + const int j = lane & 15, h = lane >> 4; + const bool useful = (j & 1) == h; + const int rr = useful ? j : (j ^ 1); + const int hpk = num_q_heads / num_kv_heads; + const int max_kv = max_blocks * block_size; + const float sm_log2 = sm_scale * 1.4426950408889634f / LUT_ONE; + + __shared__ __attribute__((aligned(16))) _Float16 sQ[32][QROW]; + + // Stage the group's queries rotated like K: lane owns 8 dims. + for (int v = w; v < 32; v += 2 * NSB) { + const int r = v / hpk, g = v % hpk; + const int qi = grp * QG + r; + float x[8]; + const bool live = v < QG * hpk && qi < num_q; + #pragma unroll + for (int i = 0; i < 8; ++i) + x[i] = live ? to_f(Q[(int64_t)qi * sq0 + + (int64_t)(kvh * hpk + g) * sq1 + lane * 8 + i]) + : 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) x[i] *= sign_of(k_signs, lane * 8 + i); + #pragma unroll + for (int hh = 1; hh < 8; hh <<= 1) { + #pragma unroll + for (int i = 0; i < 8; ++i) + if ((i & hh) == 0) { + const float a = x[i], b = x[i + hh]; + x[i] = a + b; + x[i + hh] = a - b; + } + } + #pragma unroll + for (int m = 1; m <= (KC ? 16 : 4); m <<= 1) { + const bool hi = lane & m; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const float o = __shfl_xor(x[i], m); + x[i] = hi ? o - x[i] : x[i] + o; + } + } + #pragma unroll + for (int i = 0; i < 8; ++i) + sQ[v][lane * 8 + i] = (_Float16)(x[i] * (KC ? 0.0625f : 0.125f)); + } + __syncthreads(); + + int rreq[QG], rlen[QG]; + #pragma unroll + for (int r = 0; r < QG; ++r) { + const int qi = grp * QG + r; + int rq = -1, ln = 0; + if (qi < num_q) { + rq = q_to_req[qi]; + ln = SP ? topk : q_to_klen[qi]; + if (rq < 0 || rq >= num_reqs) { + rq = 0; + ln = 0; + } + ln = min(max(ln, 0), SP ? topk : max_kv); + } + rreq[r] = __builtin_amdgcn_readfirstlane(rq); + rlen[r] = __builtin_amdgcn_readfirstlane(ln); + } + + // Per-segment sums of this lane's query (debiasing): half h sums dims + // [128 h, 128 h + 128) as two 64-dim segments; the halves swap. + const int qv = rt * 16 + j; + float bq[4]; + { + const h2v one = {(_Float16)1.0f, (_Float16)1.0f}; + float s0 = 0.f, s1 = 0.f; + #pragma unroll + for (int s = 0; s < 8; ++s) { + const h8v a = *(const h8v*)&sQ[qv][128 * h + 8 * s]; + const h8v b = *(const h8v*)&sQ[qv][128 * h + 64 + 8 * s]; + #pragma unroll + for (int i = 0; i < 8; i += 2) { + s0 = __builtin_amdgcn_fdot2((h2v){a[i], a[i + 1]}, one, s0, false); + s1 = __builtin_amdgcn_fdot2((h2v){b[i], b[i + 1]}, one, s1, false); + } + } + const float o0 = xhalf(s0), o1 = xhalf(s1); + bq[0] = (h ? o0 : s0) * BIAS; + bq[1] = (h ? o1 : s1) * BIAS; + bq[2] = (h ? s0 : o0) * BIAS; + bq[3] = (h ? s1 : o1) * BIAS; + if constexpr (KC) bq[0] += bq[1] + bq[2] + bq[3]; + } + + int seg0 = 0; + while (seg0 < QG && rreq[seg0] >= 0) { + int seg1 = seg0 + 1; + if constexpr (!SP) + while (seg1 < QG && rreq[seg1] == rreq[seg0]) ++seg1; + const int req = rreq[seg0]; + int Lmax = 0; + for (int r = seg0; r < seg1; ++r) Lmax = max(Lmax, rlen[r]); + const int qv_lo = seg0 * hpk, qv_hi = seg1 * hpk; + int tps = (Lmax + num_splits - 1) / num_splits; + if (tps < min_tps) tps = min_tps; + tps = (tps + 15) & ~15; + + const bool tile_live = qv_hi > rt * 16 && qv_lo < rt * 16 + 16; + const bool qlive = qv >= qv_lo && qv < qv_hi; + if (blockIdx.z * splits_per_block * tps >= Lmax) { + if (tile_live && si < num_splits && h == 0 && qlive) { + int qx = qv; + asm volatile("" : "+v"(qx)); + float* op = mid_o + (int64_t)(grp * QG + qx / hpk) * smo + + (int64_t)(kvh * hpk + qx % hpk) * smh + (int64_t)si * sms; + op[D] = -INFINITY; + op[D + 1] = 0.0f; + } + seg0 = seg1; + continue; + } + + if (tile_live && si < num_splits) { + int qlen = 0; + #pragma unroll + for (int r = 0; r < QG; ++r) + if (qlive && qv / hpk == r) qlen = rlen[r]; + + const int start = si * tps; + const int end = min(start + tps, Lmax); + float m = -INFINITY, l = 0.0f, cp = 0.0f; + f8v O[16]; + #pragma unroll + for (int c = 0; c < 16; ++c) O[c] = (f8v){0, 0, 0, 0, 0, 0, 0, 0}; + + const uint8_t* cache_h = cache + (int64_t)kvh * sch; + auto pickw = [&](uint32_t v, int t) { + const uint32_t ev = __builtin_amdgcn_readlane(v, t); + const uint32_t od = __builtin_amdgcn_readlane(v, t + 1); + return h ? od : ev; + }; + auto tile_ptr = [&](int pos) { + const int lb = pos / block_size; + int pb = ((cint*)block_table)[req * max_blocks + lb]; + if (pb < 0 || pb >= num_phys_blocks) pb = 0; + return cache_h + (int64_t)pb * scb + + (int64_t)(pos - lb * block_size) * sct; + }; + // Sparse mode: slot of the pos-th selected position of this row, and + // whether that position maps to a live cache slot. Dead entries read + // slot 0 of block 0 and are masked out of the softmax. + const int* irow = SP ? idx + (int64_t)(grp * QG + seg0) * sidx : nullptr; + auto sel_block = [&](int pos, int& off) { + const int t = irow[min(pos, topk - 1)]; + const int lb = max(t, 0) / block_size; + off = max(t, 0) - lb * block_size; + const int pb = (t >= 0 && lb < max_blocks) + ? ((cint*)block_table)[req * max_blocks + lb] + : -1; + return (pb >= 0 && pb < num_phys_blocks) ? pb : -1; + }; + // Exact-window tokens (attended by octave_window_attn instead): the + // last kWindow positions before the query, and the first kSinks. + const int rpos = SP && wtags ? qpos[grp * QG + seg0] : 0; + auto in_window = [&](int pos, int pb, int off) { + const int t = irow[min(pos, topk - 1)]; + const int d = rpos - t; + const int slot = pb * block_size + off; + return (d >= 0 && d < kWindow && wtags[slot & wmask] == slot) || + (t >= 0 && t < kSinks && stags[slot & smask] == slot); + }; + auto sel_ptr = [&](int pos) { + int off; + const int pb = sel_block(pos, off); + return pb < 0 ? cache_h + : cache_h + (int64_t)pb * scb + (int64_t)off * sct; + }; + + // Raw K loads, one tile ahead. Each lane of a pair owns 8 of the 16 + // k-steps of token rr: NR4 4-bit steps (8 bytes of codes each), then + // 3-bit steps (2-bit plane word, 1-bit plane word pre-shifted). Block + // K: the useful lane's slots are steps {0,1,4..9}, its partner's + // {2,3,10..15}. Compact K: {0..7} and {8..15}. The partner's reach the + // useful lane through the DPP swap. + constexpr int NR4 = F::KR / 32; // 4-bit slots per lane + constexpr int NR3 = 8 - NR4; // 3-bit slots per lane + uint2 kr_raw[NR4 > 0 ? NR4 : 1]; + uint2 kn_raw[NR3]; + // Scales of token j (lane j), raw fp16 pairs. Block K: {k0, k1}, + // {k2, k3}, {v}. Compact K: {k, v}. + uint32_t sc_a = 0, sc_b = 0, sc_v = 0; + auto load_k = [&](const uint8_t* tile, int base, float dep) { + uint32_t off = (uint32_t)(rr * (int)sct); + asm volatile("" : "+v"(off) : "v"(dep)); + const uint8_t* pk = SP ? sel_ptr(base + rr) : tile + off; + const int chunk0 = useful ? 0 : NR3; + if constexpr (NR4 > 0) { + const int rope0 = useful ? 0 : NR4; + #pragma unroll + for (int k = 0; k < NR4; ++k) + kr_raw[k] = *(const uint2*)(pk + 8 * (rope0 + k)); + } + #pragma unroll + for (int k = 0; k < NR3; ++k) { + const int c = chunk0 + k; + kn_raw[k].x = *(const uint32_t*)(pk + F::OFF_KN + 4 * c); + kn_raw[k].y = + *(const uint32_t*)(pk + F::OFF_KN + F::KN / 4 + 4 * (c >> 1)) >> + (4 * (c & 1)); + } + const uint8_t* ps = + (SP ? sel_ptr(base + j) : tile + (int64_t)j * sct) + F::OFF_SC; + if constexpr (KC) { + sc_a = *(const uint32_t*)ps; + sc_v = sc_a >> 16; + } else { + const uint2 u = *(const uint2*)ps; + sc_a = u.x; + sc_b = u.y; + sc_v = *(const uint16_t*)(ps + 8); + } + }; + // Raw V loads: lane rr's 16 dims of the even (useful) or odd tokens. + uint2 vraw[8]; + auto load_v = [&](const uint8_t* tile, int base, float dep) { + uint32_t off = (uint32_t)((useful ? 0 : 1) * (int)sct); + asm volatile("" : "+v"(off) : "v"(dep)); + const uint8_t* pv = tile + off + F::OFF_V; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const uint8_t* p = + SP ? sel_ptr(base + 2 * i + (useful ? 0 : 1)) + F::OFF_V + : pv + (int64_t)(2 * i) * sct; + if constexpr (VB == 4) { + vraw[i] = *(const uint2*)(p + 8 * rr); + } else { + vraw[i].x = *(const uint32_t*)(p + 4 * rr); + vraw[i].y = *(const uint32_t*)(p + D / 4 + 4 * (rr >> 1)); + } + } + }; + // Biased codebook bytes (x + 128) of a 16-dim step from its raw words. + auto rope_bytes = [&](uint2 raw, uint32_t o[4]) { + o[0] = octave::lut4(raw.x & 0x0F0F0F0Fu) ^ X; + o[1] = octave::lut4((raw.x >> 4) & 0x0F0F0F0Fu) ^ X; + o[2] = octave::lut4(raw.y & 0x0F0F0F0Fu) ^ X; + o[3] = octave::lut4((raw.y >> 4) & 0x0F0F0F0Fu) ^ X; + }; + auto nope_bytes = [&](uint2 raw, uint32_t o[4]) { + #pragma unroll + for (int i = 0; i < 4; ++i) + o[i] = octave::lut3(((raw.x >> (2 * i)) & 0x03030303u) | + (((raw.y >> i) & 0x01010101u) << 2)) ^ + X; + }; + + // V of a tile is loaded after its QK (softmax hides it): holding it + // through the QK, next to O, the four S tiles and the raw K, spilled. + const uint8_t* tcur = nullptr; + if (start < end) { + tcur = SP ? nullptr : tile_ptr(start); + load_k(tcur, start, 0.f); + } + + for (int base = start; base < end; base += 16) { + auto lo16 = [](uint32_t x) { + return __half2float(__ushort_as_half((unsigned short)(x & 0xFFFF))); + }; + auto hi16 = [](uint32_t x) { + return __half2float(__ushort_as_half((unsigned short)(x >> 16))); + }; + const uint32_t sa_c = sc_a, sb_c = sc_b, sv_c = sc_v; + + // QK. Iteration k feeds the useful lane's slot k and, through the + // swap, the partner's slot k. + f8v S[NG]; + #pragma unroll + for (int g = 0; g < NG; ++g) S[g] = (f8v){0, 0, 0, 0, 0, 0, 0, 0}; + int qoff = qv * QROW; + asm volatile("" : "+v"(qoff)); + constexpr int kStepOwn[8] = {0, 1, 4, 5, 6, 7, 8, 9}; + constexpr int kStepPar[8] = {2, 3, 10, 11, 12, 13, 14, 15}; + #pragma unroll + for (int k = 0; k < 8; ++k) { + uint32_t wd[4]; + if (k < NR4) + rope_bytes(kr_raw[k < NR4 ? k : 0], wd); + else + nope_bytes(kn_raw[k - NR4], wd); + const int so = KC ? k : kStepOwn[k], sp = KC ? 8 + k : kStepPar[k]; + const int go = KC ? 0 : so / 4, gp = KC ? 0 : sp / 4; + // One A and one B fragment live at a time. Each read waits for the + // previous WMMA: without the dependence the compiler hoists all Q + // fragments and spills. + asm volatile("" : "+v"(qoff) : "v"(S[gp][0])); + { + const h16v b = ld16(&sQ[0][0] + qoff + 16 * so); + const h16v a = bytes_to_h16(wd[0], wd[1], wd[2], wd[3]); + S[go] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, S[go]); + } + __builtin_amdgcn_sched_barrier(0); + asm volatile("" : "+v"(qoff) : "v"(S[go][0])); + { + const h16v b = ld16(&sQ[0][0] + qoff + 16 * sp); + const h16v a = bytes_to_h16(swap1(wd[0]), swap1(wd[1]), + swap1(wd[2]), swap1(wd[3])); + S[gp] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, S[gp]); + } + __builtin_amdgcn_sched_barrier(0); + } + + load_v(tcur, base, S[0][0]); + // Next tile's K and scales now: softmax and PV hide them. + const bool next = base + 16 < end; + const uint8_t* tn = next && !SP ? tile_ptr(base + 16) : nullptr; + if (next) load_k(tn, base + 16, S[0][0]); + + float sc[8]; + float mloc = -INFINITY; + #pragma unroll + for (int e = 0; e < 8; ++e) { + const int t = base + 2 * e + h; + bool ok = t < end && t < qlen; + if constexpr (SP) { + int off; + const int pb = sel_block(t, off); + ok = ok && pb >= 0 && !(wtags && in_window(t, pb, off)); + } + const uint32_t a = pickw(sa_c, 2 * e); + float v; + if constexpr (KC) { + v = lo16(a) * (S[0][e] - bq[0]); + } else { + const uint32_t b = pickw(sb_c, 2 * e); + v = lo16(a) * (S[0][e] - bq[0]) + hi16(a) * (S[1][e] - bq[1]) + + lo16(b) * (S[2][e] - bq[2]) + hi16(b) * (S[3][e] - bq[3]); + } + sc[e] = ok ? v * sm_log2 : -INFINITY; + mloc = fmaxf(mloc, sc[e]); + } + const float mn = fmaxf(m, fmaxf(mloc, xhalf(mloc))); + float alpha = 1.0f; + if (mn > m) + alpha = (m == -INFINITY) ? 0.0f : __builtin_amdgcn_exp2f(m - mn); + m = mn; + uint32_t pd[4]; + float psum = 0.0f, lsum = 0.0f; + #pragma unroll + for (int e = 0; e < 8; e += 2) { + const float v0 = lo16(pickw(sv_c, 2 * e)); + const float v1 = lo16(pickw(sv_c, 2 * e + 2)); + const float p0 = + sc[e] == -INFINITY ? 0.0f : __builtin_amdgcn_exp2f(sc[e] - mn); + const float p1 = sc[e + 1] == -INFINITY + ? 0.0f + : __builtin_amdgcn_exp2f(sc[e + 1] - mn); + lsum += p0 + p1; + const _Float16 h0 = (_Float16)(p0 * v0 * PSCALE); + const _Float16 h1 = (_Float16)(p1 * v1 * PSCALE); + psum += (float)h0 + (float)h1; + pd[e >> 1] = __builtin_bit_cast(uint32_t, (h2v){h0, h1}); + } + l = l * alpha + lsum; + cp = cp * alpha + psum; + if (__builtin_amdgcn_ballot_w32(alpha != 1.0f)) { + #pragma unroll + for (int c = 0; c < 16; ++c) O[c] *= alpha; + } + u8v pf; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t o = xhalf(pd[i]); + pf[i] = h ? o : pd[i]; + pf[4 + i] = h ? pd[i] : o; + } + const h16v pb16 = __builtin_bit_cast(h16v, pf); + + // Biased codebook bytes of dims 4q..4q+3 of this lane's chunk, + // token i. + auto vq = [&](int i, int q) -> uint32_t { + if constexpr (VB == 4) { + const uint32_t w = q < 2 ? vraw[i].x : vraw[i].y; + return octave::lut4((w >> (4 * (q & 1))) & 0x0F0F0F0Fu) ^ X; + } else { + const uint32_t hi = vraw[i].y >> (4 * (rr & 1)); + return octave::lut3(((vraw[i].x >> (2 * q)) & 0x03030303u) | + (((hi >> q) & 0x01010101u) << 2)) ^ + X; + } + }; + #pragma unroll + for (int q = 0; q < 4; ++q) { + uint32_t t0[4], t1[4]; + tr4(vq(0, q), vq(1, q), vq(2, q), vq(3, q), t0); + tr4(vq(4, q), vq(5, q), vq(6, q), vq(7, q), t1); + #pragma unroll + for (int b = 0; b < 4; ++b) { + const int c = 4 * q + b; + uint32_t o0 = t0[b], o1 = t1[b]; + if (c >= 2) + asm volatile("" : "+v"(o0), "+v"(o1) : "v"(O[c - 2][0])); + const h16v a = bytes_to_h16(o0, o1, swap1(o0), swap1(o1)); + O[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, pb16, O[c]); + } + } + tcur = tn; + } + + l += xhalf(l); + cp += xhalf(cp); + if (qlive) { + int qx = qv; + asm volatile("" : "+v"(qx)); + float* op = mid_o + (int64_t)(grp * QG + qx / hpk) * smo + + (int64_t)(kvh * hpk + qx % hpk) * smh + (int64_t)si * sms; + const bool any = m != -INFINITY; + const float corr = BIAS * cp; + constexpr float kOut = 1.0f / (PSCALE * LUT_ONE); + #pragma unroll + for (int e = 0; e < 8; ++e) { + float* dp = op + 16 * (2 * e + h); + #pragma unroll + for (int c = 0; c < 16; c += 4) { + float4 v; + v.x = any ? (O[c][e] - corr) * kOut : 0.0f; + v.y = any ? (O[c + 1][e] - corr) * kOut : 0.0f; + v.z = any ? (O[c + 2][e] - corr) * kOut : 0.0f; + v.w = any ? (O[c + 3][e] - corr) * kOut : 0.0f; + *(float4*)(dp + c) = v; + } + } + if (h == 0) { + op[D] = m; + op[D + 1] = l; + } + } + } + seg0 = seg1; + } + #endif // OCTAVE_WMMA_STUB +} + +} // namespace octave_wmma + +// Splits per row group and minimum tokens per split, as in the int8 kernel. +constexpr int kSqWmmaSplits = 192; +constexpr int kSqWmmaMinTps = 64; +constexpr int kSqWmmaNsb = 2; + +int octave_decode_mfma(torch::Tensor query, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor q_to_req, + torch::Tensor q_to_klen, torch::Tensor mid_o, + torch::Tensor k_signs, double sm_scale, + int64_t num_kv_splits, int64_t fmt, int64_t query_group); + +// Launches the WMMA decode when the shape is covered; returns the number of +// splits it wrote (for the reduce), or 0 when the caller must use its own +// kernel. +// With `indices` ([num_q, topk] int32 logical positions) the sparse variant +// runs and q_to_klen is ignored. +int octave_decode_wmma(torch::Tensor query, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor q_to_req, + torch::Tensor q_to_klen, torch::Tensor mid_o, + torch::Tensor k_signs, double sm_scale, + int64_t num_kv_splits, int64_t fmt, + const std::optional& indices, + const int* qpos, const int* wtags, int wmask, + const int* stags, int smask, int64_t query_group) { + using namespace octave_wmma; + // 0: gfx1100 (this file), 1: gfx942 (octave_decode_mfma_cdna.cu, dense + // only). + static const int arch = [] { + const std::string name = + at::cuda::getCurrentDeviceProperties()->gcnArchName; + if (name.rfind("gfx1100", 0) == 0) return 0; + if (name.rfind("gfx942", 0) == 0) return 1; + return -1; + }(); + if (arch == 1) + return indices.has_value() + ? 0 + : octave_decode_mfma(query, cache, block_table, q_to_req, + q_to_klen, mid_o, k_signs, sm_scale, + num_kv_splits, fmt, query_group); + const bool arch_ok = arch == 0; + const int num_q = query.size(0); + const int num_q_heads = query.size(1); + const int num_kv_heads = cache.size(1); + const int hpk = num_q_heads / num_kv_heads; + const bool eligible = + arch_ok && query.size(2) == D && query.stride(2) == 1 && hpk * QG <= 32 && + cache.size(2) % 16 == 0 && + cache.stride(2) % (fmt & octave::kCompactFlag ? 4 : 8) == 0 && + num_kv_splits >= 1; + if (!eligible) return 0; + const int ns = std::min(kSqWmmaSplits, (int)num_kv_splits); + if (num_q == 0) return ns; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const bool sp = indices.has_value(); + const int* idx = sp ? indices->data_ptr() : nullptr; + const int64_t sidx = sp ? indices->stride(0) : 0; + const int topk = sp ? (int)indices->size(1) : 0; + const int nrt = std::min(num_q, QG) * hpk > 16 ? 2 : 1; + const int spb = 2 * kSqWmmaNsb / nrt; + dim3 grid((num_q + QG - 1) / QG, num_kv_heads, (ns + spb - 1) / spb); + #define SQW_SP(B, KC, T, SP) \ + decode_wmma \ + <<>>( \ + (const T*)query.data_ptr(), (const uint8_t*)cache.data_ptr(), \ + block_table.data_ptr(), q_to_req.data_ptr(), \ + q_to_klen.data_ptr(), mid_o.data_ptr(), \ + k_signs.data_ptr(), (float)sm_scale, num_q, num_q_heads, \ + num_kv_heads, cache.size(2), block_table.size(1), \ + block_table.size(0), cache.size(0), ns, kSqWmmaMinTps, nrt, \ + query.stride(0), query.stride(1), cache.stride(0), \ + cache.stride(1), cache.stride(2), mid_o.stride(0), \ + mid_o.stride(1), mid_o.stride(2), idx, sidx, topk, qpos, wtags, \ + wmask, stags, smask) + #define SQW(B, KC, T) \ + do { \ + if (sp) \ + SQW_SP(B, KC, T, true); \ + else \ + SQW_SP(B, KC, T, false); \ + } while (0) + const bool bf = query.dtype() == at::kBFloat16; + if (fmt == 4) { + if (bf) + SQW(4, false, __hip_bfloat16); + else + SQW(4, false, half); + } else if (fmt == 3) { + if (bf) + SQW(3, false, __hip_bfloat16); + else + SQW(3, false, half); + } else { + if (bf) + SQW(3, true, __hip_bfloat16); + else + SQW(3, true, half); + } + #undef SQW + #undef SQW_SP + return ns; +} + +#endif // USE_ROCM diff --git a/csrc/attention/octave_format.cuh b/csrc/attention/octave_format.cuh new file mode 100644 index 000000000000..faf8a8848646 --- /dev/null +++ b/csrc/attention/octave_format.cuh @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Byte layout of one Octave slot (one token, one KV head). The PyTorch +// reference in vllm/v1/attention/ops/rocm_octave.py is the source of truth. +// +// [0, OFF_KN) K RoPE block (dims 0-63), rotated, 4-bit codes +// [OFF_KN, OFF_V) K NoPE blocks (dims 64-255), rotated, 3-bit codes +// [OFF_V, OFF_SC) V, rotated, VBITS-bit codes +// [OFF_SC, ...) fp16 scales: K block 0..3, V +// +// Compact K (KC): K rotated over all 256 dims like V, 3-bit codes for every +// dim (OFF_KN = 0), one K scale; slots pad to 4 bytes instead of 8. +// +// 4-bit codes: dims 8i..8i+7 in one dword, low nibbles = dims 0-3, high +// nibbles = dims 4-7. 3-bit codes: a 2-bit plane (16 dims per dword, dims +// 4k..4k+3 at bit 2k of bytes 0-3) followed by a 1-bit plane (32 dims per +// dword, dims 4k..4k+3 at bit k of bytes 0-3). A code c of a block with +// scale s stands for LUT[c] * s / LUT_ONE, LUT being the int8 Lloyd-Max +// codebook of its width. + +#pragma once + +#include + +namespace octave { + +constexpr int D = 256; // head size +constexpr int R = 64; // RoPE dims (K block 0) +constexpr int N = D - R; // NoPE dims (K blocks 1-3) +constexpr int NKB = D / 64; // K Hadamard blocks +constexpr float LUT_ONE = 64.0f; + +template +struct Format { + static constexpr int KR = KC ? 0 : R; // K dims with 4-bit codes + static constexpr int KN = D - KR; // K dims with 3-bit codes + static constexpr int OFF_KN = KR * 4 / 8; + static constexpr int OFF_V = OFF_KN + KN * 3 / 8; + static constexpr int OFF_SC = OFF_V + D * VBITS / 8; + static constexpr int NUM_KSC = KC ? 1 : NKB; + static constexpr int NUM_SC = NUM_KSC + 1; + static constexpr int ALIGN = KC ? 4 : 8; + static constexpr int SLOT = (OFF_SC + 2 * NUM_SC + ALIGN - 1) / ALIGN * ALIGN; +}; + +// Format id passed by the host: V bits, plus 16 for compact K. +constexpr int kCompactFlag = 16; + +// int8 Lloyd-Max codebooks as perm tables: byte c of the table is LUT[c]. +// 3-bit: -127 -79 -45 -14 14 45 79 127 +// 4-bit: -126 -95 -74 -58 -43 -30 -18 -6 6 18 30 43 58 74 95 126 +constexpr uint32_t LUT3_LO = 0xF2D3B181u, LUT3_HI = 0x7F4F2D0Eu; +constexpr uint32_t LUT4_0 = 0xC6B6A182u, LUT4_1 = 0xFAEEE2D5u; +constexpr uint32_t LUT4_2 = 0x2B1E1206u, LUT4_3 = 0x7E5F4A3Au; + +// Four 3-bit codes (one per byte) -> their four int8 codebook values. +__device__ __forceinline__ uint32_t lut3(uint32_t codes) { + return __builtin_amdgcn_perm(LUT3_HI, LUT3_LO, codes); +} +// Four 4-bit codes (one per byte) -> their four int8 codebook values. +__device__ __forceinline__ uint32_t lut4(uint32_t codes) { + const uint32_t lo = + __builtin_amdgcn_perm(LUT4_1, LUT4_0, codes & 0x07070707u); + const uint32_t hi = + __builtin_amdgcn_perm(LUT4_3, LUT4_2, codes & 0x07070707u); + const uint32_t m = ((codes >> 3) & 0x01010101u) * 0xFFu; + return (hi & m) | (lo & ~m); +} +template +__device__ __forceinline__ uint32_t lut(uint32_t codes) { + if constexpr (BITS == 4) + return lut4(codes); + else + return lut3(codes); +} + +} // namespace octave diff --git a/csrc/attention/octave_prefill_cdna.cu b/csrc/attention/octave_prefill_cdna.cu new file mode 100644 index 000000000000..d1638219b79f --- /dev/null +++ b/csrc/attention/octave_prefill_cdna.cu @@ -0,0 +1,666 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave prefill for CDNA3 (gfx942), head size 256, fp16. Same contract as +// octave_prefill_rdna3.cu: reads the packed cache directly, everything in the +// rotated space, phase 1 writes (O, m, l) partials over the cached prefix and +// phase 2 runs causal attention over the chunk and merges. +// +// Both phases compute transposed tiles so that every MFMA result is already in +// the layout the next MFMA consumes (16x16 MFMA, wave64: lane l holds column +// l % 16 and rows 4 (l / 16) .. 4 (l / 16) + 3): +// S^T[token][q] = K x Q^T (phase 1: int8 16x16x32, phase 2: fp16) +// O^T[dim][q] += V^T x P^T (fp16 16x16x16) +// so P^T comes straight out of the S^T accumulator, one query per lane. +// +// Phase 1: one block = 128 query vectors (token, head) of one KV head, 8 waves +// of 16 query vectors x 256 dims. 256 threads unpack the next 16-token tile +// into LDS while the current one is consumed. + +#include +#include + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + + #include "octave_format.cuh" + + #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx942__) + #define OCTAVE_CDNA_STUB + #endif + +namespace octave_cdna { + +using octave::D; +using octave::Format; +using octave::LUT_ONE; + +typedef _Float16 h4v __attribute__((ext_vector_type(4))); +typedef _Float16 h2v __attribute__((ext_vector_type(2))); +typedef float f4v __attribute__((ext_vector_type(4))); +typedef int i4v __attribute__((ext_vector_type(4))); +typedef uint32_t u4v __attribute__((ext_vector_type(4))); +typedef uint32_t u2v __attribute__((ext_vector_type(2))); + +constexpr int KT = 16; // tokens per tile +[[maybe_unused]] constexpr int FR = D / 16; // 16-dim output fragments +// Both phases run 8 waves (128 query vectors of one KV head) per block: every +// block reads the whole K/V it attends to, so more query vectors per block +// means less of that work per query. +constexpr int NW1 = 8; +constexpr int NT1 = NW1 * 64; +constexpr int QB1 = NW1 * 16; +constexpr int KT1 = 32; // phase-1 tile: two 16-token sub-tiles +[[maybe_unused]] constexpr int VST1 = KT1 + 4; +[[maybe_unused]] constexpr int KROW = D + 16; // bytes per int8 K row +[[maybe_unused]] constexpr int VST = KT + 4; // halves per V^T row +[[maybe_unused]] constexpr int HROW = D + 8; // halves per fp16 K row +// Phase 2 LDS rows, picked for conflict-free banks: a K row is 2 dwords +// mod 32, and V^T rows of 40 halves with the 4-token group XORed by +// (row >> 3) & 7 spread both the transposed stores and the MFMA reads. +[[maybe_unused]] constexpr int KROW2 = D + 4; +[[maybe_unused]] constexpr int VST2 = KT1 + 8; +[[maybe_unused]] constexpr float PSCALE = 4096.0f; + +constexpr int kMinSplitTokens = 2048; +__device__ __forceinline__ int used_splits(int ctx, int num_splits) { + if (ctx <= 0) return 0; + return min(num_splits, (ctx + kMinSplitTokens - 1) / kMinSplitTokens); +} +__device__ __forceinline__ int split_len(int ctx, int used) { + return ((ctx + used - 1) / used + 31) / 32 * 32; +} + +// Reductions over the four lanes of a query (l, l ^ 16, l ^ 32, l ^ 48). +__device__ __forceinline__ float qmax4(float v) { + v = fmaxf(v, __shfl_xor(v, 16)); + return fmaxf(v, __shfl_xor(v, 32)); +} +__device__ __forceinline__ float qsum4(float v) { + v += __shfl_xor(v, 16); + return v + __shfl_xor(v, 32); +} + +// Codes of dims 4k..4k+3 of a 64-dim block (see octave_prefill_rdna3.cu). +template +__device__ __forceinline__ uint32_t block_codes(const uint32_t* raw, + uint32_t hi, int k) { + if constexpr (BITS == 4) { + return (raw[k >> 1] >> (4 * (k & 1))) & 0x0F0F0F0Fu; + } else { + return ((raw[0] >> (2 * (k & 3))) & 0x03030303u) | + (((hi >> (k & 7)) & 0x01010101u) << 2); + } +} + +// 4 dwords, one per token (same 4 dims) -> 4 dwords, one per dim (4 tokens). +__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2, + uint32_t a3, uint32_t o[4]) { + const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u); + const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u); + const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u); + const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u); + o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u); + o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u); + o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u); + o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u); +} + +// =========================================================================== +// Phase 1: prefix +// =========================================================================== +template +__global__ __launch_bounds__(NT1) void prefix_attn( + const __half* __restrict__ q, const uint8_t* __restrict__ cache, + const int* __restrict__ block_table, const int* __restrict__ cu_seqlens_q, + const int* __restrict__ seq_lens, float* __restrict__ ws, float sm_scale, + int num_q_heads, int num_kv_heads, int block_size, int max_blocks, + int num_splits, int total_q_tokens, int64_t sq_tok, int64_t sq_head, + int64_t scb, int64_t sch, int64_t sct) { + #ifndef OCTAVE_CDNA_STUB + using F = Format; + const int seq = blockIdx.x, kvh = blockIdx.y; + const int rowtile = blockIdx.z / num_splits, split = blockIdx.z % num_splits; + const int tid = threadIdx.x; + const int w = tid >> 6, lane = tid & 63; + const int col = lane & 15, gq = lane >> 4; + const int hpk = num_q_heads / num_kv_heads; + + const int q0 = cu_seqlens_q[seq]; + const int qlen = cu_seqlens_q[seq + 1] - q0; + int ctx = seq_lens[seq] - qlen; + ctx = max(0, min(ctx, max_blocks * block_size)); + const int nqv = qlen * hpk; + if (rowtile * QB1 >= nqv) return; + + const int used = used_splits(ctx, num_splits); + if (split >= used) return; + const int k_len = split_len(ctx, used); + const int k_begin = min(ctx, split * k_len); + const int k_stop = min(ctx, k_begin + k_len); + const int ntiles = k_stop > k_begin ? (k_stop - k_begin + KT1 - 1) / KT1 : 0; + + __shared__ __attribute__((aligned(16))) int8_t sK[2][KT1][KROW]; + __shared__ __attribute__((aligned(16))) _Float16 sV[2][D * VST1]; + __shared__ float sKs[2][4][KT1], sVs[2][KT1]; + + // ---- loader mapping. Thread half sub (256 threads each) loads the 16-token + // sub-tile sub of a tile; within it, K: dims 16 kp .. + 15 of block kq of + // token tk; V: dims vd + 4 vh .. + 3 of token quad vq. + const int sub = tid >> 8, lt = tid & 255; + const int tk = lt >> 4, kq = (lt >> 2) & 3, kp = lt & 3; + const int vu = lt >> 1, vh = lt & 1, vq = vu & 3, vd = (vu >> 2) * 8; + const bool k3 = KC || kq > 0; + const int b3 = KC ? kq : kq - 1; + const int koff0 = k3 ? F::OFF_KN + 16 * b3 + 4 * kp : 8 * kp; + const int koff1 = + k3 ? F::OFF_KN + F::KN / 4 + 8 * b3 + 4 * (kp >> 1) : 8 * kp + 4; + constexpr int VW = VB == 4 ? 1 : 2; + const int tlast = k_stop - 1; + const uint8_t* const head = cache + (int64_t)kvh * sch; + const int* const bt = block_table + seq * max_blocks; + + struct Stg { + uint32_t k0, k1; + uint32_t v[4][VW]; + uint32_t s0, s1; + }; + // Sub-tiles start 16-aligned and sit in one cache block (block sizes are + // multiples of 16), so all 16 slots exist: tokens past k_stop are read as + // they are and masked (K to -inf, P to 0 by selection, never by product). + // Slot offsets inside a sub-tile are fixed per thread. + auto sub_start = [&](int i) { + return __builtin_amdgcn_readfirstlane( + min(k_begin + i * KT1 + 16 * sub, tlast & ~15)); + }; + const uint32_t ok0 = tk * (uint32_t)sct + koff0; + const uint32_t ok1 = tk * (uint32_t)sct + koff1; + uint32_t ov[4]; + #pragma unroll + for (int r = 0; r < 4; ++r) { + ov[r] = (4 * vq + r) * (uint32_t)sct + F::OFF_V + + (VB == 4 ? vd / 2 : 4 * (vd >> 4)); + } + const uint32_t os = + (lt & 15) * (uint32_t)sct + F::OFF_SC + (!KC && lt >= 16 ? 8 : 0); + auto fetch = [&](int i, Stg& S) { + const int start = sub_start(i); + const int pb = __builtin_amdgcn_readfirstlane(bt[start / block_size]); + const uint8_t* base = head + (int64_t)(pb < 0 ? 0 : pb) * scb + + (int64_t)(start % block_size) * sct; + S.k0 = *(const uint32_t*)(base + ok0); + S.k1 = *(const uint32_t*)(base + ok1); + #pragma unroll + for (int r = 0; r < 4; ++r) { + S.v[r][0] = *(const uint32_t*)(base + ov[r]); + if constexpr (VB != 4) + S.v[r][1] = *(const uint32_t*)(base + ov[r] - 4 * (vd >> 4) + D / 4 + + 4 * (vd >> 5)); + } + if (lt < 32) { + if constexpr (KC) { + S.s0 = *(const uint32_t*)(base + os); + } else { + const u2v x = *(const u2v*)(base + os); + S.s0 = x[0]; + S.s1 = x[1]; + } + } + }; + auto store = [&](const Stg& S, int bf) { + const int t0 = 16 * sub; + { + const uint32_t r4[2] = {S.k0, S.k1}; + const uint32_t r3[1] = {S.k0}; + const uint32_t hi = S.k1 >> (4 * (kp & 1)); + u4v o; + #pragma unroll + for (int x = 0; x < 4; ++x) + o[x] = k3 ? octave::lut3(block_codes<3>(r3, hi, x)) + : octave::lut4(block_codes<4>(r4, 0, x)); + *(u4v*)&sK[bf][t0 + tk][64 * kq + 16 * kp] = o; + } + uint32_t c[4]; + #pragma unroll + for (int r = 0; r < 4; ++r) { + if constexpr (VB == 4) { + c[r] = (S.v[r][0] >> (4 * vh)) & 0x0F0F0F0Fu; + } else { + const int sl = 2 * ((vd & 15) >> 2) + 2 * vh; + const int sh = ((vd & 31) >> 2) + vh; + c[r] = ((S.v[r][0] >> sl) & 0x03030303u) | + (((S.v[r][1] >> sh) & 0x01010101u) << 2); + } + c[r] = octave::lut(c[r]) ^ 0x80808080u; + } + uint32_t t[4]; + tr4(c[0], c[1], c[2], c[3], t); + // Codebook bytes fed as fp16 1024 + (x + 128), less 1152: exact x. + const h2v kofs = {(_Float16)1152.0f, (_Float16)1152.0f}; + #pragma unroll + for (int d = 0; d < 4; ++d) { + const int row = vd + 4 * vh + d; + const h2v p01 = __builtin_bit_cast( + h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07010700u)); + const h2v p23 = __builtin_bit_cast( + h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07030702u)); + *(u2v*)&sV[bf][row * VST1 + t0 + 4 * vq] = + (u2v){__builtin_bit_cast(uint32_t, p01 - kofs), + __builtin_bit_cast(uint32_t, p23 - kofs)}; + } + if (lt < 16) { + const h2v a = __builtin_bit_cast(h2v, S.s0); + sKs[bf][0][t0 + lt] = (float)a[0]; + if constexpr (!KC) { + const h2v cc = __builtin_bit_cast(h2v, S.s1); + sKs[bf][1][t0 + lt] = (float)a[1]; + sKs[bf][2][t0 + lt] = (float)cc[0]; + sKs[bf][3][t0 + lt] = (float)cc[1]; + } + } else if (lt < 32) { + sVs[bf][t0 + lt - 16] = + (float)__builtin_bit_cast(_Float16, + (uint16_t)(KC ? S.s0 >> 16 : S.s0)) * + (PSCALE / LUT_ONE); + } + }; + + // ---- this lane's query vector, int8 over dims 32 s + 8 gq .. + 7 + const int qv = rowtile * QB1 + w * 16 + col; + const bool qlive = qv < nqv; + long Q[8]; + float qsc; + { + const int qtok = qlive ? qv / hpk : 0; + const int qhead = kvh * hpk + (qlive ? qv % hpk : 0); + const __half* qrow = + q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)qhead * sq_head; + _Float16 x[8][8]; + float qmax = 1e-8f; + #pragma unroll + for (int s = 0; s < 8; ++s) { + const u4v r = qlive ? *(const u4v*)(qrow + 32 * s + 8 * gq) : (u4v){}; + __builtin_memcpy(x[s], &r, 16); + #pragma unroll + for (int e = 0; e < 8; ++e) qmax = fmaxf(qmax, fabsf((float)x[s][e])); + } + qmax = qmax4(qmax); + const float qinv = 127.0f / qmax; + qsc = qmax * (1.0f / 127.0f); + #pragma unroll + for (int s = 0; s < 8; ++s) { + uint64_t word = 0; + #pragma unroll + for (int e = 0; e < 8; ++e) { + int v = (int)lrintf((float)x[s][e] * qinv); + v = max(-127, min(127, v)); + word |= ((uint64_t)(uint8_t)(int8_t)v) << (8 * e); + } + Q[s] = (long)word; + } + } + + f4v O[FR]; + #pragma unroll + for (int c = 0; c < FR; ++c) O[c] = (f4v){0, 0, 0, 0}; + const float c1 = sm_scale * qsc * (1.4426950408889634f / LUT_ONE); + float m = -INFINITY, l = 0.f; + + // Pipeline: iteration i computes tile i from LDS while the loads of tile + // i + 1 are in flight, then stores them into the other buffer. + Stg S; + if (ntiles > 0) { + fetch(0, S); + store(S, 0); + } + __syncthreads(); + for (int i = 0; i < ntiles; ++i) { + const int bf = i & 1; + if (i + 1 < ntiles) fetch(i + 1, S); + + // S^T = K x Q^T over both sub-tiles first (independent MFMA chains), + // then one online-softmax step over the 32 tokens, then PV. + f4v Sf[2]; + #pragma unroll + for (int j = 0; j < 2; ++j) { + Sf[j] = (f4v){0, 0, 0, 0}; + const int8_t* krow = &sK[bf][16 * j + col][8 * gq]; + const int tq = 16 * j + 4 * gq; + if constexpr (KC) { + i4v Si = {0, 0, 0, 0}; + #pragma unroll + for (int s = 0; s < 8; ++s) + Si = __builtin_amdgcn_mfma_i32_16x16x32_i8( + *(const long*)(krow + 32 * s), Q[s], Si, 0, 0, 0); + #pragma unroll + for (int e = 0; e < 4; ++e) + Sf[j][e] = sKs[bf][0][tq + e] * (float)Si[e]; + } else { + #pragma unroll + for (int b = 0; b < 4; ++b) { + i4v Si = {0, 0, 0, 0}; + #pragma unroll + for (int s = 2 * b; s < 2 * b + 2; ++s) + Si = __builtin_amdgcn_mfma_i32_16x16x32_i8( + *(const long*)(krow + 32 * s), Q[s], Si, 0, 0, 0); + #pragma unroll + for (int e = 0; e < 4; ++e) + Sf[j][e] += sKs[bf][b][tq + e] * (float)Si[e]; + } + } + } + + // Online softmax in log2 units; this lane holds tokens 16 j + 4 gq + e. + float sc[2][4]; + float mloc = -INFINITY; + #pragma unroll + for (int j = 0; j < 2; ++j) + #pragma unroll + for (int e = 0; e < 4; ++e) { + const int t = k_begin + i * KT1 + 16 * j + 4 * gq + e; + sc[j][e] = t < k_stop ? Sf[j][e] * c1 : -INFINITY; + mloc = fmaxf(mloc, sc[j][e]); + } + const float mn = fmaxf(m, qmax4(mloc)); + const float alpha = __builtin_amdgcn_exp2f(m - mn); + m = mn; + h4v pb[2]; + float lsum = 0.f; + #pragma unroll + for (int j = 0; j < 2; ++j) + #pragma unroll + for (int e = 0; e < 4; ++e) { + const bool live = sc[j][e] != -INFINITY; + const float p = live ? __builtin_amdgcn_exp2f(sc[j][e] - mn) : 0.f; + lsum += p; + pb[j][e] = + live ? (_Float16)(p * sVs[bf][16 * j + 4 * gq + e]) : (_Float16)0.f; + } + l = l * alpha + lsum; + if (__ballot(alpha != 1.0f)) { + #pragma unroll + for (int c = 0; c < FR; ++c) O[c] *= alpha; + } + + // O^T += V^T x P^T: A = V^T[dim 16 c + col][tokens 16 j + 4 gq .. + 3]. + #pragma unroll + for (int j = 0; j < 2; ++j) { + const _Float16* vrow = &sV[bf][col * VST1 + 16 * j + 4 * gq]; + #pragma unroll + for (int c = 0; c < FR; ++c) + O[c] = __builtin_amdgcn_mfma_f32_16x16x16f16( + *(const h4v*)(vrow + 16 * c * VST1), pb[j], O[c], 0, 0, 0); + } + + if (i + 1 < ntiles) store(S, bf ^ 1); + __syncthreads(); + } + + // ---- partials for the phase-2 merge (natural-log max, unnormalized O) + l = qsum4(l); + if (!qlive) return; + const int qtok = qv / hpk, qhead = kvh * hpk + qv % hpk; + float* wp = ws + (int64_t)split * total_q_tokens * num_q_heads * (D + 2) + + ((int64_t)(q0 + qtok) * num_q_heads + qhead) * (D + 2); + #pragma unroll + for (int c = 0; c < FR; ++c) + #pragma unroll + for (int e = 0; e < 4; ++e) + wp[16 * c + 4 * gq + e] = O[c][e] * (1.0f / PSCALE); + if (gq == 0) { + wp[D] = m * 0.6931471805599453f; + wp[D + 1] = l; + } + #endif // OCTAVE_CDNA_STUB +} + +// =========================================================================== +// Phase 2: the chunk's own K/V (fp16, causal), merged with the prefix +// partials. Like phase 1, one block = 128 query vectors (token, head) of one +// KV head, so the GQA heads share every K/V tile; 8 waves of 16 vectors. +// =========================================================================== +__global__ void __launch_bounds__(NT1) + chunk_attn(__half* __restrict__ out, const float* __restrict__ ws, + const __half* __restrict__ q, const __half* __restrict__ k, + const __half* __restrict__ v, + const int* __restrict__ cu_seqlens_q, + const int* __restrict__ seq_lens, int num_q_heads, + int num_kv_heads, int max_ctx, float sm_scale, int num_splits, + int total_q_tokens, int64_t sq_tok, int64_t sq_head, + int64_t sk_tok, int64_t sk_head, int64_t sv_tok, int64_t sv_head, + int64_t so_tok, int64_t so_head) { + #ifndef OCTAVE_CDNA_STUB + constexpr float LOG2E = 1.4426950408889634f; + const int seq = blockIdx.x, kvh = blockIdx.y, rowtile = blockIdx.z; + const int tid = threadIdx.x, w = tid >> 6, lane = tid & 63; + const int col = lane & 15, gq = lane >> 4; + const int hpk = num_q_heads / num_kv_heads; + const int q0 = cu_seqlens_q[seq]; + const int qlen = cu_seqlens_q[seq + 1] - q0; + const int nqv = qlen * hpk; + if (rowtile * QB1 >= nqv) return; + const int ctx = max(0, min(seq_lens[seq] - qlen, max_ctx)); + const int qv = rowtile * QB1 + w * 16 + col; + const bool qlive = qv < nqv; + const int qtok = qlive ? qv / hpk : 0; + const int head = kvh * hpk + (qlive ? qv % hpk : 0); + // Keys this block and this wave need (causal: up to their last token). + const int k_end = min(qlen, (min(nqv, (rowtile + 1) * QB1) - 1) / hpk + 1); + const int wave_first = rowtile * QB1 + w * 16; + const bool wave_live = wave_first < nqv; + const int wave_last_tok = min(nqv - 1, wave_first + 15) / hpk; + + // One 32-token tile in LDS (two would not fit in 64 KB with fp16 K and V); + // the next tile waits in registers. + __shared__ __attribute__((aligned(16))) _Float16 sK[KT1 * KROW2]; + __shared__ __attribute__((aligned(16))) _Float16 sV[D * VST2]; + + // Merge the prefix splits (log2 units from here on). O^T fragment c holds + // dims 16 c + 4 gq + e of this lane's vector; l stays partial per lane. + f4v acc[FR]; + #pragma unroll + for (int c = 0; c < FR; ++c) acc[c] = (f4v){0, 0, 0, 0}; + float m = -INFINITY, l = 0.f; + const int64_t split_stride = (int64_t)total_q_tokens * num_q_heads * (D + 2); + const int used = used_splits(ctx, num_splits); + if (qlive && used > 0) { + const float* w0 = + ws + ((int64_t)(q0 + qtok) * num_q_heads + head) * (D + 2); + for (int s = 0; s < used; ++s) + m = fmaxf(m, w0[s * split_stride + D] * LOG2E); + for (int s = 0; s < used; ++s) { + const float* wp = w0 + s * split_stride; + if (wp[D] == -INFINITY) continue; + const float wgt = __builtin_amdgcn_exp2f(wp[D] * LOG2E - m); + l += wp[D + 1] * wgt; + #pragma unroll + for (int c = 0; c < FR; ++c) + #pragma unroll + for (int e = 0; e < 4; ++e) acc[c][e] += wp[16 * c + 4 * gq + e] * wgt; + } + if (gq != 0) l = 0.f; + } + + // Q^T operand: dims 16 s + 4 gq .. + 3 of this lane's vector. + h4v Qh[FR]; + { + const __half* qrow = + q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)head * sq_head; + #pragma unroll + for (int s = 0; s < FR; ++s) Qh[s] = *(const h4v*)(qrow + 16 * s + 4 * gq); + } + const float c2 = sm_scale * LOG2E; + + // Loaders. K: thread t moves dims 8 (t % 32) .. + 7 of tokens t / 32 + // and t / 32 + 16 of the 32-token tile. V: dims 4 (t % 64) .. + 3 of + // tokens 4 (t / 64) .. + 3, transposed in registers into V^T rows. Rows + // past the chunk are zeros. + const int lt_tok = tid >> 5, lt_d = (tid & 31) * 8; + const int vt_d = (tid & 63) * 4, vt_g = tid >> 6; + const __half* kbase = k + (int64_t)kvh * sk_head + lt_d; + const __half* vbase = v + (int64_t)kvh * sv_head + vt_d; + u4v Kr[2]; + u2v Vr[4]; + auto fetch = [&](int n0) { + #pragma unroll + for (int h = 0; h < 2; ++h) { + const int t = n0 + lt_tok + 16 * h; + const bool ok = t < qlen; + const int64_t r = q0 + (ok ? t : 0); + Kr[h] = ok ? *(const u4v*)(kbase + r * sk_tok) : (u4v){0, 0, 0, 0}; + } + #pragma unroll + for (int e = 0; e < 4; ++e) { + const int t = n0 + 4 * vt_g + e; + const bool ok = t < qlen; + const int64_t r = q0 + (ok ? t : 0); + Vr[e] = ok ? *(const u2v*)(vbase + r * sv_tok) : (u2v){0, 0}; + } + }; + auto store = [&]() { + #pragma unroll + for (int h = 0; h < 2; ++h) { + _Float16* kp = &sK[(lt_tok + 16 * h) * KROW2 + lt_d]; + *(u2v*)kp = (u2v){Kr[h][0], Kr[h][1]}; + *(u2v*)(kp + 4) = (u2v){Kr[h][2], Kr[h][3]}; + } + #pragma unroll + for (int d = 0; d < 4; ++d) { + const uint32_t sel = (d & 1) ? 0x07060302u : 0x05040100u; + const uint32_t lo = + __builtin_amdgcn_perm(Vr[1][d >> 1], Vr[0][d >> 1], sel); + const uint32_t hi = + __builtin_amdgcn_perm(Vr[3][d >> 1], Vr[2][d >> 1], sel); + const int row = vt_d + d; + *(u2v*)&sV[row * VST2 + 4 * (vt_g ^ ((row >> 3) & 7))] = (u2v){lo, hi}; + } + }; + + const int ntiles = (k_end + KT1 - 1) / KT1; + fetch(0); + for (int i = 0; i < ntiles; ++i) { + const int n0 = i * KT1; + store(); + __syncthreads(); + if (i + 1 < ntiles) fetch(n0 + KT1); + if (wave_live && n0 <= wave_last_tok) { + // S^T for both 16-token halves, one softmax step, then PV. + f4v sacc[2]; + #pragma unroll + for (int j = 0; j < 2; ++j) { + sacc[j] = (f4v){0, 0, 0, 0}; + const _Float16* krow = &sK[(16 * j + col) * KROW2 + 4 * gq]; + #pragma unroll + for (int s = 0; s < FR; ++s) + sacc[j] = __builtin_amdgcn_mfma_f32_16x16x16f16( + *(const h4v*)(krow + 16 * s), Qh[s], sacc[j], 0, 0, 0); + } + float sc[2][4]; + float mloc = -INFINITY; + #pragma unroll + for (int j = 0; j < 2; ++j) + #pragma unroll + for (int e = 0; e < 4; ++e) { + const int kabs = n0 + 16 * j + 4 * gq + e; + sc[j][e] = qlive && kabs <= qtok ? sacc[j][e] * c2 : -INFINITY; + mloc = fmaxf(mloc, sc[j][e]); + } + const float mn = fmaxf(m, qmax4(mloc)); + const float al = mn == -INFINITY ? 1.f : __builtin_amdgcn_exp2f(m - mn); + h4v pb[2]; + float lsum = 0.f; + #pragma unroll + for (int j = 0; j < 2; ++j) + #pragma unroll + for (int e = 0; e < 4; ++e) { + const float p = sc[j][e] == -INFINITY + ? 0.f + : __builtin_amdgcn_exp2f(sc[j][e] - mn); + lsum += p; + pb[j][e] = (_Float16)p; + } + m = mn; + l = l * al + lsum; + if (__ballot(al != 1.0f)) { + #pragma unroll + for (int c = 0; c < FR; ++c) acc[c] *= al; + } + #pragma unroll + for (int j = 0; j < 2; ++j) { + // Row 16 c + col; its group 4 j + gq sits at (4 j + gq) ^ (row >> 3). + #pragma unroll + for (int c = 0; c < FR; ++c) { + const int row = 16 * c + col; + const _Float16* vp = + &sV[row * VST2 + 4 * ((4 * j + gq) ^ ((row >> 3) & 7))]; + acc[c] = __builtin_amdgcn_mfma_f32_16x16x16f16(*(const h4v*)vp, pb[j], + acc[c], 0, 0, 0); + } + } + } + __syncthreads(); + } + + l = qsum4(l); + if (!qlive) return; + const float inv = 1.0f / (l + 1e-10f); + __half* orow = out + (int64_t)(q0 + qtok) * so_tok + (int64_t)head * so_head; + #pragma unroll + for (int c = 0; c < FR; ++c) + #pragma unroll + for (int e = 0; e < 4; ++e) + orow[16 * c + 4 * gq + e] = __float2half(acc[c][e] * inv); + #endif // OCTAVE_CDNA_STUB +} + +} // namespace octave_cdna + +// Called by octave_prefill (octave_prefill_rdna3.cu) on gfx942; same +// arguments, already validated there. +void octave_prefill_cdna(torch::Tensor out, torch::Tensor q, torch::Tensor k, + torch::Tensor v, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor cu_seqlens_q, + torch::Tensor seq_lens, int64_t max_query_len, + double sm_scale, int64_t fmt, float* ws, + int num_splits) { + using namespace octave_cdna; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const int num_seqs = seq_lens.size(0); + const int hq = q.size(1), hkv = cache.size(1); + const int hpk = hq / hkv; + const int block_size = cache.size(2); + const int max_blocks = block_table.size(1); + const int total_q = q.size(0); + const int rowtiles = (int)((max_query_len * hpk + QB1 - 1) / QB1); + + dim3 grid1(num_seqs, hkv, rowtiles * num_splits); + #define PREFIX(B, KC) \ + prefix_attn<<>>( \ + (const __half*)q.data_ptr(), cache.data_ptr(), \ + block_table.data_ptr(), cu_seqlens_q.data_ptr(), \ + seq_lens.data_ptr(), ws, (float)sm_scale, hq, hkv, block_size, \ + max_blocks, num_splits, total_q, q.stride(0), q.stride(1), \ + cache.stride(0), cache.stride(1), cache.stride(2)) + if (fmt == 4) + PREFIX(4, false); + else if (fmt == 3) + PREFIX(3, false); + else + PREFIX(3, true); + #undef PREFIX + dim3 grid2(num_seqs, hkv, rowtiles); + chunk_attn<<>>( + (__half*)out.data_ptr(), ws, (const __half*)q.data_ptr(), + (const __half*)k.data_ptr(), (const __half*)v.data_ptr(), + cu_seqlens_q.data_ptr(), seq_lens.data_ptr(), hq, hkv, + max_blocks * block_size, (float)sm_scale, num_splits, total_q, + q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0), + v.stride(1), out.stride(0), out.stride(1)); +} + +#endif // USE_ROCM diff --git a/csrc/attention/octave_prefill_rdna3.cu b/csrc/attention/octave_prefill_rdna3.cu new file mode 100644 index 000000000000..3566bf62626f --- /dev/null +++ b/csrc/attention/octave_prefill_rdna3.cu @@ -0,0 +1,809 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Octave prefill for RDNA3 (gfx11), head size 256, fp16. Reads the packed +// cache directly; no intermediate buffers. Everything is in the rotated space +// (the caller rotates Q/K/V of the chunk and rotates the output back). +// +// Phase 1 (prefix_attn): attention of the chunk's queries over the cached +// prefix, split over the prefix length; writes (O, m, l) partials. +// One block = 128 query vectors (token, head) of one KV head, so the GQA +// heads share every K/V tile. 16 compute waves in pairs: wave g of a pair +// holds Q (int8) and O for head dims [128 g, 128 g + 128). 8 loader waves +// unpack 16-token tiles into double-buffered LDS, one tile ahead. +// QK S^T[token][q] = K x Q^T with int8 WMMA. K rows hold the int8 codebook +// value of every dim. Each wave's 128 dims are two 64-dim blocks with +// their own per-token scale; the wave scales both and the pair +// exchanges the float partial through LDS. +// PV O^T[dim][q] = V^T x P^T with V as its exact fp16 codebook values; the +// V scale is folded into P. +// Phase 2 (chunk_attn): causal attention over the chunk's own K/V (fp16), +// merged with the phase-1 partials; writes the output. + +#include +#include +#include +#include +#include + +#include +#include +#include + +#if defined(USE_ROCM) + #include + #include + + #include "octave_format.cuh" + + #if defined(__HIP_DEVICE_COMPILE__) && !defined(__gfx1100__) + #define OCTAVE_PREFILL_STUB + #endif + +namespace octave_pf { + +using octave::D; +using octave::Format; +using octave::LUT_ONE; + +typedef _Float16 h16v __attribute__((ext_vector_type(16))); +typedef _Float16 h8v __attribute__((ext_vector_type(8))); +typedef _Float16 h2v __attribute__((ext_vector_type(2))); +typedef float f8v __attribute__((ext_vector_type(8))); +typedef uint32_t u8v __attribute__((ext_vector_type(8))); +typedef uint32_t u4v __attribute__((ext_vector_type(4))); +typedef uint32_t u2v __attribute__((ext_vector_type(2))); +typedef int v4i_t __attribute__((ext_vector_type(4))); +typedef int v8i_t __attribute__((ext_vector_type(8))); + +constexpr int KT = 16; // tokens per tile + +// Splits of the cached prefix used by a sequence: one per kMinSplitTokens, at +// most num_splits. Both phases must agree. +constexpr int kMinSplitTokens = 2048; +__device__ __forceinline__ int used_splits(int ctx, int num_splits) { + if (ctx <= 0) return 0; + return min(num_splits, (ctx + kMinSplitTokens - 1) / kMinSplitTokens); +} +__device__ __forceinline__ int split_len(int ctx, int used) { + return ((ctx + used - 1) / used + KT - 1) / KT * KT; +} + +// =========================================================================== +// Phase 1: prefix +// =========================================================================== +constexpr int NW = 8; // compute wave pairs (16 query vectors each) +constexpr int NL = 8; // loader waves +[[maybe_unused]] constexpr int KROW = D + 16; // bytes per K row in LDS +[[maybe_unused]] constexpr float PSCALE = 4096.0f; +// The LDS read of PV step c waits for the WMMA of step c - DEPV, which keeps +// the scheduler from hoisting every read (and spilling). +[[maybe_unused]] constexpr int DEPV = 4; +constexpr int NTHREADS = (2 * NW + NL) * 32; +constexpr int VPAD = 16; // halves after every 8 rows of V^T +constexpr int HD = D / 2; +[[maybe_unused]] constexpr int QSTEPS = HD / 16; +[[maybe_unused]] constexpr int OFRAGS = HD / 16; + +__device__ __forceinline__ int vrow(int row) { + return row * KT + (row >> 3) * VPAD; +} +__device__ __forceinline__ uint32_t xhalf(uint32_t v) { + return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false, + false); +} +__device__ __forceinline__ float xhalf(float v) { + return __uint_as_float(xhalf(__float_as_uint(v))); +} +// 4 dwords, one per token (same 4 dims) -> 4 dwords, one per dim (4 tokens). +__device__ __forceinline__ void tr4(uint32_t a0, uint32_t a1, uint32_t a2, + uint32_t a3, uint32_t o[4]) { + const uint32_t lo01 = __builtin_amdgcn_perm(a1, a0, 0x05010400u); + const uint32_t hi01 = __builtin_amdgcn_perm(a1, a0, 0x07030602u); + const uint32_t lo23 = __builtin_amdgcn_perm(a3, a2, 0x05010400u); + const uint32_t hi23 = __builtin_amdgcn_perm(a3, a2, 0x07030602u); + o[0] = __builtin_amdgcn_perm(lo23, lo01, 0x05040100u); + o[1] = __builtin_amdgcn_perm(lo23, lo01, 0x07060302u); + o[2] = __builtin_amdgcn_perm(hi23, hi01, 0x05040100u); + o[3] = __builtin_amdgcn_perm(hi23, hi01, 0x07060302u); +} +// LDS-only barrier: __syncthreads() would also wait for the loaders' global +// prefetch. +__device__ __forceinline__ void tile_barrier() { + asm volatile("s_waitcnt lgkmcnt(0)\n\ts_barrier" ::: "memory"); +} + +// Codes of dims 4k..4k+3 of a 64-dim block, one byte each. 4-bit: raw is +// the block's 32 bytes. 3-bit: raw starts at the block's 16 bytes of the +// 2-bit plane; hi0/hi1 are its two dwords of the 1-bit plane. +template +__device__ __forceinline__ uint32_t block_codes(const uint32_t* raw, + uint32_t hi0, uint32_t hi1, + int k) { + if constexpr (BITS == 4) { + return (raw[k >> 1] >> (4 * (k & 1))) & 0x0F0F0F0Fu; + } else { + const uint32_t hi = k < 8 ? hi0 : hi1; + return ((raw[k >> 2] >> (2 * (k & 3))) & 0x03030303u) | + (((hi >> (k & 7)) & 0x01010101u) << 2); + } +} + +template +__global__ __launch_bounds__(NTHREADS) void prefix_attn( + const __half* __restrict__ q, const uint8_t* __restrict__ cache, + const int* __restrict__ block_table, const int* __restrict__ cu_seqlens_q, + const int* __restrict__ seq_lens, float* __restrict__ ws, float sm_scale, + int num_q_heads, int num_kv_heads, int block_size, int max_blocks, + int num_splits, int total_q_tokens, int64_t sq_tok, int64_t sq_head, + int64_t scb, int64_t sch, int64_t sct) { + #ifndef OCTAVE_PREFILL_STUB + using F = Format; + const int seq = blockIdx.x, kvh = blockIdx.y; + const int rowtile = blockIdx.z / num_splits, split = blockIdx.z % num_splits; + const int tid = threadIdx.x; + const int w = __builtin_amdgcn_readfirstlane(tid >> 5), lane = tid & 31; + const int j = lane & 15, h = lane >> 4; + const int pr = w >> 1, g = w & 1; + const int hpk = num_q_heads / num_kv_heads; + const bool loader = w >= 2 * NW; + + const int q0 = cu_seqlens_q[seq]; + const int qlen = cu_seqlens_q[seq + 1] - q0; + int ctx = seq_lens[seq] - qlen; + ctx = max(0, min(ctx, max_blocks * block_size)); + const int nqv = qlen * hpk; + if (rowtile * NW * 16 >= nqv) return; + + const int used = used_splits(ctx, num_splits); + if (split >= used) return; + const int k_len = split_len(ctx, used); + const int k_begin = min(ctx, split * k_len); + const int k_stop = min(ctx, k_begin + k_len); + const int ntiles = k_stop > k_begin ? (k_stop - k_begin + KT - 1) / KT : 0; + + __shared__ __attribute__((aligned(16))) int8_t sK[2][KT][KROW]; + __shared__ + __attribute__((aligned(16))) _Float16 sV[2][D * KT + (D / 8) * VPAD]; + __shared__ __attribute__((aligned(16))) u4v sS[2][2 * NW][2][32]; + // K scales travel with K (block 0 = RoPE, 1-3 = NoPE), V scales with V. + __shared__ float sKs[2][4][KT], sVs[2][KT]; + + // ---- loader waves + if (loader) { + // 256 threads. K: thread lt unpacks dims 16 kp .. 16 kp + 15 of block kq + // of token tk. V: dims 4 vh .. 4 vh + 3 of the 8-dim group vd of token + // quad vq. + const int lt = tid - 2 * NW * 32; + const int tk = lt >> 4, kq = (lt >> 2) & 3, kp = lt & 3; + const int vu = lt >> 1, vh = lt & 1, vq = vu & 3, vd = (vu >> 2) * 8; + // Block sizes are multiples of KT, so a tile sits in one cache block. + // Tiles past the end read the last one; tokens past k_stop read the last + // token: their K is masked to -inf and their V meets p = 0. + const int tlast = k_stop - 1; + const uint8_t* const head = cache + (int64_t)kvh * sch; + const int* const bt = block_table + seq * max_blocks; + // Walk the tiles in order: block index and offset inside it, advanced by + // KT per tile and frozen on the last tile. + struct Pos { + int start, blk, off; + }; + auto first = [&]() { + return Pos{k_begin, k_begin / block_size, k_begin % block_size}; + }; + auto next = [&](Pos p) { + if (p.start + KT > tlast) return p; + p.start += KT; + p.off += KT; + if (p.off == block_size) { + p.off = 0; + ++p.blk; + } + return p; + }; + auto pbase = [&](Pos p, int pb) { + return head + (int64_t)(pb < 0 ? 0 : pb) * scb + (int64_t)p.off * sct; + }; + auto tok = [&](int start, int r) { + return (int64_t)min(r, tlast - start) * sct; + }; + // 4-bit: dwords 2 kp, 2 kp + 1 of the RoPE codes. 3-bit: dword kp of the + // block's 2-bit plane and dword kp / 2 of its 1-bit plane. + const bool k3 = KC || kq > 0; + const int b = KC ? kq : kq - 1; // 3-bit block index + const int koff0 = k3 ? F::OFF_KN + 16 * b + 4 * kp : 8 * kp; + const int koff1 = + k3 ? F::OFF_KN + F::KN / 4 + 8 * b + 4 * (kp >> 1) : 8 * kp + 4; + constexpr int VW = VB == 4 ? 1 : 2; // dwords of 8 V dims + struct Stg { + uint32_t k0, k1; + uint32_t v[4][VW]; + uint32_t s0, s1; + }; + auto fetch = [&](int TK, const uint8_t* bK, int TV, const uint8_t* bV, + Stg& S) { + const uint8_t* pk = bK + tok(TK, tk); + S.k0 = *(const uint32_t*)(pk + koff0); + S.k1 = *(const uint32_t*)(pk + koff1); + #pragma unroll + for (int i = 0; i < 4; ++i) { + const uint8_t* pv = bV + tok(TV, 4 * vq + i) + F::OFF_V; + if constexpr (VB == 4) { + S.v[i][0] = *(const uint32_t*)(pv + vd / 2); + } else { + S.v[i][0] = *(const uint32_t*)(pv + 4 * (vd >> 4)); + S.v[i][1] = *(const uint32_t*)(pv + D / 4 + 4 * (vd >> 5)); + } + } + // Threads 0-15: the K scales of token lt; 16-31: the V scale of token + // lt - KT. Block K: two 8-byte loads inside the slot, V scale in the + // low half of s0. Compact K: one 4-byte load {k, v}. + if (lt < 2 * KT) { + if constexpr (KC) { + S.s0 = *(const uint32_t*)((lt < KT ? bK + tok(TK, lt) + : bV + tok(TV, lt - KT)) + + F::OFF_SC); + } else { + const uint8_t* ps = lt < KT ? bK + tok(TK, lt) + F::OFF_SC + : bV + tok(TV, lt - KT) + F::OFF_SC + 8; + const u2v x = *(const u2v*)ps; + S.s0 = x[0]; + S.s1 = x[1]; + } + } + }; + auto store_k = [&](const Stg& S, int bf) { + const uint32_t r4[2] = {S.k0, S.k1}; + const uint32_t r3[1] = {S.k0}; + const uint32_t hi = S.k1 >> (4 * (kp & 1)); + u4v o; + #pragma unroll + for (int x = 0; x < 4; ++x) + o[x] = k3 ? octave::lut3(block_codes<3>(r3, hi, hi, x)) + : octave::lut4(block_codes<4>(r4, 0, 0, x)); + *(u4v*)&sK[bf][tk][64 * kq + 16 * kp] = o; + if (lt < KT) { + const h2v a = __builtin_bit_cast(h2v, S.s0); + if constexpr (KC) { + sKs[bf][0][lt] = (float)a[0]; + } else { + const h2v c = __builtin_bit_cast(h2v, S.s1); + sKs[bf][0][lt] = (float)a[0]; + sKs[bf][1][lt] = (float)a[1]; + sKs[bf][2][lt] = (float)c[0]; + sKs[bf][3][lt] = (float)c[1]; + } + } + }; + auto store_v = [&](const Stg& S, int bf) { + uint32_t c[4]; // [token] code bytes of dims vd + 4 vh .. + #pragma unroll + for (int i = 0; i < 4; ++i) { + if constexpr (VB == 4) { + c[i] = (S.v[i][0] >> (4 * vh)) & 0x0F0F0F0Fu; + } else { + const int sl = 2 * ((vd & 15) >> 2) + 2 * vh; + const int sh = ((vd & 31) >> 2) + vh; + c[i] = ((S.v[i][0] >> sl) & 0x03030303u) | + (((S.v[i][1] >> sh) & 0x01010101u) << 2); + } + c[i] = octave::lut(c[i]) ^ 0x80808080u; + } + uint32_t t[4]; + tr4(c[0], c[1], c[2], c[3], t); + // Codebook bytes fed as fp16 1024 + (x + 128), less 1152: exact x. + const h2v kofs = {(_Float16)1152.0f, (_Float16)1152.0f}; + #pragma unroll + for (int d = 0; d < 4; ++d) { + const int row = vd + 4 * vh + d, sw = ((row >> 3) & 1) * 8; + const h2v p01 = __builtin_bit_cast( + h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07010700u)); + const h2v p23 = __builtin_bit_cast( + h2v, __builtin_amdgcn_perm(0x64646464u, t[d], 0x07030702u)); + *(u2v*)&sV[bf][vrow(row) + ((4 * vq) ^ sw)] = + (u2v){__builtin_bit_cast(uint32_t, p01 - kofs), + __builtin_bit_cast(uint32_t, p23 - kofs)}; + } + if (lt >= KT && lt < 2 * KT) + sVs[bf][lt - KT] = (float)__builtin_bit_cast( + _Float16, (uint16_t)(KC ? S.s0 >> 16 : S.s0)) * + (PSCALE / LUT_ONE); + }; + // Pipeline: K of tile i + 1 and V of tile i are written in iteration i + // from registers; the block number of a tile is loaded two tiles ahead. + Stg S; + if (ntiles > 0) { + const Pos p0 = first(), p1 = next(p0), p2 = next(p1); + const uint8_t* b0 = pbase(p0, bt[p0.blk]); + const int q1 = bt[p1.blk]; + fetch(p0.start, b0, p0.start, b0, S); + const int q2 = bt[p2.blk]; + Pos p3 = next(p2); + int pn = bt[p3.blk]; // block of tile i + 3 + const uint8_t* bv = pbase(p1, q1); + const uint8_t* bk = pbase(p2, q2); + int sv = p1.start, sk = p2.start; + store_k(S, 0); + fetch(p1.start, bv, p0.start, b0, S); + tile_barrier(); + for (int i = 0; i < ntiles; ++i) { + store_k(S, (i + 1) & 1); + store_v(S, i & 1); + fetch(sk, bk, sv, bv, S); + bv = bk; + sv = sk; + bk = pbase(p3, pn); + sk = p3.start; + p3 = next(p3); + pn = bt[p3.blk]; + tile_barrier(); + } + } else { + tile_barrier(); + } + tile_barrier(); + return; + } + + // ---- compute waves: this lane's query row, int8 over this wave's dims + const int qr = pr * 16 + j; + const int qv = rowtile * NW * 16 + qr; + const bool qlive = qv < nqv; + u4v Q[QSTEPS]; + float qsc; + { + const int qtok = qlive ? qv / hpk : 0; + const int qhead = kvh * hpk + (qlive ? qv % hpk : 0); + const __half* qrow = + q + (int64_t)(q0 + qtok) * sq_tok + (int64_t)qhead * sq_head; + float qmax = 1e-8f; + for (int i = 0; i < D / 8; ++i) { + const h8v x = qlive ? *(const h8v*)(qrow + 8 * i) : (h8v){}; + #pragma unroll + for (int k = 0; k < 8; ++k) qmax = fmaxf(qmax, fabsf((float)x[k])); + } + const float qinv = 127.0f / qmax; + qsc = qmax * (1.0f / 127.0f); + #pragma unroll + for (int f = 0; f < QSTEPS; ++f) { + uint32_t wd[4]; + #pragma unroll + for (int hh = 0; hh < 2; ++hh) { + const h8v x = + qlive ? *(const h8v*)(qrow + HD * g + 16 * f + 8 * hh) : (h8v){}; + #pragma unroll + for (int d4 = 0; d4 < 2; ++d4) { + uint32_t word = 0; + #pragma unroll + for (int b = 0; b < 4; ++b) { + int v = (int)lrintf((float)x[4 * d4 + b] * qinv); + v = max(-127, min(127, v)); + word |= ((uint32_t)(uint8_t)(int8_t)v) << (8 * b); + } + wd[2 * hh + d4] = word; + } + } + Q[f] = (u4v){wd[0], wd[1], wd[2], wd[3]}; + } + } + + f8v O[OFRAGS]; + #pragma unroll + for (int c = 0; c < OFRAGS; ++c) O[c] = (f8v){0, 0, 0, 0, 0, 0, 0, 0}; + // A operands are only read in lanes with (j % 2) == h; the other lane of + // each pair reads its partner's row, so the pair costs one LDS fetch. + const int rr = ((j & 1) == h) ? j : (j ^ 1); + const int sw = ((rr >> 3) & 1) * 8; + const float c1 = sm_scale * qsc * (1.4426950408889634f / LUT_ONE); + float m = -INFINITY, l = 0.f; + f8v Sp = {0, 0, 0, 0, 0, 0, 0, 0}; // this wave's partial S, tile i-1 + + tile_barrier(); + for (int i = 0; i <= ntiles; ++i) { + f8v Sn = {0, 0, 0, 0, 0, 0, 0, 0}; + if (i < ntiles) { + // QK over this wave's two 64-dim blocks, S^T = K x Q^T, 4 k-steps + // each; block s is scaled into Sn before block s + 1 starts, so one + // int accumulator is live at a time. + uint32_t koff = (uint32_t)(rr * KROW + HD * g); + if constexpr (KC) { // one K scale: one accumulator for both blocks + v8i_t Si = {0, 0, 0, 0, 0, 0, 0, 0}; + #pragma unroll + for (int f = 0; f < 8; ++f) { + const u4v a = *(const u4v*)(&sK[i & 1][0][0] + koff + 16 * f); + Si = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32( + true, __builtin_bit_cast(v4i_t, a), true, + __builtin_bit_cast(v4i_t, Q[f]), Si, false); + } + const float* ks = &sKs[i & 1][0][0]; + #pragma unroll + for (int e = 0; e < 8; ++e) Sn[e] = ks[2 * e + h] * (float)Si[e]; + } + #pragma unroll + for (int sg = 0; sg < (KC ? 0 : 2); ++sg) { + v8i_t Si = {0, 0, 0, 0, 0, 0, 0, 0}; + #pragma unroll + for (int f = 4 * sg; f < 4 * sg + 4; ++f) { + const u4v a = *(const u4v*)(&sK[i & 1][0][0] + koff + 16 * f); + Si = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32( + true, __builtin_bit_cast(v4i_t, a), true, + __builtin_bit_cast(v4i_t, Q[f]), Si, false); + } + const float* ks = &sKs[i & 1][2 * g + sg][0]; + #pragma unroll + for (int e = 0; e < 8; ++e) Sn[e] += ks[2 * e + h] * (float)Si[e]; + if (sg == 0) asm volatile("" : "+v"(koff) : "v"(Sn[0])); + } + const u8v su = __builtin_bit_cast(u8v, Sn); + sS[i & 1][w][0][lane] = (u4v){su[0], su[1], su[2], su[3]}; + sS[i & 1][w][1][lane] = (u4v){su[4], su[5], su[6], su[7]}; + } + if (i > 0) { + const int bf = (i - 1) & 1, base = k_begin + (i - 1) * KT; + const u4v o0 = sS[bf][w ^ 1][0][lane], o1 = sS[bf][w ^ 1][1][lane]; + const f8v S = + Sp + __builtin_bit_cast(f8v, (u8v){o0[0], o0[1], o0[2], o0[3], o1[0], + o1[1], o1[2], o1[3]}); + // Scores in log2 units are S * c1 (c1 > 0). Only the last tile has + // keys past k_stop; exp2(-inf) = 0, and m = -inf only before the first + // tile, so nothing else needs a guard. + float s[8]; + #pragma unroll + for (int e = 0; e < 8; ++e) s[e] = S[e]; + if (base + KT > k_stop) { + #pragma unroll + for (int e = 0; e < 8; ++e) + if (base + 2 * e + h >= k_stop) s[e] = -INFINITY; + } + float mloc = fmaxf(fmaxf(fmaxf(s[0], s[1]), fmaxf(s[2], s[3])), + fmaxf(fmaxf(s[4], s[5]), fmaxf(s[6], s[7]))); + const float mn = fmaxf(m, fmaxf(mloc, xhalf(mloc)) * c1); + const float alpha = __builtin_amdgcn_exp2f(m - mn); + m = mn; + uint32_t pd[4]; + float lsum = 0.f; + #pragma unroll + for (int e = 0; e < 8; e += 2) { + const float p0 = __builtin_amdgcn_exp2f(fmaf(s[e], c1, -mn)); + const float p1 = __builtin_amdgcn_exp2f(fmaf(s[e + 1], c1, -mn)); + lsum += p0 + p1; + const _Float16 h0 = (_Float16)(p0 * sVs[bf][2 * e + h]); + const _Float16 h1 = (_Float16)(p1 * sVs[bf][2 * e + 2 + h]); + pd[e >> 1] = __builtin_bit_cast(uint32_t, (h2v){h0, h1}); + } + l = l * alpha + lsum; + if (__builtin_amdgcn_ballot_w32(alpha != 1.0f)) { + #pragma unroll + for (int c = 0; c < OFRAGS; ++c) O[c] *= alpha; + } + u8v pf; + // B operand in natural token order: lane half 0 holds the even tokens, + // half 1 the odd ones; interleave them per pair. + #pragma unroll + for (int k = 0; k < 4; ++k) { + const uint32_t o = xhalf(pd[k]); + const uint32_t ev = h ? o : pd[k], od = h ? pd[k] : o; + pf[2 * k] = __builtin_amdgcn_perm(od, ev, 0x05040100u); + pf[2 * k + 1] = __builtin_amdgcn_perm(od, ev, 0x07060302u); + } + const h16v pb16 = __builtin_bit_cast(h16v, pf); + #pragma unroll + for (int c = 0; c < OFRAGS; ++c) { + uint32_t voff = (uint32_t)vrow(HD * g + 16 * c + rr); + if (c >= DEPV) asm volatile("" : "+v"(voff) : "v"(O[c - DEPV][0])); + const _Float16* vr = &sV[bf][0] + voff; + const h8v lo = *(const h8v*)(vr + sw); + const h8v hi = *(const h8v*)(vr + (8 ^ sw)); + const h16v a = __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, + 8, 9, 10, 11, 12, 13, 14, 15); + O[c] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, pb16, O[c]); + } + } + Sp = Sn; + tile_barrier(); + } + + // ---- partials for the phase-2 merge (natural-log max, unnormalized O) + l += xhalf(l); + if (!qlive) return; + const int qtok = qv / hpk, qhead = kvh * hpk + qv % hpk; + float* wp = ws + (int64_t)split * total_q_tokens * num_q_heads * (D + 2) + + ((int64_t)(q0 + qtok) * num_q_heads + qhead) * (D + 2) + HD * g; + #pragma unroll + for (int c = 0; c < OFRAGS; ++c) + #pragma unroll + for (int e = 0; e < 8; ++e) + wp[16 * c + 2 * e + h] = O[c][e] * (1.0f / PSCALE); + if (g == 0 && h == 0) { + wp[D] = m * 0.6931471805599453f; + wp[D + 1] = l; + } + #endif // OCTAVE_PREFILL_STUB +} + +// =========================================================================== +// Phase 2: the chunk's own K/V (fp16, causal), merged with the prefix +// partials. 8 waves, 16 query rows each, one query head per block. +// =========================================================================== +constexpr int CW = 8; // waves +constexpr int CM = 16 * CW; // query rows per block +[[maybe_unused]] constexpr int FRAGS = D / 16; +[[maybe_unused]] constexpr int X = 8; // fp16 per 16-byte chunk + +__device__ __forceinline__ float row16_max(float v) { + v = fmaxf(v, __shfl_xor(v, 1)); + v = fmaxf(v, __shfl_xor(v, 2)); + v = fmaxf(v, __shfl_xor(v, 4)); + return fmaxf(v, __shfl_xor(v, 8)); +} +__device__ __forceinline__ float row16_sum(float v) { + v += __shfl_xor(v, 1); + v += __shfl_xor(v, 2); + v += __shfl_xor(v, 4); + return v + __shfl_xor(v, 8); +} + +__global__ void __launch_bounds__(D) + chunk_attn(__half* __restrict__ out, const float* __restrict__ ws, + const __half* __restrict__ q, const __half* __restrict__ k, + const __half* __restrict__ v, + const int* __restrict__ cu_seqlens_q, + const int* __restrict__ seq_lens, int num_q_heads, + int num_kv_heads, int max_ctx, float sm_scale, int num_splits, + int total_q_tokens, int64_t sq_tok, int64_t sq_head, + int64_t sk_tok, int64_t sk_head, int64_t sv_tok, int64_t sv_head, + int64_t so_tok, int64_t so_head) { + #ifndef OCTAVE_PREFILL_STUB + const int seq = blockIdx.x, head = blockIdx.y; + const int tid = threadIdx.x, wave = tid >> 5, lane = tid & 31; + const int lo = lane & 15, hi = lane >> 4; + const int q0 = cu_seqlens_q[seq]; + const int qlen = cu_seqlens_q[seq + 1] - q0; + const int ctx = max(0, min(seq_lens[seq] - qlen, max_ctx)); + const int tile0 = blockIdx.z * CM; + if (tile0 >= qlen) return; + const int wq0 = tile0 + wave * 16; + const int nq_wave = max(0, min(16, qlen - wq0)); + const int kvh = head / (num_q_heads / num_kv_heads); + + __shared__ __attribute__((aligned(16))) _Float16 sK[FRAGS * 2 * KT * X]; + __shared__ __attribute__((aligned(16))) _Float16 sV[D * KT]; + __shared__ __attribute__((aligned(16))) _Float16 sP[CW][16 * KT]; + + float m_st[8], l_st[8]; + f8v acc[FRAGS]; + // Merge the prefix splits (one weight per split; exp(0) = 1 with one). + const int64_t split_stride = (int64_t)total_q_tokens * num_q_heads * (D + 2); + const int used = used_splits(ctx, num_splits); + #pragma unroll + for (int i = 0; i < 8; ++i) { + m_st[i] = -INFINITY; + l_st[i] = 0.f; + } + #pragma unroll + for (int f = 0; f < FRAGS; ++f) acc[f] = (f8v){0, 0, 0, 0, 0, 0, 0, 0}; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const int qp = wq0 + 2 * i + hi; + if (qp >= qlen) continue; + const float* w0 = ws + ((int64_t)(q0 + qp) * num_q_heads + head) * (D + 2); + float mx = -INFINITY; + for (int s = 0; s < used; ++s) mx = fmaxf(mx, w0[s * split_stride + D]); + float lsum = 0.f; + for (int s = 0; s < used; ++s) { + const float* wp = w0 + s * split_stride; + if (wp[D] == -INFINITY) continue; + const float wgt = __expf(wp[D] - mx); + lsum += wp[D + 1] * wgt; + #pragma unroll + for (int f = 0; f < FRAGS; ++f) acc[f][i] += wp[16 * f + lo] * wgt; + } + m_st[i] = mx; + l_st[i] = lsum; + } + + // Q rows are re-read from global (L1/L2) in the loop instead of being held + // in registers; out-of-range rows are masked anyway. + const __half* qrow = q + (int64_t)(q0 + min(wq0 + lo, qlen - 1)) * sq_tok + + (int64_t)head * sq_head; + const int k_end = min(qlen, tile0 + min(CM, qlen - tile0)); + _Float16* sPw = &sP[wave][0]; + for (int n0 = 0; n0 < k_end; n0 += KT) { + { // K tile: [16-dim chunk][token][8] fp16 + const int tk = tid / (FRAGS), dc = (tid % FRAGS) * 2; + const bool ok = n0 + tk < qlen; + const __half* row = + k + (int64_t)(q0 + n0 + tk) * sk_tok + (int64_t)kvh * sk_head; + #pragma unroll + for (int d = 0; d < 2; ++d) { + int4 x = ok ? *(const int4*)(row + (dc + d) * X) : int4{0, 0, 0, 0}; + *(int4*)&sK[(dc + d) * (KT * X) + tk * X] = x; + } + } + { // V tile transposed: [dim][token] + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int tk = tid / (D / KT), + dbase = ((tid % (D / KT)) + p * (D / KT)) * 8; + const bool ok = n0 + tk < qlen; + int4 x = ok ? *(const int4*)(v + (int64_t)(q0 + n0 + tk) * sv_tok + + (int64_t)kvh * sv_head + dbase) + : int4{0, 0, 0, 0}; + _Float16 t8[8]; + __builtin_memcpy(t8, &x, 16); + #pragma unroll + for (int e = 0; e < 8; ++e) sV[(dbase + e) * KT + tk] = t8[e]; + } + } + __syncthreads(); + if (nq_wave > 0 && n0 <= wq0 + nq_wave - 1) { + f8v s = {0, 0, 0, 0, 0, 0, 0, 0}; + #pragma unroll 2 + for (int f = 0; f < FRAGS; ++f) { + h16v b, a; + const int4 klo = *(const int4*)&sK[(2 * f) * (KT * X) + lo * X]; + const int4 khi = *(const int4*)&sK[(2 * f + 1) * (KT * X) + lo * X]; + __builtin_memcpy(&b, &klo, 16); + __builtin_memcpy(((char*)&b) + 16, &khi, 16); + __builtin_memcpy(&a, qrow + 16 * f, 32); + s = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, s); + } + const int kabs = n0 + lo; + float p[8]; + #pragma unroll + for (int i = 0; i < 8; ++i) { + const int r = 2 * i + hi; + const bool keep = r < nq_wave && kabs < qlen && kabs <= wq0 + r; + const float sc = keep ? s[i] * sm_scale : -INFINITY; + const float mi = row16_max(sc); + const float mn = fmaxf(m_st[i], mi); + const float al = m_st[i] == -INFINITY ? 0.f : __expf(m_st[i] - mn); + p[i] = mn == -INFINITY ? 0.f : __expf(sc - mn); + l_st[i] = l_st[i] * al + row16_sum(p[i]); + m_st[i] = mn; + #pragma unroll + for (int f = 0; f < FRAGS; ++f) acc[f][i] *= al; + sPw[r * KT + lo] = (_Float16)p[i]; + } + h16v pa; + const int4 plo = *(const int4*)&sPw[lo * KT]; + const int4 phi = *(const int4*)&sPw[lo * KT + 8]; + __builtin_memcpy(&pa, &plo, 16); + __builtin_memcpy(((char*)&pa) + 16, &phi, 16); + #pragma unroll + for (int f = 0; f < FRAGS; ++f) { + h16v vb; + const int4 vlo = *(const int4*)&sV[(16 * f + lo) * KT]; + const int4 vhi = *(const int4*)&sV[(16 * f + lo) * KT + 8]; + __builtin_memcpy(&vb, &vlo, 16); + __builtin_memcpy(((char*)&vb) + 16, &vhi, 16); + acc[f] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(pa, vb, acc[f]); + } + } + __syncthreads(); + } + + #pragma unroll + for (int i = 0; i < 8; ++i) { + const int qp = wq0 + 2 * i + hi; + if (qp >= qlen) continue; + const float inv = 1.0f / (l_st[i] + 1e-10f); + __half* orow = out + (int64_t)(q0 + qp) * so_tok + (int64_t)head * so_head; + #pragma unroll + for (int f = 0; f < FRAGS; ++f) + orow[16 * f + lo] = __float2half(acc[f][i] * inv); + } + #endif // OCTAVE_PREFILL_STUB +} + +// Partials of phase 1, kept alive for the process: a graph that captured a +// pointer must keep finding it valid. Grows geometrically, so it settles. +float* workspace(int64_t need, const at::TensorOptions& opts) { + static std::mutex mtx; + static std::unordered_map> kept; + std::lock_guard lock(mtx); + std::vector& bufs = kept[opts.device().index()]; + if (bufs.empty() || bufs.back().numel() < need) { + int64_t n = 8 << 20; + while (n < need) n *= 2; + bufs.push_back(at::empty({n}, opts)); + } + return (float*)bufs.back().data_ptr(); +} + +} // namespace octave_pf + +int octave_slot_bytes(int64_t fmt); +void octave_prefill_cdna(torch::Tensor out, torch::Tensor q, torch::Tensor k, + torch::Tensor v, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor cu_seqlens_q, + torch::Tensor seq_lens, int64_t max_query_len, + double sm_scale, int64_t fmt, float* ws, + int num_splits); + +void octave_prefill(torch::Tensor out, torch::Tensor q, torch::Tensor k, + torch::Tensor v, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor cu_seqlens_q, + torch::Tensor seq_lens, int64_t max_query_len, + double sm_scale, int64_t fmt) { + using namespace octave_pf; + TORCH_CHECK(q.dtype() == at::kHalf && k.dtype() == at::kHalf && + v.dtype() == at::kHalf && out.dtype() == at::kHalf, + "octave_prefill: fp16 only"); + TORCH_CHECK(q.size(2) == D && q.stride(2) == 1 && k.stride(2) == 1 && + v.stride(2) == 1 && out.stride(2) == 1); + TORCH_CHECK(cache.dtype() == at::kByte && cache.dim() == 4 && + cache.stride(3) == 1); + TORCH_CHECK(cache.size(3) == octave_slot_bytes(fmt), + "octave_prefill: slot size mismatch"); + const int align = fmt & octave::kCompactFlag ? 4 : 8; + TORCH_CHECK(cache.stride(2) % align == 0 && cache.stride(1) % align == 0 && + cache.stride(0) % align == 0, + "octave_prefill: misaligned slots"); + // 0: gfx1100 (WMMA), 1: gfx942 (MFMA, octave_prefill_cdna.cu). + static const int arch = [] { + const std::string name = + at::cuda::getCurrentDeviceProperties()->gcnArchName; + if (name.rfind("gfx1100", 0) == 0) return 0; + if (name.rfind("gfx942", 0) == 0) return 1; + return -1; + }(); + TORCH_CHECK(arch >= 0, "octave_prefill: gfx1100 or gfx942 only"); + const int num_seqs = seq_lens.size(0); + if (num_seqs == 0 || q.size(0) == 0) return; + const at::cuda::OptionalCUDAGuard guard(device_of(q)); + auto stream = at::cuda::getCurrentCUDAStream().stream(); + const int hq = q.size(1), hkv = cache.size(1); + TORCH_CHECK(hq % hkv == 0); + const int hpk = hq / hkv; + const int block_size = cache.size(2); + const int max_blocks = block_table.size(1); + TORCH_CHECK(block_table.stride(1) == 1 && block_table.stride(0) == max_blocks, + "octave_prefill: block_table rows must be contiguous"); + const int total_q = q.size(0); + + const int rowtiles = (int)((max_query_len * hpk + 127) / 128); + const int base_blocks = std::max(1, num_seqs * hkv * rowtiles); + // Blocks whose split has no keys exit at once, so ask for plenty. + const int num_splits = + std::min(16, std::max(1, (1024 + base_blocks - 1) / base_blocks)); + float* ws = workspace((int64_t)num_splits * total_q * hq * (D + 2), + q.options().dtype(at::kFloat)); + if (arch == 1) { + octave_prefill_cdna(out, q, k, v, cache, block_table, cu_seqlens_q, + seq_lens, max_query_len, sm_scale, fmt, ws, num_splits); + return; + } + + dim3 grid1(num_seqs, hkv, rowtiles * num_splits); + #define PREFIX(B, KC) \ + prefix_attn<<>>( \ + (const __half*)q.data_ptr(), cache.data_ptr(), \ + block_table.data_ptr(), cu_seqlens_q.data_ptr(), \ + seq_lens.data_ptr(), ws, (float)sm_scale, hq, hkv, block_size, \ + max_blocks, num_splits, total_q, q.stride(0), q.stride(1), \ + cache.stride(0), cache.stride(1), cache.stride(2)) + if (fmt == 4) + PREFIX(4, false); + else if (fmt == 3) + PREFIX(3, false); + else + PREFIX(3, true); + #undef PREFIX + dim3 grid2(num_seqs, hq, (max_query_len + CM - 1) / CM); + chunk_attn<<>>( + (__half*)out.data_ptr(), ws, (const __half*)q.data_ptr(), + (const __half*)k.data_ptr(), (const __half*)v.data_ptr(), + cu_seqlens_q.data_ptr(), seq_lens.data_ptr(), hq, hkv, + max_blocks * block_size, (float)sm_scale, num_splits, total_q, + q.stride(0), q.stride(1), k.stride(0), k.stride(1), v.stride(0), + v.stride(1), out.stride(0), out.stride(1)); +} + +#else +void octave_prefill(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, + int64_t, double, int64_t) { + TORCH_CHECK(false, "octave requires ROCm"); +} +#endif diff --git a/csrc/attention/octave_torch_bindings.cpp b/csrc/attention/octave_torch_bindings.cpp new file mode 100644 index 000000000000..9446606df856 --- /dev/null +++ b/csrc/attention/octave_torch_bindings.cpp @@ -0,0 +1,82 @@ +// Op registration for the Octave KV cache kernels (legacy _C extension, full +// torch API). Built only for gfx1100 and gfx942; see CMakeLists.txt. +#include + +#include +#include + +#include "core/registration.h" + +// fmt: V bits (3 or 4), plus 16 for compact K. +void octave_cache_store(torch::Tensor key, torch::Tensor value, + torch::Tensor cache, torch::Tensor slot_mapping, + torch::Tensor k_signs, torch::Tensor v_signs, + int64_t fmt); + +void octave_decode(torch::Tensor out, torch::Tensor query, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor q_to_req, + torch::Tensor q_to_klen, torch::Tensor mid_o, + torch::Tensor k_signs, torch::Tensor v_signs, + double sm_scale, int64_t num_kv_splits, int64_t fmt, + int64_t query_group, bool use_wmma); + +void octave_decode_sparse(torch::Tensor out, torch::Tensor query, + torch::Tensor cache, torch::Tensor block_table, + torch::Tensor q_to_req, torch::Tensor indices, + torch::Tensor mid_o, torch::Tensor k_signs, + torch::Tensor v_signs, double sm_scale, + int64_t num_kv_splits, int64_t fmt, + const std::optional& positions, + const std::optional& wtab, + const std::optional& wtags, + const std::optional& stab, + const std::optional& stags); + +void octave_window_store(torch::Tensor key, torch::Tensor value, + torch::Tensor slot_mapping, torch::Tensor positions, + torch::Tensor wtab, torch::Tensor wtags, + torch::Tensor wlocks, torch::Tensor stab, + torch::Tensor stags, torch::Tensor slocks); + +void octave_rotate(torch::Tensor x, torch::Tensor signs, bool k_layout, + bool inverse); + +void octave_prefill(torch::Tensor out, torch::Tensor q, torch::Tensor k, + torch::Tensor v, torch::Tensor cache, + torch::Tensor block_table, torch::Tensor cu_seqlens_q, + torch::Tensor seq_lens, int64_t max_query_len, + double sm_scale, int64_t fmt); + +TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def( + "octave_cache_store(Tensor key, Tensor value, Tensor! cache, " + "Tensor slot_mapping, Tensor k_signs, Tensor v_signs, int fmt) -> ()"); + ops.impl("octave_cache_store", torch::kCUDA, &octave_cache_store); + ops.def( + "octave_decode(Tensor! out, Tensor query, Tensor cache, " + "Tensor block_table, Tensor q_to_req, Tensor q_to_klen, " + "Tensor! mid_o, Tensor k_signs, Tensor v_signs, float sm_scale, " + "int num_kv_splits, int fmt, int query_group, bool use_wmma) -> ()"); + ops.impl("octave_decode", torch::kCUDA, &octave_decode); + ops.def( + "octave_decode_sparse(Tensor! out, Tensor query, Tensor cache, " + "Tensor block_table, Tensor q_to_req, Tensor indices, Tensor! mid_o, " + "Tensor k_signs, Tensor v_signs, float sm_scale, int num_kv_splits, " + "int fmt, Tensor? positions=None, Tensor? wtab=None, " + "Tensor? wtags=None, Tensor? stab=None, Tensor? stags=None) -> ()"); + ops.impl("octave_decode_sparse", torch::kCUDA, &octave_decode_sparse); + ops.def( + "octave_window_store(Tensor key, Tensor value, Tensor slot_mapping, " + "Tensor positions, Tensor! wtab, Tensor! wtags, Tensor! wlocks, " + "Tensor! stab, Tensor! stags, Tensor! slocks) -> ()"); + ops.impl("octave_window_store", torch::kCUDA, &octave_window_store); + ops.def( + "octave_rotate(Tensor! x, Tensor signs, bool k_layout, bool inverse) " + "-> ()"); + ops.impl("octave_rotate", torch::kCUDA, &octave_rotate); + ops.def( + "octave_prefill(Tensor! out, Tensor q, Tensor k, Tensor v, Tensor cache, " + "Tensor block_table, Tensor cu_seqlens_q, Tensor seq_lens, " + "int max_query_len, float sm_scale, int fmt) -> ()"); + ops.impl("octave_prefill", torch::kCUDA, &octave_prefill); +} diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index 351bcf24f9ff..bd4a080c25ff 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -420,7 +420,75 @@ __global__ void reshape_and_cache_flash_kernel( } } +// Asymmetric K/V head widths. K and V are views of one packed buffer, so they +// share a head stride and the head-contiguous path below cannot apply; this +// kernel therefore walks heads explicitly and takes V's width as its one extra +// argument. It is kept separate so the symmetric kernel's codegen is unchanged. template +__global__ void reshape_and_cache_flash_kernel_asym( + const scalar_t* __restrict__ key, // [num_tokens, num_heads, head_size] + const scalar_t* __restrict__ value, // [num_tokens, num_heads, v_head_size] + // NHD or HND: [num_blocks, block_size, num_heads, head_size], addressed + // through the strides passed by the launcher + cache_t* __restrict__ key_cache, + cache_t* __restrict__ value_cache, // same, with v_head_size + const int64_t* __restrict__ slot_mapping, // [num_tokens] + const int64_t block_stride, const int64_t page_stride, + const int64_t head_stride, const int64_t key_stride, + const int64_t value_stride, const int num_heads, const int head_size, + const int v_head_size, const int block_size, const float* k_scale, + const float* v_scale, const int kv_scale_stride) { + const int64_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + // NOTE: slot_idx can be -1 if the token is padded + if (slot_idx < 0) { + return; + } + const int64_t block_idx = slot_idx / block_size; + const int64_t block_offset = slot_idx % block_size; + + const scalar_t* __restrict__ key_src = key + token_idx * key_stride; + const scalar_t* __restrict__ value_src = value + token_idx * value_stride; + + cache_t* __restrict__ key_dst = + key_cache + block_idx * block_stride + block_offset * page_stride; + cache_t* __restrict__ value_dst = + value_cache + block_idx * block_stride + block_offset * page_stride; + + constexpr int VEC_SIZE = (sizeof(scalar_t) == 2) ? 8 : 4; + + const int lane = threadIdx.x & 31; // 0..31 within warp + const int warp_id = threadIdx.x >> 5; // warp index within block + const int warps_per_block = blockDim.x >> 5; + + for (int head = warp_id; head < num_heads; head += warps_per_block) { + const scalar_t* __restrict__ k_src_h = key_src + head * head_size; + const scalar_t* __restrict__ v_src_h = value_src + head * v_head_size; + + cache_t* __restrict__ k_dst_h = + key_dst + static_cast(head) * head_stride; + cache_t* __restrict__ v_dst_h = + value_dst + static_cast(head) * head_stride; + + float k_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) + ? 0.f + : k_scale[head * kv_scale_stride]; + float v_scale_val = (kv_dt == Fp8KVCacheDataType::kAuto) + ? 0.f + : v_scale[head * kv_scale_stride]; + + CopyWithScaleOp k_op{k_scale_val}; + CopyWithScaleOp v_op{v_scale_val}; + + vectorize_with_alignment(k_src_h, k_dst_h, head_size, lane, 32, + k_op); + vectorize_with_alignment(v_src_h, v_dst_h, v_head_size, lane, 32, + v_op); + } +} + +template __global__ void concat_and_cache_mla_kernel( const scalar_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] const scalar_t* __restrict__ k_pe, // [num_tokens, pe_dim] @@ -453,6 +521,9 @@ __global__ void concat_and_cache_mla_kernel( block_idx * block_stride + block_offset * entry_stride + i + offset; if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { dst[dst_idx] = src[src_idx]; + } else if constexpr (skip_scale) { + dst[dst_idx] = + fp8::scaled_convert(src[src_idx], 1.0f); } else { dst[dst_idx] = fp8::scaled_convert(src[src_idx], *scale); @@ -536,6 +607,7 @@ __global__ void concat_and_cache_ds_mla_kernel( const int kv_lora_rank, // const int pe_dim, // const int block_size, // + const bool has_rope_tail, // const float* scale // ) { const int64_t token_idx = blockIdx.x; @@ -552,15 +624,15 @@ __global__ void concat_and_cache_ds_mla_kernel( // For the NoPE part, each tile of 128 elements is handled by half of one warp // (16 threads). There are 4 total tiles, so 2 warps (64 threads). // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. - // The RoPE part (last 64 elements) is handled by another 1 warp (32 threads). + // The optional RoPE tail is handled by another 1 warp (32 threads). // So in total, we use 3 warps (96 threads) per block. - // Cast kv_cache to 16_bit for RoPE values - scalar_t* kv_cache_16bit = - reinterpret_cast(&kv_cache[dst_idx_start]); - - // Zero the reserved RoPE tail for NoPE rows. if (threadIdx.x >= 64) { + if (!has_rope_tail) { + return; + } + scalar_t* kv_cache_16bit = + reinterpret_cast(&kv_cache[dst_idx_start]); // Each thread handles two elements of RoPE const int8_t pe_idx_start = (threadIdx.x - 64) * 2; // RoPE values start after the packed 8-bit NoPE values and the @@ -822,6 +894,20 @@ void reshape_and_cache( // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. +#define CALL_RESHAPE_AND_CACHE_FLASH_ASYM(KV_T, CACHE_T, KV_DTYPE) \ + vllm::reshape_and_cache_flash_kernel_asym \ + <<>>( \ + reinterpret_cast(key.data_ptr()), \ + reinterpret_cast(value.data_ptr()), \ + reinterpret_cast(key_cache.data_ptr()), \ + reinterpret_cast(value_cache.data_ptr()), \ + slot_mapping.const_data_ptr(), block_stride, page_stride, \ + head_stride, key_stride, value_stride, num_heads, head_size, \ + v_head_size, block_size, \ + reinterpret_cast(k_scale.data_ptr()), \ + reinterpret_cast(v_scale.data_ptr()), \ + kv_scale_stride); + #define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE) \ vllm::reshape_and_cache_flash_kernel \ <<>>( \ @@ -859,6 +945,7 @@ void reshape_and_cache_flash( int num_tokens = slot_mapping.size(0); int num_heads = key.size(1); int head_size = key.size(2); + int v_head_size = value.size(2); const torch::stable::accelerator::DeviceGuard device_guard( key.get_device_index()); @@ -892,8 +979,8 @@ void reshape_and_cache_flash( int64_t block_stride = key_cache.stride(0); int64_t page_stride = key_cache.stride(1); int64_t head_stride = key_cache.stride(2); - STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); + STD_TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0)); STD_TORCH_CHECK(k_scale.sizes().equals(v_scale.sizes()), "k_scale and v_scale must have the same shape"); STD_TORCH_CHECK(k_scale.numel() == 1 || k_scale.numel() == num_heads, @@ -901,17 +988,38 @@ void reshape_and_cache_flash( int kv_scale_stride = (k_scale.numel() > 1) ? 1 : 0; dim3 grid(num_tokens); - dim3 block(std::min(num_heads * head_size, 512)); - - DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, - CALL_RESHAPE_AND_CACHE_FLASH); + // Only models with v_head_size != head_size take the kernel below; everything + // else keeps using the kernel above, so its codegen is untouched. + if (v_head_size == head_size) { + dim3 block(std::min(num_heads * head_size, 512)); + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, + CALL_RESHAPE_AND_CACHE_FLASH); + } else { + // V is a view of K's buffer, so the two must agree on geometry; these + // checks stay out of the symmetric path so its host-side work is unchanged. + STD_TORCH_CHECK(value_cache.stride(1) == page_stride, + "key_cache and value_cache must share the page stride"); + STD_TORCH_CHECK(key_cache.size(0) == value_cache.size(0) && + key_cache.size(1) == value_cache.size(1), + "key_cache and value_cache must share block geometry"); + STD_TORCH_CHECK(key.size(1) == value.size(1), + "key and value must have the same number of heads"); + STD_TORCH_CHECK(value_cache.stride(2) == head_stride, + "key_cache and value_cache must share a head stride"); + // Only the strided walk applies here, and V's width is the single extra + // argument. One warp per head, and at least one warp so the head loop + // always advances. + dim3 block(32 * std::min(num_heads, 16)); + DISPATCH_BY_KV_CACHE_DTYPE(key.scalar_type(), kv_cache_dtype, + CALL_RESHAPE_AND_CACHE_FLASH_ASYM); + } } // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. // KV_DTYPE is the real data type of kv-cache. -#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ - vllm::concat_and_cache_mla_kernel \ +#define LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, SKIP_SCALE) \ + vllm::concat_and_cache_mla_kernel \ <<>>( \ reinterpret_cast(kv_c.data_ptr()), \ reinterpret_cast(k_pe.data_ptr()), \ @@ -920,6 +1028,13 @@ void reshape_and_cache_flash( kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ reinterpret_cast(scale.data_ptr())); +#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE) \ + if (skip_scale) { \ + LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, true) \ + } else { \ + LAUNCH_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE, false) \ + } + // KV_T is the data type of key and value tensors. // CACHE_T is the stored data type of kv-cache. #define CALL_CONCAT_AND_CACHE_DS_MLA(KV_T, CACHE_T, KV_DTYPE) \ @@ -930,7 +1045,7 @@ void reshape_and_cache_flash( reinterpret_cast(kv_cache.data_ptr()), \ slot_mapping.const_data_ptr(), block_stride, entry_stride, \ kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size, \ - reinterpret_cast(scale.data_ptr())); + has_rope_tail, reinterpret_cast(scale.data_ptr())); void concat_and_cache_mla( torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] @@ -938,7 +1053,8 @@ void concat_and_cache_mla( torch::stable::Tensor& kv_cache, // [num_blocks, block_size, (kv_lora_rank // + pe_dim)] torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] - const std::string& kv_cache_dtype, torch::stable::Tensor& scale) { + const std::string& kv_cache_dtype, torch::stable::Tensor& scale, + bool skip_scale) { // NOTE(woosuk): In vLLM V1, key.size(0) can be different from // slot_mapping.size(0) because of padding for CUDA graphs. // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because @@ -953,6 +1069,8 @@ void concat_and_cache_mla( int kv_lora_rank = kv_c.size(1); int pe_dim = k_pe.size(1); int block_size = kv_cache.size(1); + const int64_t row_bytes = kv_cache.size(2) * kv_cache.element_size(); + const bool has_rope_tail = row_bytes == 656; const bool is_nvfp4_ds_mla = kv_cache_dtype == "nvfp4_ds_mla"; if (kv_cache_dtype == "fp8_ds_mla") { @@ -960,8 +1078,9 @@ void concat_and_cache_mla( "kv_lora_rank must be 512 for fp8_ds_mla"); STD_TORCH_CHECK(pe_dim == 64 || pe_dim == 0, "pe_dim must be 64 or 0 for fp8_ds_mla"); - STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(), - "kv_cache.size(2) must be 656 bytes for fp8_ds_mla"); + STD_TORCH_CHECK(has_rope_tail || (pe_dim == 0 && row_bytes == 528), + "fp8_ds_mla requires 656-byte rows, or 528-byte rows with " + "pe_dim == 0"); STD_TORCH_CHECK(kv_c.element_size() == 2, "kv_c.element_size() must be 2 for fp8_ds_mla"); STD_TORCH_CHECK(k_pe.element_size() == 2, @@ -999,8 +1118,8 @@ void concat_and_cache_mla( // For the NoPE part, each tile of 128 elements is handled by half of one // warp (16 threads). There are 4 total tiles, so 2 warps (64 threads). // Lanes 0 and 16 of each warp write the scale values for that warp's tiles. - // The RoPE part (last 64 elements) is handled by another 1 warp (32 - // threads). So in total, we use 3 warps (96 threads) per block. + // The optional RoPE tail is handled by another 1 warp (32 threads). + // So in total, we use 3 warps (96 threads) per block. dim3 block(96); DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, CALL_CONCAT_AND_CACHE_DS_MLA); @@ -1030,6 +1149,9 @@ void concat_and_cache_mla( } } +#undef CALL_CONCAT_AND_CACHE_MLA +#undef LAUNCH_CONCAT_AND_CACHE_MLA + void concat_and_cache_mla_grouped( torch::stable::Tensor& kv_c, // [num_layers, num_tokens, kv_lora_rank] torch::stable::Tensor& k_pe, // [num_layers, num_tokens, pe_dim] @@ -1822,6 +1944,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( STD_TORCH_CHECK(dst.scalar_type() == torch::headeronly::ScalarType::BFloat16, "dst must be bfloat16"); STD_TORCH_CHECK(head_dim == 576, "head_dim must be 576 for MLA"); + STD_TORCH_CHECK(src_cache.size(2) * src_cache.element_size() >= 656, + "src_cache rows must contain at least 656 bytes for MLA"); int64_t block_table_stride = block_table.stride(0); int64_t cache_block_stride = src_cache.stride(0); @@ -1864,7 +1988,7 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size) { + int64_t batch_size, std::optional seq_starts) { // The kernel lives in nvfp4_ds_mla_cache_kernels.cu, which is only built when // the target list contains an SM100 arch. #if !defined(USE_ROCM) && defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100 @@ -1887,6 +2011,13 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( STD_TORCH_CHECK( workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int, "workspace_starts must be int32"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); + } STD_TORCH_CHECK(src_cache.device() == dst.device(), "src_cache and dst must be on the same device"); STD_TORCH_CHECK(src_cache.device() == block_table.device(), @@ -1914,6 +2045,8 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( src_cache.const_data_ptr(), dst.data_ptr(), block_table.const_data_ptr(), workspace_starts.const_data_ptr(), + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr, static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, dst_entry_stride, stream); diff --git a/csrc/cpu/block_table.cpp b/csrc/cpu/block_table.cpp new file mode 100644 index 000000000000..ac8f21d44d6e --- /dev/null +++ b/csrc/cpu/block_table.cpp @@ -0,0 +1,73 @@ +#include "cpu_types.hpp" + +#include +#include + +namespace cpu_utils { + +void gather_block_table_cpu(const torch::Tensor& block_table, + torch::Tensor& output, + const torch::Tensor& idx_mapping, + const torch::Tensor& num_blocks, + const int64_t num_reqs_padded) { + const int32_t* block_table_ptr = block_table.data_ptr(); + int32_t* output_ptr = output.data_ptr(); + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int32_t* num_blocks_ptr = num_blocks.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + const int64_t block_table_stride = block_table.stride(0); + const int64_t output_stride = output.stride(0); + for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) { + const int64_t req_idx = idx_mapping_ptr[batch_idx]; + const int64_t request_num_blocks = num_blocks_ptr[req_idx]; + std::memcpy(output_ptr + batch_idx * output_stride, + block_table_ptr + req_idx * block_table_stride, + request_num_blocks * sizeof(int32_t)); + } + + for (int64_t batch_idx = num_reqs; batch_idx < num_reqs_padded; ++batch_idx) { + std::memset(output_ptr + batch_idx * output_stride, 0, + output.size(1) * sizeof(int32_t)); + } +} + +void compute_slot_mapping_cpu( + const torch::Tensor& block_table, torch::Tensor& slot_mapping, + const torch::Tensor& idx_mapping, const torch::Tensor& query_start_loc, + const torch::Tensor& positions, const int64_t block_size, + const bool mapping_enabled, const int64_t pad_slot_id) { + const int32_t* block_table_ptr = block_table.data_ptr(); + int64_t* slot_mapping_ptr = slot_mapping.data_ptr(); + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int32_t* query_start_loc_ptr = query_start_loc.data_ptr(); + const int64_t* positions_ptr = positions.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + const int64_t num_tokens = query_start_loc_ptr[num_reqs]; + const int64_t block_table_stride = block_table.stride(0); + std::fill(slot_mapping_ptr + num_tokens, + slot_mapping_ptr + slot_mapping.numel(), pad_slot_id); + + for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) { + const int64_t req_state_idx = idx_mapping_ptr[batch_idx]; + const int64_t start = query_start_loc_ptr[batch_idx]; + const int64_t end = query_start_loc_ptr[batch_idx + 1]; + if (!mapping_enabled || req_state_idx < 0) { + std::fill(slot_mapping_ptr + start, slot_mapping_ptr + end, pad_slot_id); + continue; + } + + for (int64_t token_idx = start; token_idx < end; ++token_idx) { + const int64_t position = positions_ptr[token_idx]; + const int64_t block_idx = position / block_size; + const int64_t block_offset = position % block_size; + const int32_t block_number = + block_table_ptr[req_state_idx * block_table_stride + block_idx]; + slot_mapping_ptr[token_idx] = + static_cast(block_number) * block_size + block_offset; + } + } +} + +} // namespace cpu_utils diff --git a/csrc/cpu/cpu_types_scalar.hpp b/csrc/cpu/cpu_types_scalar.hpp index d94cdbc076e8..76b3aec5d45e 100644 --- a/csrc/cpu/cpu_types_scalar.hpp +++ b/csrc/cpu/cpu_types_scalar.hpp @@ -502,7 +502,7 @@ inline FP16Vec16::FP16Vec16(const FP32Vec16& v) { [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); }); } -inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) { +inline FP16Vec8::FP16Vec8(const FP32Vec8& v) { unroll_loop( [&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); }); } diff --git a/csrc/cpu/input_batch.cpp b/csrc/cpu/input_batch.cpp new file mode 100644 index 000000000000..8f04e05b4d2b --- /dev/null +++ b/csrc/cpu/input_batch.cpp @@ -0,0 +1,295 @@ +#include "cpu_types.hpp" + +#include +#include +#include + +namespace cpu_utils { + +void prepare_prefill_inputs_cpu(torch::Tensor& input_ids, + torch::Tensor& next_prefill_tokens, + const torch::Tensor& idx_mapping, + const torch::Tensor& query_start_loc, + const torch::Tensor& all_token_ids, + const torch::Tensor& prefill_len, + const torch::Tensor& num_computed_tokens) { + int32_t* input_ids_ptr = input_ids.data_ptr(); + int32_t* next_prefill_tokens_ptr = next_prefill_tokens.data_ptr(); + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int32_t* query_start_loc_ptr = query_start_loc.data_ptr(); + const int32_t* all_token_ids_ptr = all_token_ids.data_ptr(); + const int32_t* prefill_len_ptr = prefill_len.data_ptr(); + const int32_t* num_computed_tokens_ptr = + num_computed_tokens.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + const int64_t num_lookahead = next_prefill_tokens.size(0); + const int64_t next_prefill_tokens_stride = next_prefill_tokens.stride(0); + const int64_t all_token_ids_stride = all_token_ids.stride(0); + + for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) { + const int64_t req_state_idx = idx_mapping_ptr[batch_idx]; + const int32_t request_prefill_len = prefill_len_ptr[req_state_idx]; + const int32_t num_computed = num_computed_tokens_ptr[req_state_idx]; + if (num_computed >= request_prefill_len) { + continue; + } + + const int32_t query_start = query_start_loc_ptr[batch_idx]; + const int32_t query_end = query_start_loc_ptr[batch_idx + 1]; + const int32_t query_len = query_end - query_start; + const int32_t* request_tokens = + all_token_ids_ptr + req_state_idx * all_token_ids_stride; + std::memcpy(input_ids_ptr + query_start, request_tokens + num_computed, + query_len * sizeof(int32_t)); + + const int64_t lookahead_start = + static_cast(num_computed) + query_len; + for (int64_t lookahead = 0; lookahead < num_lookahead; ++lookahead) { + const int64_t pos = lookahead_start + lookahead; + next_prefill_tokens_ptr[lookahead * next_prefill_tokens_stride + + req_state_idx] = + pos < request_prefill_len ? request_tokens[pos] : 0; + } + } +} + +torch::Tensor get_prompt_logprobs_token_ids_cpu( + const int64_t num_tokens, const torch::Tensor& query_start_loc, + const torch::Tensor& idx_mapping, const torch::Tensor& num_computed_tokens, + const torch::Tensor& all_token_ids) { + torch::Tensor token_ids = + torch::empty({num_tokens}, idx_mapping.options().dtype(torch::kInt64)); + + int64_t* token_ids_ptr = token_ids.data_ptr(); + const int32_t* query_start_loc_ptr = query_start_loc.data_ptr(); + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int32_t* num_computed_tokens_ptr = + num_computed_tokens.data_ptr(); + const int32_t* all_token_ids_ptr = all_token_ids.data_ptr(); + const int64_t all_token_ids_stride = all_token_ids.stride(0); + + for (int64_t batch_idx = 0; batch_idx < idx_mapping.size(0); ++batch_idx) { + const int64_t req_state_idx = idx_mapping_ptr[batch_idx]; + const int32_t query_start = query_start_loc_ptr[batch_idx]; + const int32_t query_end = query_start_loc_ptr[batch_idx + 1]; + const int32_t target_start = num_computed_tokens_ptr[req_state_idx] + 1; + const int32_t* request_tokens = + all_token_ids_ptr + req_state_idx * all_token_ids_stride; + + std::transform(request_tokens + target_start, + request_tokens + target_start + query_end - query_start, + token_ids_ptr + query_start, [](int32_t token_id) { + return static_cast(token_id); + }); + } + return token_ids; +} + +void prepare_pos_seq_lens_cpu(const torch::Tensor& idx_mapping, + const torch::Tensor& query_start_loc, + const torch::Tensor& num_computed_tokens, + torch::Tensor& pos, torch::Tensor& seq_lens) { + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int32_t* query_start_loc_ptr = query_start_loc.data_ptr(); + const int32_t* num_computed_tokens_ptr = + num_computed_tokens.data_ptr(); + int64_t* pos_ptr = pos.data_ptr(); + int32_t* seq_lens_ptr = seq_lens.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + std::fill(seq_lens_ptr + num_reqs, seq_lens_ptr + seq_lens.numel(), 0); + + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + const int64_t req_state_idx = idx_mapping_ptr[req_idx]; + const int32_t num_computed = num_computed_tokens_ptr[req_state_idx]; + const int64_t start = query_start_loc_ptr[req_idx]; + const int64_t end = query_start_loc_ptr[req_idx + 1]; + + seq_lens_ptr[req_idx] = num_computed + static_cast(end - start); + std::iota(pos_ptr + start, pos_ptr + end, + static_cast(num_computed)); + } +} + +void prepare_mrope_positions_cpu(torch::Tensor& positions, + const torch::Tensor& prefill_positions, + const torch::Tensor& prefill_delta, + const torch::Tensor& idx_mapping, + const torch::Tensor& query_start_loc, + const torch::Tensor& prefill_lens, + const torch::Tensor& num_computed_tokens) { + int64_t* positions_ptr = positions.data_ptr(); + const int32_t* prefill_positions_ptr = prefill_positions.data_ptr(); + const int32_t* prefill_delta_ptr = prefill_delta.data_ptr(); + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int32_t* query_start_loc_ptr = query_start_loc.data_ptr(); + const int32_t* prefill_lens_ptr = prefill_lens.data_ptr(); + const int32_t* num_computed_tokens_ptr = + num_computed_tokens.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + const int64_t num_dims = positions.size(0); + const int64_t positions_stride0 = positions.stride(0); + const int64_t positions_stride1 = positions.stride(1); + const int64_t prefill_positions_stride0 = prefill_positions.stride(0); + const int64_t prefill_positions_stride1 = prefill_positions.stride(1); + + for (int64_t batch_idx = 0; batch_idx < num_reqs; ++batch_idx) { + const int64_t req_state_idx = idx_mapping_ptr[batch_idx]; + const int32_t query_start = query_start_loc_ptr[batch_idx]; + const int32_t query_end = query_start_loc_ptr[batch_idx + 1]; + const int32_t query_len = query_end - query_start; + const int32_t num_computed = num_computed_tokens_ptr[req_state_idx]; + const int32_t prompt_len = std::min( + query_len, std::max(prefill_lens_ptr[req_state_idx] - num_computed, 0)); + const int32_t delta = prefill_delta_ptr[req_state_idx]; + + for (int64_t dim = 0; dim < num_dims; ++dim) { + int64_t* output = positions_ptr + dim * positions_stride0 + + query_start * positions_stride1; + const int32_t* prompt = + prefill_positions_ptr + + (req_state_idx * num_dims + dim) * prefill_positions_stride0; + + for (int32_t token_idx = 0; token_idx < prompt_len; ++token_idx) { + output[token_idx * positions_stride1] = + prompt[(num_computed + token_idx) * prefill_positions_stride1]; + } + for (int32_t token_idx = prompt_len; token_idx < query_len; ++token_idx) { + output[token_idx * positions_stride1] = + num_computed + token_idx + delta; + } + } + } +} + +torch::Tensor combine_sampled_and_draft_tokens_cpu( + torch::Tensor& input_ids, const torch::Tensor& idx_mapping, + const torch::Tensor& last_sampled_tokens, + const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens, + const torch::Tensor& prefill_len, const torch::Tensor& draft_tokens, + const torch::Tensor& cu_num_logits, const int64_t num_logits, + const int64_t num_new_sampled_tokens) { + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + const int64_t* last_sampled_tokens_ptr = + last_sampled_tokens.data_ptr(); + const int32_t* query_start_loc_ptr = query_start_loc.data_ptr(); + const int32_t* seq_lens_ptr = seq_lens.data_ptr(); + const int32_t* prefill_len_ptr = prefill_len.data_ptr(); + const int64_t* draft_tokens_ptr = draft_tokens.data_ptr(); + const int32_t* cu_num_logits_ptr = cu_num_logits.data_ptr(); + int32_t* input_ids_ptr = input_ids.data_ptr(); + + torch::Tensor logits_indices = + torch::empty({num_logits}, input_ids.options().dtype(torch::kInt64)); + int64_t* logits_indices_ptr = logits_indices.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + const int64_t draft_tokens_stride = draft_tokens.stride(0); + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + const int64_t req_state_idx = idx_mapping_ptr[req_idx]; + const int64_t logits_offset = cu_num_logits_ptr[req_idx]; + const int64_t logits_offset_end = cu_num_logits_ptr[req_idx + 1]; + const int64_t request_num_logits = logits_offset_end - logits_offset; + const int64_t num_draft_tokens = + request_num_logits - num_new_sampled_tokens; + const int64_t query_end = query_start_loc_ptr[req_idx + 1]; + const int64_t logits_start = query_end - request_num_logits; + + std::iota(logits_indices_ptr + logits_offset, + logits_indices_ptr + logits_offset_end, logits_start); + + const int32_t seq_len = seq_lens_ptr[req_idx]; + const int32_t request_prefill_len = prefill_len_ptr[req_state_idx]; + if (seq_len <= request_prefill_len) { + continue; + } + + const int64_t first_logit_seq_pos = seq_len - request_num_logits; + if (num_new_sampled_tokens > 0 && + first_logit_seq_pos >= request_prefill_len) { + input_ids_ptr[logits_start] = + static_cast(last_sampled_tokens_ptr[req_state_idx]); + } + + const int64_t draft_start = query_end - num_draft_tokens; + for (int64_t i = 0; i < num_draft_tokens; ++i) { + input_ids_ptr[draft_start + i] = static_cast( + draft_tokens_ptr[req_state_idx * draft_tokens_stride + i]); + } + } + + return logits_indices; +} + +void post_update_cpu(const torch::Tensor& idx_mapping, + torch::Tensor& num_computed_tokens, + torch::Tensor& last_sampled_tokens, + const c10::optional& output_bin_counts, + const torch::Tensor& sampled_tokens, + const torch::Tensor& num_sampled, + const torch::Tensor& num_rejected, + const c10::optional& query_start_loc, + torch::Tensor& all_token_ids, torch::Tensor& total_len) { + const int64_t* idx_mapping_ptr = idx_mapping.data_ptr(); + int32_t* num_computed_tokens_ptr = num_computed_tokens.data_ptr(); + int64_t* last_sampled_tokens_ptr = last_sampled_tokens.data_ptr(); + int32_t* output_bin_counts_ptr = + output_bin_counts ? output_bin_counts->data_ptr() : nullptr; + const int64_t* sampled_tokens_ptr = sampled_tokens.data_ptr(); + const int32_t* num_sampled_ptr = num_sampled.data_ptr(); + const int32_t* num_rejected_ptr = num_rejected.data_ptr(); + const int32_t* query_start_loc_ptr = + query_start_loc ? query_start_loc->data_ptr() : nullptr; + int32_t* all_token_ids_ptr = all_token_ids.data_ptr(); + int32_t* total_len_ptr = total_len.data_ptr(); + + const int64_t num_reqs = idx_mapping.size(0); + const int64_t last_sampled_tokens_stride = last_sampled_tokens.stride(0); + const int64_t output_bin_counts_stride = + output_bin_counts ? output_bin_counts->stride(0) : 0; + const int64_t sampled_tokens_stride = sampled_tokens.stride(0); + const int64_t all_token_ids_stride = all_token_ids.stride(0); + + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + const int64_t req_state_idx = idx_mapping_ptr[req_idx]; + if (req_state_idx < 0) { + continue; + } + + const int32_t request_total_len = total_len_ptr[req_state_idx]; + const int32_t request_num_sampled = num_sampled_ptr[req_idx]; + const int64_t* request_sampled_tokens = + sampled_tokens_ptr + req_idx * sampled_tokens_stride; + int32_t* request_all_token_ids = + all_token_ids_ptr + req_state_idx * all_token_ids_stride; + + if (request_num_sampled > 0) { + last_sampled_tokens_ptr[req_state_idx * last_sampled_tokens_stride] = + request_sampled_tokens[request_num_sampled - 1]; + total_len_ptr[req_state_idx] = request_total_len + request_num_sampled; + } + + for (int32_t i = 0; i < request_num_sampled; ++i) { + const int64_t token_id = request_sampled_tokens[i]; + request_all_token_ids[request_total_len + i] = + static_cast(token_id); + if (output_bin_counts_ptr != nullptr) { + ++output_bin_counts_ptr[req_state_idx * output_bin_counts_stride + + token_id]; + } + } + + const int32_t query_len = + query_start_loc_ptr == nullptr + ? 0 + : query_start_loc_ptr[req_idx + 1] - query_start_loc_ptr[req_idx]; + const int32_t computed_delta = query_len - num_rejected_ptr[req_idx]; + if (computed_delta != 0) { + num_computed_tokens_ptr[req_state_idx] += computed_delta; + } + } +} + +} // namespace cpu_utils diff --git a/csrc/cpu/sampling/common.h b/csrc/cpu/sampling/common.h new file mode 100644 index 000000000000..d18da0079c1e --- /dev/null +++ b/csrc/cpu/sampling/common.h @@ -0,0 +1,279 @@ +#pragma once + +#include "../utils.hpp" + +#include +#include +#include +#include +#include +#include + +namespace cpu_sampling { + +// Draft proposals and target verification must draw from disjoint noise +// streams even when they use the same request seed and token position. +inline constexpr int64_t DRAFT_NOISE_SALT = 1 << 30; + +inline uint32_t murmur3_rotl32(uint32_t value, int shift) { + return (value << shift) | (value >> (32 - shift)); +} + +inline uint32_t murmur3_mix(uint32_t hash, uint32_t key) { + key *= 0xcc9e2d51; + key = murmur3_rotl32(key, 15); + key *= 0x1b873593; + hash ^= key; + hash = murmur3_rotl32(hash, 13); + return hash * 5 + 0xe6546b64; +} + +inline uint32_t murmur3_fmix32(uint32_t hash) { + hash ^= hash >> 16; + hash *= 0x85ebca6b; + hash ^= hash >> 13; + hash *= 0xc2b2ae35; + return hash ^ (hash >> 16); +} + +inline uint32_t murmur3_hash32_prefix(int64_t seed, int64_t position, + uint32_t domain) { + // Seed and position are shared by every vocabulary entry in a row. Mix this + // prefix once so the inner loop only needs to mix the token ID and finalize. + const uint64_t seed_bits = static_cast(seed); + uint32_t hash = domain; + hash = murmur3_mix(hash, static_cast(seed_bits)); + hash = murmur3_mix(hash, static_cast(seed_bits >> 32)); + hash = murmur3_mix(hash, static_cast(position)); + return hash; +} + +inline uint32_t murmur3_hash32(int64_t seed, int64_t position, + uint32_t token_id, uint32_t domain = 0) { + uint32_t hash = murmur3_hash32_prefix(seed, position, domain); + hash = murmur3_mix(hash, token_id); + return murmur3_fmix32(hash ^ 16); +} + +inline float murmur3_uniform32(int64_t seed, int64_t position, + uint32_t token_id) { + const uint32_t random = murmur3_hash32(seed, position, token_id); + // Splitting the hash makes both integer-to-float conversions exact: + // u = round_fp32((random + 0.5) / 2^32). + // Converting random directly to FP32 could discard up to eight low bits. + const int32_t high = static_cast(random >> 16); + const int32_t low = static_cast(random & 0xffff); + return static_cast(high) * 1.52587890625e-05f + + (static_cast(low) + 0.5f) * 2.3283064365386963e-10f; +} + +inline double murmur3_uniform64(int64_t seed, int64_t position, + uint32_t token_id) { + const uint64_t low = murmur3_hash32(seed, position, token_id); + const uint64_t high = murmur3_hash32(seed, position, token_id, 0x9e3779b9); + // Combine two hashes and retain the 53 bits representable by an FP64 + // significand: u = (random53 + 0.5) / 2^53. + const uint64_t random53 = ((high << 32) | low) >> 11; + double uniform = (static_cast(random53) + 0.5) * 0x1p-53; + if (uniform == 1.0) uniform -= 0x1p-53; + return uniform; +} + +inline float log1p_neg_stable(float value) { + // This is inherited from the triton impl: + // For small x, evaluate the degree-8 expansion + // log(1-x) = -x * (1 + x/2 + x^2/3 + ... + x^7/8) + // instead of forming 1-x, which loses the low bits of x. + float polynomial = 1.0f / 8.0f; + polynomial = 1.0f / 7.0f + value * polynomial; + polynomial = 1.0f / 6.0f + value * polynomial; + polynomial = 1.0f / 5.0f + value * polynomial; + polynomial = 1.0f / 4.0f + value * polynomial; + polynomial = 1.0f / 3.0f + value * polynomial; + polynomial = 1.0f / 2.0f + value * polynomial; + polynomial = 1.0f + value * polynomial; + const float series = -value * polynomial; + const float direct = std::log(std::max(1.0f - value, 5.960464477539063e-08f)); + return value < 0.25f ? series : direct; +} + +inline float gumbel_noise32(int64_t seed, int64_t position, uint32_t token_id) { + const float uniform = murmur3_uniform32(seed, position, token_id); + return -std::log(-log1p_neg_stable(uniform)); +} + +inline double gumbel_noise64(int64_t seed, int64_t position, + uint32_t token_id) { + const double uniform = murmur3_uniform64(seed, position, token_id); + return -std::log(-std::log(uniform)); +} + +template +void scalar_greedy_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, float& best_score, + int64_t& best_idx) { + best_score = -std::numeric_limits::infinity(); + best_idx = token_begin; + for (int64_t token_id = token_begin; token_id < token_end; ++token_id) { + const float score = static_cast(logits[token_id]); + if (score > best_score) { + best_score = score; + best_idx = token_id; + } + } +} + +template +void scalar_gumbel_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, int64_t seed, int64_t position, + float temperature, score_t& best_score, + int64_t& best_idx) { + best_score = -std::numeric_limits::infinity(); + best_idx = token_begin; + + for (int64_t token_id = token_begin; token_id < token_end; ++token_id) { + float logit = static_cast(logits[token_id]); + if constexpr (APPLY_TEMPERATURE) logit /= temperature; + + score_t noise; + if constexpr (std::is_same_v) { + noise = gumbel_noise64(seed, position, token_id); + } else { + noise = gumbel_noise32(seed, position, token_id); + } + const score_t score = static_cast(logit) + noise; + if (score > best_score) { + best_score = score; + best_idx = token_id; + } + } +} + +template +class SamplingImpl; + +template +class SamplingMainLoop { + public: + using scalar_t = typename sampling_impl_t::scalar_type; + + void operator()(int64_t* output, const scalar_t* logits, + int64_t logits_stride, const int64_t* expanded_idx_mapping, + const float* temperature, const int64_t* seeds, + const int64_t* positions, int64_t num_tokens, + int64_t vocab_size, bool apply_temperature, bool is_drafting, + bool use_fp64) const { + if (num_tokens == 0) return; + + // Resolve runtime options before entering the vocabulary loop so its + // temperature and score-type branches are compile-time specializations. + if (apply_temperature) { + if (use_fp64) { + execute(output, logits, logits_stride, + expanded_idx_mapping, temperature, seeds, + positions, num_tokens, vocab_size, is_drafting); + } else { + execute(output, logits, logits_stride, + expanded_idx_mapping, temperature, seeds, + positions, num_tokens, vocab_size, is_drafting); + } + } else if (use_fp64) { + execute(output, logits, logits_stride, + expanded_idx_mapping, temperature, seeds, + positions, num_tokens, vocab_size, is_drafting); + } else { + execute(output, logits, logits_stride, expanded_idx_mapping, + temperature, seeds, positions, num_tokens, + vocab_size, is_drafting); + } + } + + private: + template + void execute(int64_t* output, const scalar_t* logits, int64_t logits_stride, + const int64_t* expanded_idx_mapping, const float* temperature, + const int64_t* seeds, const int64_t* positions, + int64_t num_tokens, int64_t vocab_size, bool is_drafting) const { + // Parallelize across logit rows first; num_tokens is not necessarily the + // request count. When there are fewer rows than workers, split each row to + // expose roughly one task per worker. Keep ranges large enough to amortize + // queue overhead and align them to an ISA-chosen block multiple (32 on + // NEON) so range boundaries do not create extra scalar tails. + constexpr int64_t MIN_RANGE_SIZE = 2048; + constexpr int64_t RANGE_ALIGNMENT = sampling_impl_t::RangeAlignment; + static_assert(RANGE_ALIGNMENT > 0); + + const int max_workers = cpu_utils::get_max_threads(); + const int64_t target_ranges_per_row = + std::max(1, (max_workers + num_tokens - 1) / num_tokens); + const int64_t unaligned_range_size = std::max( + MIN_RANGE_SIZE, + (vocab_size + target_ranges_per_row - 1) / target_ranges_per_row); + const int64_t range_size = (unaligned_range_size + RANGE_ALIGNMENT - 1) / + RANGE_ALIGNMENT * RANGE_ALIGNMENT; + const int64_t num_ranges = (vocab_size + range_size - 1) / range_size; + const int64_t num_tasks = num_tokens * num_ranges; + + std::vector range_scores(num_tasks); + std::vector range_indices(num_tasks); + alignas(64) cpu_utils::Counter task_counter; + const sampling_impl_t sampling_impl; + +#pragma omp parallel + { + // Phase 1: find one winner per vocabulary range of each row. + for (;;) { + const int64_t task_id = task_counter.acquire_counter(); + if (task_id >= num_tasks) break; + + const int64_t token_idx = task_id / num_ranges; + const int64_t range_idx = task_id % num_ranges; + const int64_t token_begin = range_idx * range_size; + const int64_t token_end = + std::min(token_begin + range_size, vocab_size); + const int64_t req_state_idx = expanded_idx_mapping[token_idx]; + const bool valid_request = req_state_idx >= 0; + const float temp = valid_request ? temperature[req_state_idx] : 0.0f; + const scalar_t* row = logits + token_idx * logits_stride; + + if (temp == 0.0f) { + float best_score; + sampling_impl.greedy_argmax(row, token_begin, token_end, best_score, + range_indices[task_id]); + range_scores[task_id] = best_score; + } else { + int64_t position = positions[token_idx]; + if (is_drafting) position += DRAFT_NOISE_SALT; + sampling_impl.template gumbel_argmax( + row, token_begin, token_end, seeds[req_state_idx], position, temp, + range_scores[task_id], range_indices[task_id]); + } + } + + // Every range winner must be visible before its row is reduced. +#pragma omp barrier + // Phase 1 found the best token in each vocabulary range. Phase 2 assigns + // each row to one worker, which compares its range winners to find the + // best token for the complete vocabulary. Ranges are visited in order, + // and the strict comparison retains the first token when scores tie. +#pragma omp for schedule(static) + for (int64_t token_idx = 0; token_idx < num_tokens; ++token_idx) { + score_t best_score = -std::numeric_limits::infinity(); + int64_t best_idx = 0; + const int64_t first_range = token_idx * num_ranges; + // Vectorizing this isn't worthwhile because each row has few ranges. + for (int64_t range_idx = 0; range_idx < num_ranges; ++range_idx) { + const int64_t task_id = first_range + range_idx; + const score_t score = range_scores[task_id]; + if (score > best_score) { + best_score = score; + best_idx = range_indices[task_id]; + } + } + output[token_idx] = best_idx; + } + } + } +}; + +} // namespace cpu_sampling diff --git a/csrc/cpu/sampling/neon.h b/csrc/cpu/sampling/neon.h new file mode 100644 index 000000000000..b4808d5c3c09 --- /dev/null +++ b/csrc/cpu/sampling/neon.h @@ -0,0 +1,395 @@ +#pragma once + +#include "common.h" + +#include +#include + +#include +#include +#include + +namespace cpu_sampling::neon { + +FORCE_INLINE uint32x4_t murmur3_hash32x4(uint32x4_t hash, + uint32x4_t token_ids) { + // `hash` contains the seed/position prefix. Mix four token IDs and finalize + // four independent hashes, lane-for-lane matching murmur3_hash32(). + token_ids = vmulq_n_u32(token_ids, 0xcc9e2d51); + token_ids = vorrq_u32(vshlq_n_u32(token_ids, 15), vshrq_n_u32(token_ids, 17)); + token_ids = vmulq_n_u32(token_ids, 0x1b873593); + hash = veorq_u32(hash, token_ids); + hash = vorrq_u32(vshlq_n_u32(hash, 13), vshrq_n_u32(hash, 19)); + hash = vaddq_u32(vmulq_n_u32(hash, 5), vdupq_n_u32(0xe6546b64)); + hash = veorq_u32(hash, vdupq_n_u32(16)); + hash = veorq_u32(hash, vshrq_n_u32(hash, 16)); + hash = vmulq_n_u32(hash, 0x85ebca6b); + hash = veorq_u32(hash, vshrq_n_u32(hash, 13)); + hash = vmulq_n_u32(hash, 0xc2b2ae35); + return veorq_u32(hash, vshrq_n_u32(hash, 16)); +} + +FORCE_INLINE float32x4_t murmur3_uniform32x4(uint32x4_t hash, + uint32x4_t token_ids) { + const uint32x4_t random = murmur3_hash32x4(hash, token_ids); + // Evaluate u = round_fp32((random + 0.5) / 2^32) without first rounding the + // entire 32-bit integer to FP32. + const uint32x4_t high = vshrq_n_u32(random, 16); + const uint32x4_t low = vandq_u32(random, vdupq_n_u32(0xffff)); + return vmlaq_n_f32(vmulq_n_f32(vcvtq_f32_u32(high), 1.52587890625e-05f), + vaddq_f32(vcvtq_f32_u32(low), vdupq_n_f32(0.5f)), + 2.3283064365386963e-10f); +} + +// Adapted from data and inline_logf in Arm Optimized Routines: +// math/aarch64/advsimd/logf.c at commit +// 59afb6fed4275e2fc249df49fce43ee5e75af178. +// https://github.com/ARM-software/optimized-routines +struct PositiveNormalLogData { + float32x4_t c2; + float32x4_t c4; + float32x4_t c6; + float32x4_t ln2; + uint32x4_t offset; + uint32x4_t mantissa_mask; + float32x4_t odd_coefficients; +}; + +constexpr PositiveNormalLogData kPositiveNormalLogData = { + .c2 = {-0x1.4f9934p-3f, -0x1.4f9934p-3f, -0x1.4f9934p-3f, -0x1.4f9934p-3f}, + .c4 = {-0x1.00187cp-2f, -0x1.00187cp-2f, -0x1.00187cp-2f, -0x1.00187cp-2f}, + .c6 = {-0x1.ffffc8p-2f, -0x1.ffffc8p-2f, -0x1.ffffc8p-2f, -0x1.ffffc8p-2f}, + .ln2 = {0x1.62e43p-1f, 0x1.62e43p-1f, 0x1.62e43p-1f, 0x1.62e43p-1f}, + .offset = {0x3f2aaaab, 0x3f2aaaab, 0x3f2aaaab, 0x3f2aaaab}, + .mantissa_mask = {0x007fffff, 0x007fffff, 0x007fffff, 0x007fffff}, + .odd_coefficients = {0x1.5a9aa2p-3f, 0x1.961348p-3f, 0x1.555d7cp-2f, + -0x1.3e737cp-3f}, +}; + +// The midpoint conversion produces u = round_fp32((hash + 0.5) / 2^32), so +// 2^-33 <= u <= 1. The direct inner call is clamped to [2^-24, 1], and the +// value passed to the outer call is either the positive series result or +// -log(max(1-u, 2^-24)), approximately in [2^-33, 24*ln(2)]. Both inputs are +// therefore finite positive normal FP32 values (the smallest normal is +// 2^-126), so no special-value path is needed. +FORCE_INLINE float32x4_t positive_normal_log(float32x4_t value) { + const PositiveNormalLogData& data = kPositiveNormalLogData; + // Reduce x to x = 2^n * (1+r), evaluate the log1p polynomial in r, then + // reconstruct log(x) = n*log(2) + log1p(r). + const uint32x4_t offset_value = + vsubq_u32(vreinterpretq_u32_f32(value), data.offset); + const float32x4_t exponent = + vcvtq_f32_s32(vshrq_n_s32(vreinterpretq_s32_u32(offset_value), 23)); + const uint32x4_t mantissa = + vaddq_u32(vandq_u32(offset_value, data.mantissa_mask), data.offset); + const float32x4_t reduced = + vsubq_f32(vreinterpretq_f32_u32(mantissa), vdupq_n_f32(1.0f)); + const float32x4_t reduced_squared = vmulq_f32(reduced, reduced); + + float32x4_t first = + vfmaq_laneq_f32(data.c2, reduced, data.odd_coefficients, 0); + float32x4_t second = + vfmaq_laneq_f32(data.c4, reduced, data.odd_coefficients, 1); + float32x4_t third = + vfmaq_laneq_f32(data.c6, reduced, data.odd_coefficients, 2); + first = vfmaq_laneq_f32(first, reduced_squared, data.odd_coefficients, 3); + second = vfmaq_f32(second, first, reduced_squared); + third = vfmaq_f32(third, second, reduced_squared); + const float32x4_t leading = vfmaq_f32(reduced, data.ln2, exponent); + return vfmaq_f32(leading, third, reduced_squared); +} + +FORCE_INLINE float32x4_t log1p_neg_stable(float32x4_t value) { + // Each lane evaluates log(1-x). Use the degree-8 series near zero and the + // direct logarithm elsewhere; vbsl selects per lane without scalar branches. + float32x4_t polynomial = vdupq_n_f32(1.0f / 8.0f); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 7.0f), value, polynomial); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 6.0f), value, polynomial); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 5.0f), value, polynomial); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 4.0f), value, polynomial); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 3.0f), value, polynomial); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f / 2.0f), value, polynomial); + polynomial = vfmaq_f32(vdupq_n_f32(1.0f), value, polynomial); + const float32x4_t series = vnegq_f32(vmulq_f32(value, polynomial)); + const float32x4_t direct = + positive_normal_log(vmaxq_f32(vsubq_f32(vdupq_n_f32(1.0f), value), + vdupq_n_f32(5.960464477539063e-08f))); + return vbslq_f32(vcltq_f32(value, vdupq_n_f32(0.25f)), series, direct); +} + +FORCE_INLINE float32x4x2_t load_logits(const float* logits) { + return {vld1q_f32(logits), vld1q_f32(logits + 4)}; +} + +FORCE_INLINE float32x4x2_t load_logits(const c10::Half* logits) { + const auto values = vld1q_f16(reinterpret_cast(logits)); + return {vcvt_f32_f16(vget_low_f16(values)), + vcvt_f32_f16(vget_high_f16(values))}; +} + +FORCE_INLINE float32x4x2_t load_logits(const c10::BFloat16* logits) { +#ifdef ARM_BF16_SUPPORT + const auto values = vld1q_bf16(reinterpret_cast(logits)); + return {vcvt_f32_bf16(vget_low_bf16(values)), + vcvt_f32_bf16(vget_high_bf16(values))}; +#else + const auto values = vld1q_u16(reinterpret_cast(logits)); + return { + vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(vget_low_u16(values)), 16)), + vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(vget_high_u16(values)), 16)), + }; +#endif +} + +template +FORCE_INLINE void greedy_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, float& best_score, + int64_t& best_idx) { + constexpr int64_t VEC_SIZE = 4; + constexpr int64_t UNROLLED_SIZE = VEC_SIZE * 4; + const uint32x4_t first_offsets = {0, 1, 2, 3}; + const uint32x4_t second_offsets = + vaddq_u32(first_offsets, vdupq_n_u32(VEC_SIZE)); + const uint32x4_t third_offsets = + vaddq_u32(second_offsets, vdupq_n_u32(VEC_SIZE)); + const uint32x4_t fourth_offsets = + vaddq_u32(third_offsets, vdupq_n_u32(VEC_SIZE)); + + float32x4_t maximum = vdupq_n_f32(-std::numeric_limits::infinity()); + uint32x4_t maximum_indices = vdupq_n_u32(static_cast(token_begin)); + + int64_t token_id = token_begin; + const int64_t vector_end = + token_end - (token_end - token_begin) % UNROLLED_SIZE; + for (; token_id < vector_end; token_id += UNROLLED_SIZE) { + // Reduce 16 logits into four lanes. Lane j carries the winner among token + // offsets j, j+4, j+8, and j+12, together with its absolute token ID. + const float32x4x2_t first = load_logits(logits + token_id); + const float32x4x2_t second = load_logits(logits + token_id + 2 * VEC_SIZE); + + const float32x4_t first_half = vmaxnmq_f32(first.val[0], first.val[1]); + const uint32x4_t first_half_offsets = vbslq_u32( + vceqq_f32(first.val[0], first_half), first_offsets, second_offsets); + const float32x4_t second_half = vmaxnmq_f32(second.val[0], second.val[1]); + const uint32x4_t second_half_offsets = vbslq_u32( + vceqq_f32(second.val[0], second_half), third_offsets, fourth_offsets); + const float32x4_t block_maximum = vmaxnmq_f32(first_half, second_half); + const uint32x4_t block_offsets = + vbslq_u32(vceqq_f32(first_half, block_maximum), first_half_offsets, + second_half_offsets); + const uint32x4_t block_indices = + vaddq_u32(vdupq_n_u32(static_cast(token_id)), block_offsets); + const uint32x4_t update = vcgtq_f32(block_maximum, maximum); + maximum = vbslq_f32(update, block_maximum, maximum); + maximum_indices = vbslq_u32(update, block_indices, maximum_indices); + } + + if (token_end - token_id >= 2 * VEC_SIZE) { + const float32x4x2_t values = load_logits(logits + token_id); + const float32x4_t block_maximum = vmaxnmq_f32(values.val[0], values.val[1]); + const uint32x4_t block_offsets = vbslq_u32( + vceqq_f32(values.val[0], block_maximum), first_offsets, second_offsets); + const uint32x4_t block_indices = + vaddq_u32(vdupq_n_u32(static_cast(token_id)), block_offsets); + const uint32x4_t update = vcgtq_f32(block_maximum, maximum); + maximum = vbslq_f32(update, block_maximum, maximum); + maximum_indices = vbslq_u32(update, block_indices, maximum_indices); + token_id += 2 * VEC_SIZE; + } + + best_score = vmaxvq_f32(maximum); + // More than one lane can contain the global maximum. Select the smallest + // associated token ID to preserve first-occurrence argmax semantics. + const uint32x4_t matching_indices = + vbslq_u32(vceqq_f32(maximum, vdupq_n_f32(best_score)), maximum_indices, + vdupq_n_u32(std::numeric_limits::max())); + best_idx = vminvq_u32(matching_indices); + + for (; token_id < token_end; ++token_id) { + const float score = static_cast(logits[token_id]); + if (score > best_score) { + best_score = score; + best_idx = token_id; + } + } +} + +template +FORCE_INLINE void score_gumbel_block(const scalar_t* logits, int64_t token_id, + uint32_t token_offset, uint32x4_t hash, + uint32x4_t offsets, float temperature, + float32x4_t& winning_scores, + uint32x4_t& winning_token_offsets) { + constexpr int64_t VEC_SIZE = 8; + // For each of eight tokens, compute + // + // u = (Murmur3(seed, position, token) + 0.5) / 2^32 + // g = -log(-log(1-u)) + // scaled_logit = APPLY_TEMPERATURE ? logit / temperature : logit + // score = scaled_logit + g. + // + // The two four-lane score vectors are folded lane-wise. Output lane j is the + // winner between token offsets j and j+4; token_offset makes that ID relative + // to the caller's outer block (0, 8, 16, or 24 in the 32-token loop). + const uint32x4_t first_tokens = + vaddq_u32(vdupq_n_u32(static_cast(token_id)), offsets); + const uint32x4_t second_tokens = + vaddq_u32(first_tokens, vdupq_n_u32(VEC_SIZE / 2)); + const float32x4x2_t uniforms = {murmur3_uniform32x4(hash, first_tokens), + murmur3_uniform32x4(hash, second_tokens)}; + const float32x4x2_t inner_log = { + vnegq_f32(log1p_neg_stable(uniforms.val[0])), + vnegq_f32(log1p_neg_stable(uniforms.val[1]))}; + const float32x4x2_t noise = { + vnegq_f32(positive_normal_log(inner_log.val[0])), + vnegq_f32(positive_normal_log(inner_log.val[1]))}; + + float32x4x2_t logit = load_logits(logits + token_id); + if constexpr (APPLY_TEMPERATURE) { + // Match Triton's temperature division. Its cached-logits path divides by + // the same temperature on reload and relies on reproducing this value + // bitwise. + const float32x4_t divisor = vdupq_n_f32(temperature); + logit.val[0] = vdivq_f32(logit.val[0], divisor); + logit.val[1] = vdivq_f32(logit.val[1], divisor); + } + const float32x4_t first_scores = vaddq_f32(logit.val[0], noise.val[0]); + const float32x4_t second_scores = vaddq_f32(logit.val[1], noise.val[1]); + winning_scores = vmaxnmq_f32(first_scores, second_scores); + const uint32x4_t first_token_offsets = + vaddq_u32(offsets, vdupq_n_u32(token_offset)); + const uint32x4_t second_token_offsets = + vaddq_u32(first_token_offsets, vdupq_n_u32(VEC_SIZE / 2)); + winning_token_offsets = vbslq_u32(vceqq_f32(first_scores, winning_scores), + first_token_offsets, second_token_offsets); +} + +template +void gumbel_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, int64_t seed, int64_t position, + float temperature, float& best_score, int64_t& best_idx) { + constexpr int64_t VEC_SIZE = 8; + constexpr int64_t UNROLLED_SIZE = VEC_SIZE * 4; + const int64_t vector_end = + token_end - (token_end - token_begin) % UNROLLED_SIZE; + const uint32x4_t hash = vdupq_n_u32(murmur3_hash32_prefix(seed, position, 0)); + const uint32x4_t offsets = {0, 1, 2, 3}; + + best_score = -std::numeric_limits::infinity(); + best_idx = token_begin; + + int64_t token_id = token_begin; + for (; token_id < vector_end; token_id += UNROLLED_SIZE) { + // Keep four independent eight-token pipelines in flight to hide the log + // latency. + float32x4_t first_maximum; + uint32x4_t first_token_offsets; + score_gumbel_block( + logits, token_id, 0, hash, offsets, temperature, first_maximum, + first_token_offsets); + float32x4_t second_maximum; + uint32x4_t second_token_offsets; + score_gumbel_block( + logits, token_id + VEC_SIZE, VEC_SIZE, hash, offsets, temperature, + second_maximum, second_token_offsets); + float32x4_t third_maximum; + uint32x4_t third_token_offsets; + score_gumbel_block( + logits, token_id + 2 * VEC_SIZE, 2 * VEC_SIZE, hash, offsets, + temperature, third_maximum, third_token_offsets); + float32x4_t fourth_maximum; + uint32x4_t fourth_token_offsets; + score_gumbel_block( + logits, token_id + 3 * VEC_SIZE, 3 * VEC_SIZE, hash, offsets, + temperature, fourth_maximum, fourth_token_offsets); + + // reduce the four 8 token pipelines to find the block's argmax + const float32x4_t first_half_maximum = + vmaxnmq_f32(first_maximum, second_maximum); + const uint32x4_t first_half_token_offsets = + vbslq_u32(vceqq_f32(first_maximum, first_half_maximum), + first_token_offsets, second_token_offsets); + const float32x4_t second_half_maximum = + vmaxnmq_f32(third_maximum, fourth_maximum); + const uint32x4_t second_half_token_offsets = + vbslq_u32(vceqq_f32(third_maximum, second_half_maximum), + third_token_offsets, fourth_token_offsets); + const float32x4_t block_maximum = + vmaxnmq_f32(first_half_maximum, second_half_maximum); + const uint32x4_t block_token_offsets = + vbslq_u32(vceqq_f32(first_half_maximum, block_maximum), + first_half_token_offsets, second_half_token_offsets); + + const float score = vmaxnmvq_f32(block_maximum); + if (score > best_score) { + const uint32x4_t matching_offsets = vbslq_u32( + vceqq_f32(block_maximum, vdupq_n_f32(score)), block_token_offsets, + vdupq_n_u32(std::numeric_limits::max())); + best_score = score; + best_idx = token_id + vminvq_u32(matching_offsets); + } + } + + const int64_t vector_tail_end = token_end - (token_end - token_id) % VEC_SIZE; + for (; token_id < vector_tail_end; token_id += VEC_SIZE) { + float32x4_t maximum; + uint32x4_t token_offsets; + score_gumbel_block(logits, token_id, 0, hash, + offsets, temperature, + maximum, token_offsets); + const float score = vmaxnmvq_f32(maximum); + if (score > best_score) { + const uint32x4_t matching_offsets = + vbslq_u32(vceqq_f32(maximum, vdupq_n_f32(score)), token_offsets, + vdupq_n_u32(std::numeric_limits::max())); + best_score = score; + best_idx = token_id + vminvq_u32(matching_offsets); + } + } + + for (; token_id < token_end; ++token_id) { + float logit = static_cast(logits[token_id]); + if constexpr (APPLY_TEMPERATURE) logit /= temperature; + const float score = logit + gumbel_noise32(seed, position, token_id); + if (score > best_score) { + best_score = score; + best_idx = token_id; + } + } +} + +} // namespace cpu_sampling::neon + +namespace cpu_sampling { + +template +class SamplingImpl { + public: + using scalar_type = scalar_t; + static constexpr int64_t RangeAlignment = 32; + + void greedy_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, float& best_score, + int64_t& best_idx) const { + neon::greedy_argmax(logits, token_begin, token_end, best_score, best_idx); + } + + template + void gumbel_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, int64_t seed, int64_t position, + float temperature, score_t& best_score, + int64_t& best_idx) const { + if constexpr (std::is_same_v) { + // TODO: do we need to vectorize this? + scalar_gumbel_argmax( + logits, token_begin, token_end, seed, position, temperature, + best_score, best_idx); + } else { + neon::gumbel_argmax( + logits, token_begin, token_end, seed, position, temperature, + best_score, best_idx); + } + } +}; + +} // namespace cpu_sampling diff --git a/csrc/cpu/sampling/vec.h b/csrc/cpu/sampling/vec.h new file mode 100644 index 000000000000..bb9f28355d78 --- /dev/null +++ b/csrc/cpu/sampling/vec.h @@ -0,0 +1,30 @@ +#pragma once + +#include "common.h" + +namespace cpu_sampling { + +template +class SamplingImpl { + public: + using scalar_type = scalar_t; + static constexpr int64_t RangeAlignment = 1; + + void greedy_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, float& best_score, + int64_t& best_idx) const { + scalar_greedy_argmax(logits, token_begin, token_end, best_score, best_idx); + } + + template + void gumbel_argmax(const scalar_t* logits, int64_t token_begin, + int64_t token_end, int64_t seed, int64_t position, + float temperature, score_t& best_score, + int64_t& best_idx) const { + scalar_gumbel_argmax( + logits, token_begin, token_end, seed, position, temperature, best_score, + best_idx); + } +}; + +} // namespace cpu_sampling diff --git a/csrc/cpu/sampling_kernels.cpp b/csrc/cpu/sampling_kernels.cpp index 1d30b92a8d9a..ccca5b12ce9a 100644 --- a/csrc/cpu/sampling_kernels.cpp +++ b/csrc/cpu/sampling_kernels.cpp @@ -1,7 +1,23 @@ #include "cpu_types.hpp" +#if defined(__aarch64__) + #include "sampling/neon.h" + +template +using SelectedSamplingImpl = + cpu_sampling::SamplingImpl; +#else + #include "sampling/vec.h" + +template +using SelectedSamplingImpl = + cpu_sampling::SamplingImpl; +#endif + #include +#include +#include #include #include #include @@ -166,3 +182,133 @@ torch::Tensor greedy_argmax(const torch::Tensor& logits) { logits_contig.size(1)); return output; } + +torch::Tensor gumbel_sample_cpu(const torch::Tensor& logits, + const torch::Tensor& expanded_idx_mapping, + const torch::Tensor& temperature, + const torch::Tensor& seeds, + const torch::Tensor& positions, + bool apply_temperature, bool is_drafting, + bool use_fp64) { + TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [num_tokens, vocab]"); + TORCH_CHECK(logits.stride(1) == 1, + "the logits vocabulary dimension must be contiguous"); + + auto output = + torch::empty({logits.size(0)}, logits.options().dtype(torch::kInt64)); + VLLM_DISPATCH_FLOATING_TYPES(logits.scalar_type(), "gumbel_sample_cpu", [&] { + using sampling_impl_t = SelectedSamplingImpl; + cpu_sampling::SamplingMainLoop mainloop; + mainloop(output.data_ptr(), logits.data_ptr(), + logits.stride(0), expanded_idx_mapping.data_ptr(), + temperature.data_ptr(), seeds.data_ptr(), + positions.data_ptr(), logits.size(0), logits.size(1), + apply_temperature, is_drafting, use_fp64); + }); + return output; +} + +FORCE_INLINE void fill_logits(float* logits, int64_t size, float value) { + constexpr int64_t VEC_SIZE = vec_op::FP32Vec16::VEC_ELEM_NUM; + const vec_op::FP32Vec16 values(value); + const int64_t vectorized_size = size - size % VEC_SIZE; + int64_t offset = 0; + for (; offset < vectorized_size; offset += VEC_SIZE) { + values.save(logits + offset); + } + std::fill(logits + offset, logits + size, value); +} + +void apply_logit_bias_cpu( + torch::Tensor& logits, const torch::Tensor& expanded_idx_mapping, + const torch::Tensor& positions, const torch::Tensor& num_allowed_token_ids, + const torch::Tensor& allowed_token_ids, const torch::Tensor& num_logit_bias, + const torch::Tensor& logit_bias_token_ids, const torch::Tensor& logit_bias, + const torch::Tensor& min_lens, const torch::Tensor& num_stop_token_ids, + const torch::Tensor& restore_when_all_masked, + const torch::Tensor& stop_token_ids, bool check_all_masked_rows) { + constexpr int64_t MAX_ALLOWED_TOKEN_IDS = 1024; + constexpr int64_t MAX_STOP_TOKEN_IDS = 128; + const float negative_infinity = -std::numeric_limits::infinity(); + const int64_t num_tokens = logits.size(0); + const int64_t vocab_size = logits.size(1); + + float* logits_ptr = logits.data_ptr(); + const int64_t* idx_mapping_ptr = expanded_idx_mapping.data_ptr(); + const int64_t* positions_ptr = positions.data_ptr(); + const int32_t* num_allowed_ptr = num_allowed_token_ids.data_ptr(); + const int32_t* allowed_ptr = allowed_token_ids.data_ptr(); + const int32_t* num_bias_ptr = num_logit_bias.data_ptr(); + const int32_t* bias_token_ids_ptr = logit_bias_token_ids.data_ptr(); + const float* bias_ptr = logit_bias.data_ptr(); + const int32_t* min_lens_ptr = min_lens.data_ptr(); + const int32_t* num_stop_ptr = num_stop_token_ids.data_ptr(); + const int32_t* restore_ptr = restore_when_all_masked.data_ptr(); + const int32_t* stop_token_ids_ptr = stop_token_ids.data_ptr(); + + const int64_t logits_stride = logits.stride(0); + const int64_t allowed_stride = allowed_token_ids.stride(0); + const int64_t bias_ids_stride = logit_bias_token_ids.stride(0); + const int64_t bias_stride = logit_bias.stride(0); + const int64_t stop_stride = stop_token_ids.stride(0); + +#pragma omp parallel for schedule(static) if (num_tokens > 1) + for (int64_t token_idx = 0; token_idx < num_tokens; ++token_idx) { + const int64_t req_idx = idx_mapping_ptr[token_idx]; + float* row = logits_ptr + token_idx * logits_stride; + + const int32_t num_allowed = num_allowed_ptr[req_idx]; + if (num_allowed > 0) { + std::array saved_logits; + const int32_t* row_allowed = allowed_ptr + req_idx * allowed_stride; + for (int32_t i = 0; i < num_allowed; ++i) { + saved_logits[i] = row[row_allowed[i]]; + } + fill_logits(row, vocab_size, negative_infinity); + for (int32_t i = 0; i < num_allowed; ++i) { + row[row_allowed[i]] = saved_logits[i]; + } + } + + const int32_t num_bias = num_bias_ptr[req_idx]; + const int32_t* row_bias_ids = + bias_token_ids_ptr + req_idx * bias_ids_stride; + const float* row_bias = bias_ptr + req_idx * bias_stride; + for (int32_t i = 0; i < num_bias; ++i) { + row[row_bias_ids[i]] += row_bias[i]; + } + + const int32_t num_stop = num_stop_ptr[req_idx]; + if (num_stop == 0 || + positions_ptr[token_idx] + 1 >= min_lens_ptr[req_idx]) { + continue; + } + + const int32_t* row_stop_ids = stop_token_ids_ptr + req_idx * stop_stride; + if (!check_all_masked_rows || restore_ptr[req_idx] == 0) { + for (int32_t i = 0; i < num_stop; ++i) { + row[row_stop_ids[i]] = negative_infinity; + } + continue; + } + + std::array saved_stop_logits; + for (int32_t i = 0; i < num_stop; ++i) { + saved_stop_logits[i] = row[row_stop_ids[i]]; + row[row_stop_ids[i]] = negative_infinity; + } + + const bool all_masked = + std::all_of(row, row + vocab_size, [negative_infinity](float value) { + return value == negative_infinity; + }); + if (all_masked) { + for (int32_t i = 0; i < num_stop; ++i) { + const float value = saved_stop_logits[i]; + if (std::isfinite(value)) { + row[row_stop_ids[i]] = value; + } + } + } + } +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index c68d5c6317df..cffa908841e3 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -406,9 +406,77 @@ void dynamic_per_token_scaled_fp8_quant( torch::Tensor& out, const torch::Tensor& input, torch::Tensor& scale, const std::optional& scale_ub); +torch::Tensor gumbel_sample_cpu(const torch::Tensor& logits, + const torch::Tensor& expanded_idx_mapping, + const torch::Tensor& temperature, + const torch::Tensor& seeds, + const torch::Tensor& positions, + bool apply_temperature, bool is_drafting, + bool use_fp64); + +void apply_logit_bias_cpu( + torch::Tensor& logits, const torch::Tensor& expanded_idx_mapping, + const torch::Tensor& positions, const torch::Tensor& num_allowed_token_ids, + const torch::Tensor& allowed_token_ids, const torch::Tensor& num_logit_bias, + const torch::Tensor& logit_bias_token_ids, const torch::Tensor& logit_bias, + const torch::Tensor& min_lens, const torch::Tensor& num_stop_token_ids, + const torch::Tensor& restore_when_all_masked, + const torch::Tensor& stop_token_ids, bool check_all_masked_rows); + void init_cpu_memory_env(std::vector node_ids); namespace cpu_utils { +// Kernels used only by standalone CPU MRV2. +#ifdef VLLM_STANDALONE_CPU_MRV2 +void gather_block_table_cpu(const torch::Tensor& block_table, + torch::Tensor& output, + const torch::Tensor& idx_mapping, + const torch::Tensor& num_blocks, + const int64_t num_reqs_padded); +void compute_slot_mapping_cpu( + const torch::Tensor& block_table, torch::Tensor& slot_mapping, + const torch::Tensor& idx_mapping, const torch::Tensor& query_start_loc, + const torch::Tensor& positions, const int64_t block_size, + const bool mapping_enabled, const int64_t pad_slot_id); +void prepare_prefill_inputs_cpu(torch::Tensor& input_ids, + torch::Tensor& next_prefill_tokens, + const torch::Tensor& idx_mapping, + const torch::Tensor& query_start_loc, + const torch::Tensor& all_token_ids, + const torch::Tensor& prefill_len, + const torch::Tensor& num_computed_tokens); +torch::Tensor get_prompt_logprobs_token_ids_cpu( + int64_t num_tokens, const torch::Tensor& query_start_loc, + const torch::Tensor& idx_mapping, const torch::Tensor& num_computed_tokens, + const torch::Tensor& all_token_ids); +void prepare_pos_seq_lens_cpu(const torch::Tensor& idx_mapping, + const torch::Tensor& query_start_loc, + const torch::Tensor& num_computed_tokens, + torch::Tensor& pos, torch::Tensor& seq_lens); +void prepare_mrope_positions_cpu(torch::Tensor& positions, + const torch::Tensor& prefill_positions, + const torch::Tensor& prefill_delta, + const torch::Tensor& idx_mapping, + const torch::Tensor& query_start_loc, + const torch::Tensor& prefill_lens, + const torch::Tensor& num_computed_tokens); +torch::Tensor combine_sampled_and_draft_tokens_cpu( + torch::Tensor& input_ids, const torch::Tensor& idx_mapping, + const torch::Tensor& last_sampled_tokens, + const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens, + const torch::Tensor& prefill_len, const torch::Tensor& draft_tokens, + const torch::Tensor& cu_num_logits, const int64_t num_logits, + const int64_t num_new_sampled_tokens); +void post_update_cpu(const torch::Tensor& idx_mapping, + torch::Tensor& num_computed_tokens, + torch::Tensor& last_sampled_tokens, + const c10::optional& output_bin_counts, + const torch::Tensor& sampled_tokens, + const torch::Tensor& num_sampled, + const torch::Tensor& num_rejected, + const c10::optional& query_start_loc, + torch::Tensor& all_token_ids, torch::Tensor& total_len); +#endif void eagle_prepare_inputs_padded_kernel_impl( const torch::Tensor& cu_num_draft_tokens, const torch::Tensor& valid_sampled_tokens_count, @@ -999,6 +1067,62 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); +#ifdef VLLM_STANDALONE_CPU_MRV2 + ops.def( + "gather_block_table_cpu(Tensor block_table, Tensor(a1!) output, " + "Tensor idx_mapping, Tensor num_blocks, SymInt num_reqs_padded) -> ()", + &cpu_utils::gather_block_table_cpu); + + ops.def( + "compute_slot_mapping_cpu(" + "Tensor block_table, Tensor(a1!) slot_mapping, Tensor idx_mapping, " + "Tensor query_start_loc, Tensor positions, SymInt block_size, " + "bool mapping_enabled, SymInt pad_slot_id) -> ()", + &cpu_utils::compute_slot_mapping_cpu); + + ops.def( + "prepare_prefill_inputs_cpu(Tensor(a0!) input_ids, " + "Tensor(a1!) next_prefill_tokens, Tensor idx_mapping, " + "Tensor query_start_loc, Tensor all_token_ids, Tensor prefill_len, " + "Tensor num_computed_tokens) -> ()", + &cpu_utils::prepare_prefill_inputs_cpu); + + ops.def( + "get_prompt_logprobs_token_ids_cpu(" + "SymInt num_tokens, Tensor query_start_loc, Tensor idx_mapping, " + "Tensor num_computed_tokens, Tensor all_token_ids) -> Tensor", + &cpu_utils::get_prompt_logprobs_token_ids_cpu); + + ops.def( + "prepare_pos_seq_lens_cpu(Tensor idx_mapping, Tensor query_start_loc, " + "Tensor num_computed_tokens, Tensor(a3!) pos, " + "Tensor(a4!) seq_lens) -> ()", + &cpu_utils::prepare_pos_seq_lens_cpu); + + ops.def( + "prepare_mrope_positions_cpu(" + "Tensor(a0!) positions, Tensor prefill_positions, Tensor prefill_delta, " + "Tensor idx_mapping, Tensor query_start_loc, Tensor prefill_lens, " + "Tensor num_computed_tokens) -> ()", + &cpu_utils::prepare_mrope_positions_cpu); + + ops.def( + "combine_sampled_and_draft_tokens_cpu(" + "Tensor(a0!) input_ids, Tensor idx_mapping, Tensor last_sampled_tokens, " + "Tensor query_start_loc, Tensor seq_lens, Tensor prefill_len, " + "Tensor draft_tokens, Tensor cu_num_logits, SymInt num_logits, " + "SymInt num_new_sampled_tokens) -> Tensor", + &cpu_utils::combine_sampled_and_draft_tokens_cpu); + + ops.def( + "post_update_cpu(Tensor idx_mapping, " + "Tensor(a1!) num_computed_tokens, Tensor(a2!) last_sampled_tokens, " + "Tensor(a3!)? output_bin_counts, Tensor sampled_tokens, " + "Tensor num_sampled, Tensor num_rejected, Tensor? query_start_loc, " + "Tensor(a8!) all_token_ids, Tensor(a9!) total_len) -> ()", + &cpu_utils::post_update_cpu); +#endif + // Fused sampling kernels ops.def("fused_gumbel_argmax(Tensor logits, Tensor seeds) -> Tensor"); ops.impl("fused_gumbel_argmax", torch::kCPU, &fused_gumbel_argmax); @@ -1006,6 +1130,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("greedy_argmax(Tensor logits) -> Tensor"); ops.impl("greedy_argmax", torch::kCPU, &greedy_argmax); + ops.def( + "gumbel_sample_cpu(Tensor logits, Tensor expanded_idx_mapping, " + "Tensor temperature, Tensor seeds, Tensor positions, " + "bool apply_temperature, bool is_drafting, bool use_fp64) -> Tensor"); + ops.impl("gumbel_sample_cpu", torch::kCPU, &gumbel_sample_cpu); + + ops.def( + "apply_logit_bias_cpu(Tensor(a0!) logits, " + "Tensor expanded_idx_mapping, Tensor positions, " + "Tensor num_allowed_token_ids, Tensor allowed_token_ids, " + "Tensor num_logit_bias, Tensor logit_bias_token_ids, " + "Tensor logit_bias, Tensor min_lens, Tensor num_stop_token_ids, " + "Tensor restore_when_all_masked, Tensor stop_token_ids, " + "bool check_all_masked_rows) -> ()", + &apply_logit_bias_cpu); + // Speculative decoding kernels ops.def( "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp index 60d3117f24db..0f52c12a6a7c 100644 --- a/csrc/fs_io.cpp +++ b/csrc/fs_io.cpp @@ -71,7 +71,11 @@ inline int _store_block(const char* tmp_path, const char* dest_path, if (rename(tmp_path, dest_path) != 0) { const int err = errno; unlink(tmp_path); - return err; + // Another writer published this block after the access() above; some + // shared filesystems fail rename(2) with EEXIST instead of replacing. + // File names are content hashes, so the file already in place holds this + // block and the store is done. + return err == EEXIST ? 0 : err; } return 0; diff --git a/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu index 09484daf3315..11e76e6d9873 100644 --- a/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu +++ b/csrc/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu @@ -48,9 +48,11 @@ * ``qkv`` tensor. Caches (bf16) are scatter-written by slot. */ +#include #include #include #include +#include #include "torch_utils.h" @@ -271,11 +273,43 @@ __device__ __forceinline__ void storeElemsFp8( // Match scaled_fp8_quant(q_out): materialize q in scalar_t before applying the // inverse dequantization scale and converting it to E4M3. +#ifndef USE_ROCM +// storeScaledQElemsFp8 with inv_scale = 1, from the packed model-dtype values: +// E4M3 of a scalar_t equals E4M3 of its float, the max takes NaN and values +// below -448 to -448 as storeElemsFp8's fmaxf does, and the cvt saturates the +// top at 448. +template +__device__ __forceinline__ void storeUnitScaleQElemsFp8( + uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) { + using Converter = vllm::_typeConvert; + __nv_fp8x2_storage_t out2[kElemsPerLane / 2]; + #pragma unroll + for (int i = 0; i < kElemsPerLane / 2; i++) { + auto const v = + Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1])); + if constexpr (std::is_same_v) { + out2[i] = __nv_cvt_bfloat16raw2_to_fp8x2( + __hmax2(v, __float2bfloat162_rn(-448.0f)), __NV_SATFINITE, __NV_E4M3); + } else { + out2[i] = __nv_cvt_halfraw2_to_fp8x2( + __hmax2(v, __float2half2_rn(-448.0f)), __NV_SATFINITE, __NV_E4M3); + } + } + *reinterpret_cast(dst) = *reinterpret_cast(out2); +} +#endif + template __device__ __forceinline__ void storeScaledQElemsFp8( uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane], float const inv_scale) { using Converter = vllm::_typeConvert; +#ifndef USE_ROCM + if (inv_scale == 1.0f) { + storeUnitScaleQElemsFp8(dst, elems); + return; + } +#endif float scaled[kElemsPerLane]; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -424,27 +458,14 @@ __device__ __forceinline__ uint16_t clearNegZeroNibbles(uint16_t packed) { return packed & (((packed & 0x7777u) + 0x7777u) | 0x7777u); } -// Quantize one K or V head row of one token into its NVFP4 page, where -// ``scale`` is the dequantization (global) scale. Every lane of the warp must -// call this (the scale-group reductions shuffle). +// The E2M1 codes (low nibble first) of four values x, already rounded to the +// model dtype, of a 16-value group with absolute maximum group_max, and the +// group's E4M3 scale in ``sf_byte``, where ``scale`` is the dequantization +// (global) scale. template -__device__ __forceinline__ void storeNvfp4CacheElems( - uint8_t* __restrict__ page, int const kv, int const head, int const token, - int const block_size, int const laneId, float const (&elems)[kElemsPerLane], - float const scale) { - // Quantize the model-dtype value that the unfused path reads back from qkv. - using Converter = vllm::_typeConvert; - float x[kElemsPerLane]; - float group_max = 0.0f; - #pragma unroll - for (int i = 0; i < kElemsPerLane; i++) { - x[i] = static_cast(Converter::convert(elems[i])); - group_max = fmaxf(group_max, fabsf(x[i])); - } - group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32)); - group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32)); - - uint8_t sf_byte; +__device__ __forceinline__ uint16_t +nvfp4EncodeElems(float const (&x)[kElemsPerLane], float const group_max, + float const scale, uint8_t& sf_byte) { float sf_rcp; float q[kElemsPerLane]; // (x / scale) / sf_value, E2M1-encoded below if (scale == 1.0f) { @@ -516,8 +537,49 @@ __device__ __forceinline__ void storeNvfp4CacheElems( for (int i = 0; i < kElemsPerLane; i++) codes |= e2m1Code(q[i]) << (4 * i); packed = static_cast(codes); #endif - packed = clearNegZeroNibbles(packed); + return clearNegZeroNibbles(packed); +} + +// Offset, within a K (kv 0) or V (kv 1) head's scale block, of the E4M3 scale +// of a token's 16-value group: linear for K, the SM100 trtllm-gen 4x4 +// token-quad swizzle for V. +__device__ __forceinline__ int nvfp4ScaleOffset(int const kv, int const token, + int const group) { + if (kv == 0) return token * kNvfp4ScaleBytes + group; + constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4; + int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad; + int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4; + return swizzled_t * kNvfp4ScaleBytes + swizzled_s; +} + +// The E2M1 codes of this lane's four values of a K or V head row, rounded to +// the model dtype, and in ``sf_byte`` its 16-value group's E4M3 scale, where +// ``scale`` is the dequantization (global) scale. Every lane of the warp must +// call this (the scale-group reductions shuffle). +template +__device__ __forceinline__ uint16_t nvfp4QuantizeElems( + float const (&elems)[kElemsPerLane], float const scale, uint8_t& sf_byte) { + // Quantize the model-dtype value that the unfused path reads back from qkv. + using Converter = vllm::_typeConvert; + float x[kElemsPerLane]; + float group_max = 0.0f; + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + x[i] = static_cast(Converter::convert(elems[i])); + group_max = fmaxf(group_max, fabsf(x[i])); + } + group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 1, 32)); + group_max = fmaxf(group_max, __shfl_xor_sync(FINAL_MASK, group_max, 2, 32)); + + return nvfp4EncodeElems(x, group_max, scale, sf_byte); +} +// Store the codes and scale from nvfp4QuantizeElems of lane laneId's four +// values of a token's K (kv 0) or V (kv 1) head row. +__device__ __forceinline__ void nvfp4StoreElems( + uint8_t* __restrict__ page, int const kv, int const head, int const token, + int const block_size, int const laneId, uint16_t const packed, + uint8_t const sf_byte) { uint8_t* const slot = page + static_cast(2 * head + kv) * block_size * (kNvfp4DataBytes + kNvfp4ScaleBytes); @@ -526,19 +588,57 @@ __device__ __forceinline__ void storeNvfp4CacheElems( *reinterpret_cast(data) = packed; if (laneId % kNvfp4LanesPerGroup == 0) { - int const group = laneId / kNvfp4LanesPerGroup; - int offset; - if (kv == 0) { - offset = token * kNvfp4ScaleBytes + group; - } else { - constexpr int kGroupsPerQuad = kNvfp4ScaleBytes / 4; - int const swizzled_t = (token / 4) * 4 + group / kGroupsPerQuad; - int const swizzled_s = (group % kGroupsPerQuad) * 4 + token % 4; - offset = swizzled_t * kNvfp4ScaleBytes + swizzled_s; - } - slot[block_size * kNvfp4DataBytes + offset] = sf_byte; + slot[block_size * kNvfp4DataBytes + + nvfp4ScaleOffset(kv, token, laneId / kNvfp4LanesPerGroup)] = sf_byte; } } + +// nvfp4QuantizeElems for one lane holding a whole 16-value group x (already +// rounded to the model dtype): the group max is order-independent, so this is +// the same codes and scale without the shuffles. Low nibble first. +template +__device__ __forceinline__ uint2 nvfp4QuantizeGroup(float const (&x)[16], + float const scale, + uint8_t& sf_byte) { + float group_max = 0.0f; + #pragma unroll + for (int i = 0; i < 16; i++) group_max = fmaxf(group_max, fabsf(x[i])); + uint32_t codes[4]; + #pragma unroll + for (int k = 0; k < 4; k++) { + float const xk[kElemsPerLane] = {x[4 * k], x[4 * k + 1], x[4 * k + 2], + x[4 * k + 3]}; + codes[k] = nvfp4EncodeElems(xk, group_max, scale, sf_byte); + } + return make_uint2(codes[0] | codes[1] << 16, codes[2] | codes[3] << 16); +} + +// Store group ``group`` (its 8 code bytes and its scale) of a token's K +// (kv 0) or V (kv 1) head row into its NVFP4 page. +__device__ __forceinline__ void nvfp4StoreGroup( + uint8_t* __restrict__ page, int const kv, int const head, int const token, + int const block_size, int const group, uint2 const codes, + uint8_t const sf_byte) { + uint8_t* const slot = page + static_cast(2 * head + kv) * + block_size * + (kNvfp4DataBytes + kNvfp4ScaleBytes); + *reinterpret_cast(slot + token * kNvfp4DataBytes + 8 * group) = codes; + slot[block_size * kNvfp4DataBytes + nvfp4ScaleOffset(kv, token, group)] = + sf_byte; +} + +// Quantize one K or V head row of one token into its NVFP4 page, where +// ``scale`` is the dequantization (global) scale. Every lane of the warp must +// call this (the scale-group reductions shuffle). +template +__device__ __forceinline__ void storeNvfp4CacheElems( + uint8_t* __restrict__ page, int const kv, int const head, int const token, + int const block_size, int const laneId, float const (&elems)[kElemsPerLane], + float const scale) { + uint8_t sf_byte; + uint16_t const packed = nvfp4QuantizeElems(elems, scale, sf_byte); + nvfp4StoreElems(page, kv, head, token, block_size, laneId, packed, sf_byte); +} #endif // USE_ROCM // ──────────────────────────────────────────────────────────────────────────── @@ -553,18 +653,21 @@ __device__ __forceinline__ void storeNvfp4CacheElems( // Slots per token: // Q : nq (always — norm+RoPE) // K : nkv (always — norm+RoPE; +K-cache insert) -// V : nkv only if kInsertKV (V-cache insert; no warps in dense) +// V : nkv only if kInsertKV (V-cache insert) // IQ: niq only if kProcessIndex (norm+RoPE) // IK: 1 only if kProcessIndex (norm+RoPE; +index-cache insert) // cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx: // indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte). +// kNvfp4Idx (implies kFp8Idx): the index-K cache is instead packed NVFP4 pages +// [nb, bs, 72] (bs*64 B E2M1, then bs*8 B linear E4M3 scales, no global scale), +// the layout of the MSA Q8KV4 indexer; index-Q stays e4m3. // kHasIndex means the qkv row is laid out as sparse [q|k|v|index_q|index_k]. // kProcessIndex controls whether this launch actually norms/ropes the index // branch and writes index_q/index_k outputs. Skip-index-topk reuse layers keep // kHasIndex=true but set kProcessIndex=false. template + bool kProcessIndex, bool kFp8Idx, bool kNvfp4Idx> __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse) scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr @@ -588,6 +691,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( // NHD/HND layout choice is captured by the head/token strides. int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, int64_t const kv_s_dim, + // NVFP4 index cache only: page stride in bytes. + int64_t const idx_s_block, // NVFP4 main cache only: per-layer K/V dequantization (global) scales. float const* __restrict__ kv_k_scale, float const* __restrict__ kv_v_scale) { @@ -604,8 +709,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( static_assert(!kProcessIndex || kHasIndex, "index processing requires sparse row layout"); + static_assert(!kNvfp4Idx || (kFp8Idx && kInsertKV), + "NVFP4 index cache requires fp8 index-Q and cache insert"); - // Slot layout (compile-time gated: dense has neither V nor index slots). + // Slot layout (compile-time gated: V slots only when inserting, index + // slots only on the sparse layer). int const v_slots = kInsertKV ? nkv : 0; int const idx_slots = kProcessIndex ? niq + 1 : 0; int const slots_per_token = nq + nkv + v_slots + idx_slots; @@ -743,7 +851,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( } } - // ── Cache inserts (sparse serving only). ─────────────────────────────── + // ── Cache inserts. ───────────────────────────────────────────────────── if constexpr (kInsertKV) { // Guard (not early-return) so every thread reaches the PDL trigger below. int64_t sm = -1; @@ -754,7 +862,16 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( } if (sm >= 0) { // skip padded / unscheduled tokens if (isIK) { - if constexpr (kFp8Idx) { + if constexpr (kNvfp4Idx) { +#ifndef USE_ROCM + // A single-head K side: data [bs, 64] then linear scales [bs, 8]. + storeNvfp4CacheElems( + reinterpret_cast(index_cache) + + (sm / block_size) * idx_s_block, + /*kv=*/0, /*head=*/0, static_cast(sm % block_size), + block_size, laneId, elems, /*scale=*/1.0f); +#endif + } else if constexpr (kFp8Idx) { storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems); } else { storeElems(index_cache + sm * kHeadDim + dim_base, elems); @@ -779,6 +896,15 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( int64_t const off = b * kv_s_block + head * kv_s_head + t * kv_s_token + (kv * kHeadDim + dim_base) * kv_s_dim; + if constexpr (!kHasIndex) { + // Dense layers stand in for reshape_and_cache_flash, which + // quantizes the k already materialized in the model dtype. + using Converter = vllm::_typeConvert; +#pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + elems[i] = Converter::convert(Converter::convert(elems[i])); + } + } storeCacheElems(kv_cache + off, elems); } } @@ -794,6 +920,754 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel( #endif } +#ifndef USE_ROCM +// ──────────────────────────────────────────────────────────────────────────── +// Persistent row-group kernel +// ──────────────────────────────────────────────────────────────────────────── +// Writes the same bytes as the kernel above, for every instantiation of it, +// when rotary_dim == 64. A token's work is a group of warps with fixed roles: +// the q heads kQHeads (4 or 8) per warp, then K, V and index_q four heads per +// warp, then index_k. A branch of fewer than four heads, and index_k, takes +// one warp per head in the layout of the kernel above (lane l holds dims +// 4l .. 4l + 3); the others run four-row steps, where lane (g, v) holds dims +// 32c + 4v .. 32c + 4v + 3 (c = 0..3) of row g, the dims lane 8c + v holds +// above, so the stores call the same helpers with that lane id. A CTA runs +// some or all warps of the warp groups of tokens x, x + gridDim.x, ..., and +// each warp copies the qkv rows and cos/sin row of its next tokens into a +// kStages-deep shared-memory ring with cp.async. +namespace rg { + +constexpr int kRows = 4; +constexpr int kChunks = 4; +constexpr int kRotaryDim = 64; +constexpr int kStages = 3; // ring depth: two tokens in flight per warp +// A ring row: 256 B padded to 320 B, so that rows 1 apart start 16 banks apart +// and a 4-row step's 8-byte loads are conflict-free. +constexpr int kRowPitch = kHeadDim * 2 + 64; +constexpr int kCosSinBytes = kRotaryDim * 2; +enum Role : int { kQ = 0, kK, kV, kIQ, kIK }; + +// SM100's packed FP32 instructions where they exist; the same RN results. +__device__ __forceinline__ float2 mul2(float2 const a, float2 const b) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __fmul2_rn(a, b); + #else + return make_float2(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y)); + #endif +} + +__device__ __forceinline__ float2 fma2(float2 const a, float2 const b, + float2 const c) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __ffma2_rn(a, b, c); + #else + return make_float2(__fmaf_rn(a.x, b.x, c.x), __fmaf_rn(a.y, b.y, c.y)); + #endif +} + +__device__ __forceinline__ float2 add2(float2 const a, float2 const b) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + return __fadd2_rn(a, b); + #else + return make_float2(__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y)); + #endif +} + +// A head row as one lane holds it, paired for packed FP32 math: +// x[k][i] = element i of chunks 2k and 2k + 1. +using Row = float2[2][4]; + +template +__device__ __forceinline__ void loadRow(uint8_t const* row, int const v, + Row& x) { + float e[kChunks][kElemsPerLane]; + #pragma unroll + for (int c = 0; c < kChunks; c++) { + loadElems(reinterpret_cast(row + 64 * c + 8 * v), + e[c]); + } + #pragma unroll + for (int k = 0; k < 2; k++) { + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + x[k][i] = make_float2(e[2 * k][i], e[2 * k + 1][i]); + } + } +} + +__device__ __forceinline__ void chunkElems(Row const& x, int const c, + float (&e)[kElemsPerLane]) { + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + e[i] = c % 2 ? x[c / 2][i].y : x[c / 2][i].x; + } +} + +// normAndRope (Gemma RMSNorm, then NeoX RoPE on dims [0, 64)) with its +// rounding as compiled: an FFMA sum of squares per lane of the one-warp-per-row +// layout (one chunk here), the xor-16/8/4/2/1 tree (16 pairs chunks c and +// c ^ 2, 8 pairs c and c ^ 1), FFMA(sum, 1/128, eps), rsqrtf, x * rms * w, and +// x * cos -/+ RN(partner * sin) as one FFMA. +__device__ __forceinline__ void normRope(Row& x, Row const& w1, float const eps, + float const (&cos4)[4], + float const (&sin4)[4]) { + float2 sq[2]; + #pragma unroll + for (int k = 0; k < 2; k++) { + sq[k] = mul2(x[k][0], x[k][0]); // = FFMA(x, x, 0) + #pragma unroll + for (int i = 1; i < 4; i++) sq[k] = fma2(x[k][i], x[k][i], sq[k]); + } + float2 const xor16 = add2(sq[0], sq[1]); + float sum = __fadd_rn(xor16.x, xor16.y); + sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 4, 32)); + sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 2, 32)); + sum = __fadd_rn(sum, __shfl_xor_sync(FINAL_MASK, sum, 1, 32)); + float const rms_rcp = + rsqrtf(__fmaf_rn(sum, 1.0f / static_cast(kHeadDim), eps)); + float2 const r2 = make_float2(rms_rcp, rms_rcp); + #pragma unroll + for (int k = 0; k < 2; k++) { + #pragma unroll + for (int i = 0; i < 4; i++) x[k][i] = mul2(mul2(x[k][i], r2), w1[k][i]); + } + // x[0][i] = (x1, x2): dims 4v + i and 32 + 4v + i. + #pragma unroll + for (int i = 0; i < 4; i++) { + float2 const x_sin = mul2(x[0][i], make_float2(sin4[i], sin4[i])); + x[0][i] = make_float2(__fmaf_rn(x[0][i].x, cos4[i], -x_sin.y), + __fmaf_rn(x[0][i].y, cos4[i], x_sin.x)); + } +} + +__device__ __forceinline__ void cpAsync16(uint32_t const dst, void const* src) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst), + "l"(src) + : "memory"); +} + +__device__ __forceinline__ void cpAsyncCommit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void cpAsyncWait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(kPending) : "memory"); +} + +// Warp roles of a token's warp group: the q heads kQHeads at a time, then K, +// V and index_q four heads at a time, then index_k. +// A branch of fewer than kRows heads (and index_k) takes one warp per head in +// the layout of the kernel above instead of a part-empty four-row step. +struct Roles { + int nq, nkv, niq; + bool kv_one, iq_one; // one warp per head + int q_warps, kv_warps, v_warps, iq_warps, ik_warps; + int warps; // one token's warp group + + __host__ __device__ Roles(int const q_heads, int const nq_, int const nkv_, + int const niq_, bool const insert_kv, + bool const process_index) + : nq(nq_), nkv(nkv_), niq(niq_) { + kv_one = nkv < kRows; + iq_one = niq < kRows; + q_warps = (nq + q_heads - 1) / q_heads; + kv_warps = kv_one ? nkv : (nkv + kRows - 1) / kRows; + v_warps = insert_kv ? kv_warps : 0; + iq_warps = process_index ? (iq_one ? niq : (niq + kRows - 1) / kRows) : 0; + ik_warps = process_index ? 1 : 0; + warps = q_warps + kv_warps + v_warps + iq_warps + ik_warps; + } + + // Shared-memory offset of warp w's ring: the q warps' rings, then the + // others'. + __host__ __device__ int ringOffset(int const w, int const q_ring, + int const side_ring) const { + return min(w, q_warps) * q_ring + max(w - q_warps, 0) * side_ring; + } +}; + +// CTA size limit, for 80 registers per thread. +constexpr int kMaxWarps = 24; + +template +constexpr int stageBytes() { + return kQHeads * kRowPitch + kCosSinBytes; +} + +// What a warp writes besides the caches, fixed per launch so that the stores +// of the other modes are not compiled into the loop: q (kQFp8Unit: only +// q_fp8_out, unit scale; kQFp8: only q_fp8_out; kQAny: decided per store), +// index_q (kIQOut: index_q_out; kIQInPlace: back into qkv; kNone), and the +// in-place k / index_k rewrite (kInPlace or kNone). +enum Out : int { + kNone = 0, + kInPlace, + kQFp8Unit, + kQFp8, + kQAny, + kIQOut, + kIQInPlace +}; + +// One warp's loop over its CTA's tokens: rows [head0, head0 + nrows) of one +// branch, kSteps four-row steps per token. +template +__device__ __forceinline__ void runWarp( + int const head0, int const nrows, int const qkv_head0, uint8_t* const ring, + int const stage_bytes, scalar_t* __restrict__ qkv, + scalar_t* __restrict__ q_out, uint8_t* __restrict__ q_fp8_out, + out_idx_t* __restrict__ index_q_out, scalar_t const* __restrict__ norm_w, + scalar_t const* __restrict__ cos_sin_cache, + int64_t const* __restrict__ positions, + int64_t const* __restrict__ slot_mapping, cache_t* __restrict__ kv_cache, + out_idx_t* __restrict__ index_cache, float const eps, + float const q_fp8_inv_scale, int const num_tokens, int const nq, + int const niq, int const qkv_row, int const block_size, + int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, + int64_t const kv_s_dim, int64_t const idx_s_block, float const kv_scale, + bool const write_qkv) { + constexpr bool kNorm = kRole != kV; + constexpr bool kSlots = + kInsertKV && (kRole == kK || kRole == kV || kRole == kIK); + // kOneRow: a single row in the layout of the kernel above (lane l holds + // dims 4l .. 4l + 3), for branches of one head. + constexpr int kStepRows = kOneRow ? 1 : kRows; + constexpr int kLaneChunks = kOneRow ? 1 : kChunks; + int const lane = threadIdx.x % 32; + int const g = kOneRow ? 0 : lane / 8; + int const v = kOneRow ? lane : lane % 8; + constexpr int kCopies = (kSteps * kStepRows * 16 + 31) / 32; + uint32_t const ring_smem = + static_cast(__cvta_generic_to_shared(ring)); + + // Not produced by the predecessor: loaded before the PDL wait, used after. + uint2 w_raw[kChunks]; + if constexpr (kNorm && !kOneRow) { + #pragma unroll + for (int c = 0; c < kChunks; c++) { + w_raw[c] = __ldg(reinterpret_cast(norm_w + 32 * c + 4 * v)); + } + } + + // PDL: qkv comes from the predecessor (the qkv GEMM). + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); + #endif + + Row w1; + if constexpr (kNorm && !kOneRow) { + float e[kChunks][kElemsPerLane]; + #pragma unroll + for (int c = 0; c < kChunks; c++) { + loadElems(reinterpret_cast(&w_raw[c]), e[c]); + } + #pragma unroll + for (int k = 0; k < 2; k++) { + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + w1[k][i] = make_float2(1.0f + e[2 * k][i], 1.0f + e[2 * k + 1][i]); + } + } + } + + int const stride = gridDim.x; + int const t0 = blockIdx.x; + int const n = (num_tokens - t0 + stride - 1) / stride; + // Lane j holds the position (slot) of the token of iteration base + j. + auto fetch = [&](int64_t const* src, int const base, int64_t const fill) { + int64_t const t = t0 + static_cast(base + lane) * stride; + return t < num_tokens ? static_cast(__ldg( + reinterpret_cast(src) + t)) + : fill; + }; + int pos_base = 0, slot_base = 0; + int64_t pos_l = 0, slot_l = -1; + if constexpr (kNorm) pos_l = fetch(positions, 0, 0); + if constexpr (kSlots) slot_l = fetch(slot_mapping, 0, -1); + + scalar_t const* const src0 = + qkv + static_cast(t0) * qkv_row + qkv_head0 * kHeadDim; + int64_t const src_step = static_cast(stride) * qkv_row; + auto issue_rows = [&](int const it) { + if (it >= n) return; + scalar_t const* const src = src0 + it * src_step; + uint32_t const stage = ring_smem + (it % kStages) * stage_bytes; + #pragma unroll + for (int k = 0; k < kCopies; k++) { + int const r = lane / 16 + 2 * k; + if (r < nrows && (!kOneRow || lane < 16)) { + cpAsync16(stage + r * kRowPitch + lane % 16 * 16, + src + r * kHeadDim + lane % 16 * 8); + } + } + }; + // The cos/sin row of iteration it, then the commit that closes its group. + auto issue_cos_sin = [&](int const it) { + if constexpr (kNorm) { + if (it < n) { + if (it - pos_base == 32) { + pos_base = it; + pos_l = fetch(positions, pos_base, 0); + } + int64_t const pos = __shfl_sync( + FINAL_MASK, static_cast(pos_l), it - pos_base, 32); + if (lane < kCosSinBytes / 16) { + cpAsync16(ring_smem + (it % kStages) * stage_bytes + + kSteps * kStepRows * kRowPitch + lane * 16, + cos_sin_cache + pos * kRotaryDim + lane * 8); + } + } + } + cpAsyncCommit(); + }; + + // The prologue's rows do not wait for the positions. + #pragma unroll + for (int it = 0; it < kStages - 1; it++) issue_rows(it); + #pragma unroll + for (int it = 0; it < kStages - 1; it++) issue_cos_sin(it); + + // This lane's dims 4v .. 4v + 3 (+ 32c) of row g of the first step, at the + // first token; each iteration advances them by stride tokens. Dims + // 32c + 4v of a row are lane 8c + v's dims in the kernel above. + scalar_t* in_place = qkv + static_cast(t0) * qkv_row + + (qkv_head0 + g) * kHeadDim + 4 * v; + int const head_rows = kRole == kQ ? nq : kRole == kIQ ? niq : 0; + int64_t const out_lane = + (static_cast(t0) * head_rows + head0 + g) * kHeadDim + 4 * v; + int64_t const out_step = static_cast(stride) * head_rows * kHeadDim; + int64_t out = out_lane; + + for (int it = 0; it < n; it++) { + issue_rows(it + kStages - 1); + issue_cos_sin(it + kStages - 1); + cpAsyncWait(); + __syncwarp(); + uint8_t const* const stage = ring + (it % kStages) * stage_bytes; + float cos4[4], sin4[4]; + if constexpr (kNorm) { + scalar_t const* const cs = reinterpret_cast( + stage + kSteps * kStepRows * kRowPitch); + if constexpr (!kOneRow) { + loadElems(cs + 4 * v, cos4); + loadElems(cs + kRotaryDim / 2 + 4 * v, sin4); + } + } + // Slot of the token: block and offset with one 32-bit division (the host + // checks that every slot fits). + int sm = -1; + uint32_t blk = 0, tok = 0; + if constexpr (kSlots) { + if (it - slot_base == 32) { + slot_base = it; + slot_l = fetch(slot_mapping, slot_base, -1); + } + sm = static_cast(__shfl_sync( + FINAL_MASK, static_cast(slot_l), it - slot_base, 32)); + blk = static_cast(sm) / static_cast(block_size); + tok = static_cast(sm) - blk * block_size; + } + + #pragma unroll + for (int step = 0; step < kSteps; step++) { + int const r = step * kStepRows + g; // row of the warp + int const head = head0 + r; + bool const row_ok = r < nrows; + int const so = step * kStepRows * kHeadDim; // step offset, elements + float e[kChunks][kElemsPerLane]; + if constexpr (kOneRow) { + loadElems( + reinterpret_cast(stage + r * kRowPitch) + 4 * v, + e[0]); + if constexpr (kNorm) { + normAndRope( + e[0], lane, eps, norm_w, /*do_rope=*/true, kRotaryDim, + reinterpret_cast(stage + kSteps * kRowPitch), + /*apply_norm=*/true); + } + } else { + Row x; + loadRow(stage + r * kRowPitch, v, x); + if constexpr (kNorm) normRope(x, w1, eps, cos4, sin4); + #pragma unroll + for (int c = 0; c < kChunks; c++) chunkElems(x, c, e[c]); + } + if constexpr (kRole == kQ) { + if (row_ok) { + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (kOut == kQFp8Unit) { + storeUnitScaleQElemsFp8(q_fp8_out + out + so + 32 * c, + e[c]); + } else if constexpr (kOut == kQFp8) { + storeScaledQElemsFp8(q_fp8_out + out + so + 32 * c, + e[c], q_fp8_inv_scale); + } else { + if (q_out != nullptr) { + storeElems(q_out + out + so + 32 * c, e[c]); + } else if (q_fp8_out == nullptr && write_qkv) { + storeElems(in_place + so + 32 * c, e[c]); + } + if (q_fp8_out != nullptr) { + storeScaledQElemsFp8(q_fp8_out + out + so + 32 * c, + e[c], q_fp8_inv_scale); + } + } + } + } + } else if constexpr (kRole == kIQ) { + if (row_ok) { + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (kOut == kIQOut) { + if constexpr (kFp8Idx) { + storeIndexQElemsFp8( + reinterpret_cast(index_q_out + out + so + 32 * c), + e[c]); + } else { + storeElems(index_q_out + out + so + 32 * c, e[c]); + } + } else if constexpr (kOut == kIQInPlace) { + storeElems(in_place + so + 32 * c, e[c]); + } + } + } + } else if constexpr (kRole == kK || kRole == kIK) { + if constexpr (kOut == kInPlace) { + if (row_ok) { + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + storeElems(in_place + so + 32 * c, e[c]); + } + } + } + } + + if constexpr (kInsertKV && kRole == kIK) { + if constexpr (kNvfp4Idx) { + uint8_t* const page = + reinterpret_cast(index_cache) + blk * idx_s_block; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + uint8_t sf; + uint16_t const packed = + nvfp4QuantizeElems(e[c], 1.0f, sf); + if (sm >= 0 && row_ok) { + nvfp4StoreElems(page, /*kv=*/0, /*head=*/0, tok, block_size, + kOneRow ? lane : 8 * c + v, packed, sf); + } + } + } else if (sm >= 0 && row_ok) { + int64_t const off = static_cast(sm) * kHeadDim + 4 * v; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (kFp8Idx) { + storeElemsFp8( + reinterpret_cast(index_cache) + off + 32 * c, e[c]); + } else { + storeElems(index_cache + off + 32 * c, e[c]); + } + } + } + } else if constexpr (kInsertKV && (kRole == kK || kRole == kV)) { + constexpr int kv = kRole == kK ? 0 : 1; + if constexpr (kNvfp4 && !kOneRow) { + // One 16-value scale group per lane, regrouped through the ring + // row: lane (g, v) takes dims 16v .. 16v + 15 of row g. + uint8_t* const row = + ring + (it % kStages) * stage_bytes + r * kRowPitch; + if constexpr (kRole == kK) { + // The k in the model dtype, as the helper would quantize it (V + // rows already hold their model-dtype values). + #pragma unroll + for (int c = 0; c < kChunks; c++) { + storeElems( + reinterpret_cast(row + 64 * c + 8 * v), e[c]); + } + } + __syncwarp(); + float x[16]; + #pragma unroll + for (int k = 0; k < 4; k++) { + float q[kElemsPerLane]; + loadElems( + reinterpret_cast(row + 32 * v + 8 * k), q); + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) x[4 * k + i] = q[i]; + } + uint8_t sf; + uint2 const codes = nvfp4QuantizeGroup(x, kv_scale, sf); + if (sm >= 0 && row_ok) { + nvfp4StoreGroup( + reinterpret_cast(kv_cache) + blk * kv_s_block, kv, + head, tok, block_size, v, codes, sf); + } + } else if constexpr (kNvfp4) { + uint8_t* const page = + reinterpret_cast(kv_cache) + blk * kv_s_block; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + uint8_t sf; + uint16_t const packed = + nvfp4QuantizeElems(e[c], kv_scale, sf); + if (sm >= 0 && row_ok) { + nvfp4StoreElems(page, kv, head, tok, block_size, + kOneRow ? lane : 8 * c + v, packed, sf); + } + } + } else if (sm >= 0 && row_ok) { + cache_t* const dst = kv_cache + blk * kv_s_block + head * kv_s_head + + tok * kv_s_token + + (kv * kHeadDim + 4 * v) * kv_s_dim; + #pragma unroll + for (int c = 0; c < kLaneChunks; c++) { + if constexpr (!kHasIndex) { + using Converter = vllm::_typeConvert; + #pragma unroll + for (int i = 0; i < kElemsPerLane; i++) { + e[c][i] = Converter::convert(Converter::convert(e[c][i])); + } + } + storeCacheElems(dst + 32 * c * kv_s_dim, + e[c]); + } + } + } + } + in_place += src_step; + out += out_step; + __syncwarp(); + } +} + +template +__global__ void __launch_bounds__(kMaxWarps * 32) fusedMiniMaxM3RowGroupKernel( + scalar_t* __restrict__ qkv, scalar_t* __restrict__ q_out, + uint8_t* __restrict__ q_fp8_out, out_idx_t* __restrict__ index_q_out, + scalar_t const* __restrict__ q_norm_w, + scalar_t const* __restrict__ k_norm_w, + scalar_t const* __restrict__ iq_norm_w, + scalar_t const* __restrict__ ik_norm_w, + scalar_t const* __restrict__ cos_sin_cache, + int64_t const* __restrict__ positions, + int64_t const* __restrict__ slot_mapping, + int64_t const* __restrict__ index_slot_mapping, + cache_t* __restrict__ kv_cache, out_idx_t* __restrict__ index_cache, + float const eps, float const q_fp8_inv_scale, int const rotary_dim, + int const num_tokens, int const nq, int const nkv, int const niq, + int const block_size, int64_t const kv_s_block, int64_t const kv_s_head, + int64_t const kv_s_token, int64_t const kv_s_dim, int64_t const idx_s_block, + float const* __restrict__ kv_k_scale, float const* __restrict__ kv_v_scale, + bool const write_qkv) { + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + extern __shared__ __align__(16) uint8_t smem[]; + Roles const roles(kQHeads, nq, nkv, niq, kInsertKV, kProcessIndex); + int const qkv_row = (nq + 2 * nkv + (kHasIndex ? niq + 1 : 0)) * kHeadDim; + // CTA (x, y) runs warps [y * cta_warps, (y + 1) * cta_warps) of the warp + // groups of tokens x, x + gridDim.x, ... + int const cta_warps = blockDim.x / 32; + int w = blockIdx.y * cta_warps + threadIdx.x / 32; + if (w >= roles.warps) return; + constexpr int kQStage = stageBytes(); + constexpr int kSideStage = stageBytes(); + int const q_ring = kStages * kQStage; + int const side_ring = kStages * kSideStage; + uint8_t* const ring_w = + smem + roles.ringOffset(w, q_ring, side_ring) - + roles.ringOffset(blockIdx.y * cta_warps, q_ring, side_ring); + #define RG_ARGS(NORM_W, SLOTS, CACHE_SCALE) \ + qkv, q_out, q_fp8_out, index_q_out, NORM_W, cos_sin_cache, positions, \ + SLOTS, kv_cache, index_cache, eps, q_fp8_inv_scale, num_tokens, nq, \ + niq, qkv_row, block_size, kv_s_block, kv_s_head, kv_s_token, \ + kv_s_dim, idx_s_block, CACHE_SCALE, write_qkv + #define RG_RUN(ROLE, STEPS, OUT) \ + runWarp + #define RG_RUN1(ROLE, OUT) \ + runWarp + if (w < roles.q_warps) { + int const h0 = w * kQHeads; + constexpr int kSteps = kQHeads / kRows; + uint8_t* const ring = ring_w; + int const rows = min(kQHeads, nq - h0); + if (q_out == nullptr && q_fp8_out != nullptr) { + if (q_fp8_inv_scale == 1.0f) { + RG_RUN(kQ, kSteps, kQFp8Unit) + (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f)); + } else { + RG_RUN(kQ, kSteps, kQFp8) + (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f)); + } + } else { + RG_RUN(kQ, kSteps, kQAny) + (h0, rows, h0, ring, kQStage, RG_ARGS(q_norm_w, nullptr, 1.0f)); + } + } else { + uint8_t* const ring = ring_w; + w -= roles.q_warps; + if (w < roles.kv_warps) { + int const h0 = roles.kv_one ? w : w * kRows; + float const scale = kNvfp4 ? *kv_k_scale : 1.0f; + int const rows = roles.kv_one ? 1 : min(kRows, nkv - h0); + if (rows == 1) { + if (write_qkv) { + RG_RUN1(kK, kInPlace) + (h0, 1, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } else { + RG_RUN1(kK, kNone) + (h0, 1, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } + } else if (write_qkv) { + RG_RUN(kK, 1, kInPlace) + (h0, rows, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } else { + RG_RUN(kK, 1, kNone) + (h0, rows, nq + h0, ring, kSideStage, + RG_ARGS(k_norm_w, slot_mapping, scale)); + } + } else if ((w -= roles.kv_warps) < roles.v_warps) { + if constexpr (kInsertKV) { + int const h0 = roles.kv_one ? w : w * kRows; + int const rows = roles.kv_one ? 1 : min(kRows, nkv - h0); + float const scale = kNvfp4 ? *kv_v_scale : 1.0f; + if (rows == 1) { + RG_RUN1(kV, kNone) + (h0, 1, nq + nkv + h0, ring, kSideStage, + RG_ARGS(nullptr, slot_mapping, scale)); + } else { + RG_RUN(kV, 1, kNone) + (h0, rows, nq + nkv + h0, ring, kSideStage, + RG_ARGS(nullptr, slot_mapping, scale)); + } + } + } else if ((w -= roles.v_warps) < roles.iq_warps) { + if constexpr (kProcessIndex) { + int const h0 = roles.iq_one ? w : w * kRows; + int const rows = roles.iq_one ? 1 : min(kRows, niq - h0); + int const qh0 = nq + 2 * nkv + h0; + if (rows == 1) { + if (index_q_out != nullptr) { + RG_RUN1(kIQ, kIQOut) + (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else if (write_qkv) { + RG_RUN1(kIQ, kIQInPlace) + (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else { + RG_RUN1(kIQ, kNone) + (h0, 1, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } + } else if (index_q_out != nullptr) { + RG_RUN(kIQ, 1, kIQOut) + (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else if (write_qkv) { + RG_RUN(kIQ, 1, kIQInPlace) + (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } else { + RG_RUN(kIQ, 1, kNone) + (h0, rows, qh0, ring, kSideStage, RG_ARGS(iq_norm_w, nullptr, 1.0f)); + } + } + } else if constexpr (kProcessIndex) { + if (write_qkv) { + RG_RUN1(kIK, kInPlace) + (0, 1, nq + 2 * nkv + niq, ring, kSideStage, + RG_ARGS(ik_norm_w, index_slot_mapping, 1.0f)); + } else { + RG_RUN1(kIK, kNone) + (0, 1, nq + 2 * nkv + niq, ring, kSideStage, + RG_ARGS(ik_norm_w, index_slot_mapping, 1.0f)); + } + } + } + #undef RG_RUN + #undef RG_RUN1 + #undef RG_ARGS + #if __CUDA_ARCH__ >= 900 + cudaTriggerProgrammaticLaunchCompletion(); + #endif + #endif +} + +// Shared memory of the largest CTA of cta_warps warps. +template +__host__ inline int smemBytes(Roles const& r, int const cta_warps) { + int const q_ring = kStages * stageBytes(); + int const side_ring = kStages * stageBytes(); + int bytes = 0; + for (int w0 = 0; w0 < r.warps; w0 += cta_warps) { + int const w1 = std::min(w0 + cta_warps, r.warps); + bytes = std::max(bytes, r.ringOffset(w1, q_ring, side_ring) - + r.ringOffset(w0, q_ring, side_ring)); + } + return bytes; +} + +// Launch with the grid this instantiation's occupancy allows: the CTAs walk +// the tokens with stride gridDim.x. +template +void launch(cudaLaunchConfig_t config, int const num_tokens, Args... args) { + auto* const kernel = + fusedMiniMaxM3RowGroupKernel; + // Per CTA shape (warps, shared memory), computed once: the resident CTAs + // per SM. + static bool const opted_in = + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + get_device_prop()->sharedMemPerBlockOptin) == + cudaSuccess; + static int smem_of[kMaxWarps + 1] = {}; + static int resident_of[kMaxWarps + 1] = {}; + int const warps = config.blockDim.x / 32; + int const smem = static_cast(config.dynamicSmemBytes); + if (smem_of[warps] != smem) { + (void)opted_in; + int n = 0; + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&n, kernel, warps * 32, smem); + resident_of[warps] = n > 0 ? n : 1; + smem_of[warps] = smem; + } + int const ctas_per_token = config.gridDim.y; + config.gridDim.x = std::max( + 1, std::min(num_tokens, get_device_prop()->multiProcessorCount * + resident_of[warps] / ctas_per_token)); + cudaLaunchKernelEx(&config, kernel, args...); +} + +} // namespace rg +#endif // USE_ROCM + +#ifndef USE_ROCM +// q heads per q warp: four (one step, shorter warps) for a launch of at most +// 128 q rows, or while a token has at most 32 q heads and there are few q rows +// per SM; eight otherwise (fewer, longer warps keep more CTAs resident). +inline int rgQHeads(int const num_tokens, int const nq) { + static int const num_sms = get_device_prop()->multiProcessorCount; + int64_t const q_rows = static_cast(num_tokens) * nq; + return q_rows <= 128 || (nq <= 32 && q_rows < 64 * num_sms) ? 4 : 8; +} +// Warps per CTA: a token's whole warp group once there are enough tokens to +// give every SM CTAs; before that, eight, which spreads a token over SMs. +inline int rgCtaWarps(int const num_tokens) { + static int const num_sms = get_device_prop()->multiProcessorCount; + return 2 * num_tokens >= 3 * num_sms ? rg::kMaxWarps : 8; +} +#endif + // ──────────────────────────────────────────────────────────────────────────── // Launch wrapper // ──────────────────────────────────────────────────────────────────────────── @@ -809,9 +1683,10 @@ void launchFusedMiniMaxM3( float const q_fp8_inv_scale, int const rotary_dim, int const num_tokens, int const nq, int const nkv, int const niq, int const block_size, int64_t const kv_s_block, int64_t const kv_s_head, int64_t const kv_s_token, - int64_t const kv_s_dim, bool const has_index, bool const insert_kv, - bool const process_index, bool const fp8_idx, float const* kv_k_scale, - float const* kv_v_scale, cudaStream_t stream) { + int64_t const kv_s_dim, int64_t const idx_s_block, bool const has_index, + bool const insert_kv, bool const process_index, bool const fp8_idx, + bool const nvfp4_idx, float const* kv_k_scale, float const* kv_v_scale, + bool const write_qkv, bool const slots_fit_int32, cudaStream_t stream) { // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the // void* pointers per instantiation in the LAUNCH macro. // Slot count must match the kernel's compile-time gating. @@ -843,73 +1718,437 @@ void launchFusedMiniMaxM3( config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \ - cudaLaunchKernelEx( \ - &config, \ - fusedMiniMaxM3QNormRopeKVInsertKernel, \ - qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), \ - q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ + // The row-group kernel: same bytes, for rotary_dim == 64 on SM80+ (cp.async). + int const q_heads = rgQHeads(num_tokens, nq); + rg::Roles const roles(q_heads, nq, nkv, niq, insert_kv, process_index); + auto const rg_smem = [&](int const warps) { + return q_heads == 8 ? rg::smemBytes<8>(roles, warps) + : rg::smemBytes<4>(roles, warps); + }; + int cta_warps = + std::min({roles.warps, rgCtaWarps(num_tokens), rg::kMaxWarps}); + // Eight warps take at most 63 KB; a whole group can exceed the 99 KB per + // block of SM 8.6 / 8.9 / 12.x. + static int const max_block_smem = get_device_prop()->sharedMemPerBlockOptin; + if (rg_smem(cta_warps) > max_block_smem) cta_warps = std::min(roles.warps, 8); + auto const aligned = [](void const* p, uintptr_t const n) { + return reinterpret_cast(p) % n == 0; + }; + bool const row_groups = + sm_version >= 80 && rotary_dim == rg::kRotaryDim && slots_fit_int32 && + aligned(qkv, 16) && aligned(cos_sin_cache, 16) && aligned(q_norm_w, 8) && + aligned(k_norm_w, 8) && aligned(iq_norm_w, 8) && aligned(ik_norm_w, 8); + cudaLaunchConfig_t rg_config = config; + rg_config.blockDim = dim3(cta_warps * 32); + rg_config.gridDim = dim3(1, (roles.warps + cta_warps - 1) / cta_warps); + rg_config.dynamicSmemBytes = rg_smem(cta_warps); + + #define LAUNCH_ARGS(OUT_T) \ + qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), q_norm_w, \ + k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ slot_mapping, index_slot_mapping, kv_cache, \ reinterpret_cast(index_cache), eps, q_fp8_inv_scale, \ rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ - kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale) + kv_s_head, kv_s_token, kv_s_dim, idx_s_block, kv_k_scale, kv_v_scale + #define LAUNCH_RG(QH, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, \ + OUT_T) \ + rg::launch(rg_config, num_tokens, \ + LAUNCH_ARGS(OUT_T), write_qkv) + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T) \ + do { \ + if (!row_groups) { \ + cudaLaunchKernelEx( \ + &config, \ + fusedMiniMaxM3QNormRopeKVInsertKernel< \ + scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \ + PROCESS_INDEX, FP8, NVFP4_IDX>, \ + LAUNCH_ARGS(OUT_T)); \ + } else if (q_heads == 8) { \ + LAUNCH_RG(8, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T); \ + } else { \ + LAUNCH_RG(4, HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T); \ + } \ + } while (0) #else // ROCm: standard kernel launch syntax (no PDL/stream serialization). // clang-format off - #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, OUT_T) \ + #define LAUNCH(HAS_INDEX, INSERT, PROCESS_INDEX, FP8, NVFP4_IDX, OUT_T) \ fusedMiniMaxM3QNormRopeKVInsertKernel< \ scalar_t, cache_t, kv_dt, kNvfp4, OUT_T, HAS_INDEX, INSERT, \ - PROCESS_INDEX, FP8><<>>( \ + PROCESS_INDEX, FP8, NVFP4_IDX><<>>( \ qkv, q_out, q_fp8_out, reinterpret_cast(index_q_out), \ q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \ slot_mapping, index_slot_mapping, kv_cache, \ reinterpret_cast(index_cache), eps, q_fp8_inv_scale, \ rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, \ - kv_s_head, kv_s_token, kv_s_dim, kv_k_scale, kv_v_scale) + kv_s_head, kv_s_token, kv_s_dim, idx_s_block, kv_k_scale, \ + kv_v_scale) // clang-format on #endif if constexpr (kNvfp4) { // NVFP4 main cache: sparse serving only (the host checks insert_kv). if (!process_index) { - LAUNCH(true, true, false, false, scalar_t); + LAUNCH(true, true, false, false, false, scalar_t); + } else if (nvfp4_idx) { + LAUNCH(true, true, true, true, true, uint8_t); } else if (fp8_idx) { - LAUNCH(true, true, true, true, uint8_t); + LAUNCH(true, true, true, true, false, uint8_t); } else { - LAUNCH(true, true, true, false, scalar_t); + LAUNCH(true, true, true, false, false, scalar_t); } } else if (has_index) { if (!process_index) { if (insert_kv) { - LAUNCH(true, true, false, false, scalar_t); + LAUNCH(true, true, false, false, false, scalar_t); } else { - LAUNCH(true, false, false, false, scalar_t); + LAUNCH(true, false, false, false, false, scalar_t); } } else if (insert_kv) { - if (fp8_idx) { - LAUNCH(true, true, true, true, + if (nvfp4_idx) { + LAUNCH(true, true, true, true, true, + uint8_t); // sparse serving, NVFP4 index-K, fp8 index_q + } else if (fp8_idx) { + LAUNCH(true, true, true, true, false, uint8_t); // sparse serving, fp8 index outputs } else { - LAUNCH(true, true, true, false, scalar_t); // sparse serving, bf16 + LAUNCH(true, true, true, false, false, + scalar_t); // sparse serving, bf16 } } else { if (fp8_idx) { - LAUNCH(true, false, true, true, + LAUNCH(true, false, true, true, false, uint8_t); // sparse profiling, fp8 index_q } else { - LAUNCH(true, false, true, false, scalar_t); // sparse profiling, bf16 + LAUNCH(true, false, true, false, false, + scalar_t); // sparse profiling, bf16 } } + } else if (insert_kv) { + LAUNCH(false, true, false, false, false, scalar_t); // dense, KV insert } else { - // Dense layer: never has an index branch and never inserts here (the - // generic Attention layer owns the KV insert). - LAUNCH(false, false, false, false, scalar_t); + LAUNCH(false, false, false, false, false, scalar_t); // dense } #undef LAUNCH +#undef LAUNCH_RG +#undef LAUNCH_ARGS +} + +#ifndef USE_ROCM +// ──────────────────────────────────────────────────────────────────────────── +// DSpark draft context K/V: every draft layer in one launch (bf16 in). +// ──────────────────────────────────────────────────────────────────────────── +// The draft projects the target's context states to the K/V of all its layers +// with one GEMM, ``all_kv [N, L, 2, nkv, 128]``. blockIdx.z picks the layer, +// blockIdx.y the head, blockIdx.x a run of tokens. Eight lanes own one token; +// lane l holds two 8-dim chunks of K and of V, A = dims [8l, 8l + 8) and +// B = dims [64 + 8l, 64 + 8l + 8), so with rotary_dim == 64 every lane rotates +// its A chunk (partner lane l ^ 4) and no lane idles through the RoPE. K gets +// the layer's Gemma k-norm (rounded to bf16 like the unfused norm) and partial +// NeoX RoPE (fp32 math rounded once, like rotary_embedding); K and V go into +// the token's cache row like reshape_and_cache_flash (fp8: e4m3(bf16(x) * +// (1 / scale)); a unit scale converts the bf16 pairs directly). +// Layer caches are separate allocations of one layout, logical shape +// [nb, nkv, bs, 2*128] with K at [..., :128] and V at [..., 128:]; they are +// addressed as layer 0's cache plus a per-layer element offset (likewise the +// fp8 k/v scales from layer 0's k scale), so the launch is CUDA-graph safe. +constexpr int kCtxLanesPerRow = 8; +constexpr int kCtxElems = 8; // bf16 per 16-byte chunk +constexpr int kCtxHalfVecs = kHeadDim / 2 / kCtxElems; // 8: chunk B offset +constexpr int kCtxBlockSize = 256; +constexpr int kCtxRowsPerBlock = kCtxBlockSize / kCtxLanesPerRow; // 32 + +__device__ __forceinline__ float bf16lo(uint32_t v) { + return __uint_as_float(v << 16); +} +__device__ __forceinline__ float bf16hi(uint32_t v) { + return __uint_as_float(v & 0xffff0000u); +} +__device__ __forceinline__ uint32_t packBf16(float lo, float hi) { + __nv_bfloat162 const r = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&r); +} +// Loads of what the PDL predecessor may still be writing: coherent (an .nc +// load may be hoisted above griddepcontrol.wait). +__device__ __forceinline__ uint4 ldgV4(uint4 const* p) { + uint4 r; + asm volatile("ld.global.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} +// Low 32 bits of an int64 (positions and slots fit). +__device__ __forceinline__ int ldgLo32(int64_t const* p) { + int r; + asm volatile("ld.global.b32 %0, [%1];" : "=r"(r) : "l"(p)); + return r; +} +// Same approximation as flashinfer's gemma_rmsnorm. +__device__ __forceinline__ float rsqrtApproxFtz(float x) { + float r; + asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} +// e4m3 of four floats, low byte first. +__device__ __forceinline__ uint32_t packE4m3(float a, float b, float c, + float d) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890 + uint32_t r; + asm("{\n" + " .reg .b16 lo, hi;\n" + " cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" + " cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" + " mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(r) + : "f"(a), "f"(b), "f"(c), "f"(d)); + return r; + #else + uint32_t const lo = + __nv_cvt_float2_to_fp8x2(make_float2(a, b), __NV_SATFINITE, __NV_E4M3); + uint32_t const hi = + __nv_cvt_float2_to_fp8x2(make_float2(c, d), __NV_SATFINITE, __NV_E4M3); + return lo | (hi << 16); + #endif +} +// Eight bf16 (4 packed pairs) * inv_scale -> e4m3, low byte first. +__device__ __forceinline__ uint2 bf16x8ToE4m3(uint4 const p, + float const inv_scale) { + uint32_t const w[4] = {p.x, p.y, p.z, p.w}; + uint32_t r[2]; + #pragma unroll + for (int i = 0; i < 2; i++) { + r[i] = packE4m3(bf16lo(w[2 * i]) * inv_scale, bf16hi(w[2 * i]) * inv_scale, + bf16lo(w[2 * i + 1]) * inv_scale, + bf16hi(w[2 * i + 1]) * inv_scale); + } + return make_uint2(r[0], r[1]); +} + #if ((defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && \ + __CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000) || \ + (defined(__CUDA_ARCH_SPECIFIC__) && __CUDA_ARCH_SPECIFIC__ >= 1000)) && \ + defined(CUDART_VERSION) && CUDART_VERSION >= 13010 + #define VLLM_CTX_KV_BF16X2_TO_E4M3X2 + #endif +// Eight bf16 -> e4m3 with a unit scale (exact); one convert per pair where +// the arch and toolkit (PTX ISA 9.1+) convert bf16x2 -> e4m3x2. +__device__ __forceinline__ uint2 bf16x8ToE4m3(uint4 const p) { + #ifdef VLLM_CTX_KV_BF16X2_TO_E4M3X2 + uint32_t r[2]; + uint32_t const w[4] = {p.x, p.y, p.z, p.w}; + #pragma unroll + for (int i = 0; i < 2; i++) { + asm("{\n" + " .reg .b16 lo, hi;\n" + " cvt.rn.satfinite.e4m3x2.bf16x2 lo, %1;\n" + " cvt.rn.satfinite.e4m3x2.bf16x2 hi, %2;\n" + " mov.b32 %0, {lo, hi};\n" + "}" + : "=r"(r[i]) + : "r"(w[2 * i]), "r"(w[2 * i + 1])); + } + return make_uint2(r[0], r[1]); + #else + return bf16x8ToE4m3(p, 1.0f); + #endif } +template +__global__ void __launch_bounds__(kCtxBlockSize) + fusedMiniMaxM3DSparkContextKVKernel( + uint4 const* __restrict__ all_kv, // bf16 [N, L, 2, nkv, 128] + uint4 const* __restrict__ k_norm_w, // bf16 [L, 128] + uint4 const* __restrict__ cos_sin_cache, // bf16 [max_pos, rot] + int64_t const* __restrict__ positions, // [N] + int64_t const* __restrict__ slot_mapping, // [N], < 0: no write + uint8_t* __restrict__ kv_cache, // layer 0's cache, bytes + int64_t const* __restrict__ cache_offsets, // [L] elements from it + float const* __restrict__ scales, // layer 0's k scale + int64_t const* __restrict__ scale_offsets, // [2, L] k, v; floats + float const eps, int const rotary_dim, int const num_tokens, + int const num_layers, int const block_shift, int const s_token, + int const s_layer, int const s_side, int const s_head, + int const kv_s_block, int const kv_s_head, int const kv_s_token) { + // Strides are in 16-byte (8 x bf16) units for all_kv / cos_sin, elements + // for the cache. + constexpr int kElemBytes = kFp8 ? 1 : 2; + int const row0 = blockIdx.x * kCtxRowsPerBlock; + // Warps with no token exit; the others run their tail rows too (the + // shuffles need every lane) and skip the store. + if (row0 + threadIdx.x / 32 * (32 / kCtxLanesPerRow) >= num_tokens) return; + int const head = blockIdx.y; + int const layer = blockIdx.z; + int const lane = threadIdx.x % kCtxLanesPerRow; + int const dim0 = lane * kCtxElems; // chunk A; chunk B is 64 dims later + int const t = row0 + threadIdx.x / kCtxLanesPerRow; + + int const half = rotary_dim / 2; + bool const in_rope = dim0 < rotary_dim; + bool const first_half = dim0 < half; + int const cs_vec = + in_rope ? (first_half ? dim0 : dim0 - half) / kCtxElems : 0; + int const lane_xor = half / kCtxElems; + int const rot_vecs = rotary_dim / kCtxElems; + // x * cos - partner * sin in the first half, + in the second. + uint32_t const sin_sign = first_half ? 0x80008000u : 0u; + + uint8_t* const cache = + kv_cache + (cache_offsets[layer] + head * kv_s_head + dim0) * kElemBytes; + float k_scale = 1.0f, v_scale = 1.0f; + if constexpr (kFp8) { + int64_t const k_off = scale_offsets[layer]; + int64_t const v_off = scale_offsets[num_layers + layer]; + k_scale = scales[k_off]; + v_scale = scales[v_off]; + } + float g[2][kCtxElems]; // 1 + w, chunks A and B + #pragma unroll + for (int c = 0; c < 2; c++) { + uint4 const w = + k_norm_w[layer * (kHeadDim / kCtxElems) + c * kCtxHalfVecs + lane]; + uint32_t const wp[4] = {w.x, w.y, w.z, w.w}; + #pragma unroll + for (int i = 0; i < 4; i++) { + g[c][2 * i] = 1.0f + bf16lo(wp[i]); + g[c][2 * i + 1] = 1.0f + bf16hi(wp[i]); + } + } + + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); + #endif + + bool const valid = t < num_tokens; + int const ts = valid ? t : 0; + int const pos = ldgLo32(positions + ts); + int const slot = valid ? ldgLo32(slot_mapping + ts) : -1; + uint4 const* const src = + all_kv + ts * s_token + layer * s_layer + head * s_head + lane; + uint4 kr[2], vr[2]; + #pragma unroll + for (int c = 0; c < 2; c++) { + kr[c] = ldgV4(src + c * kCtxHalfVecs); + vr[c] = ldgV4(src + s_side + c * kCtxHalfVecs); + } + // Lanes past the rotary dims load (and later discard) a valid row too. + uint4 const* const cs = cos_sin_cache + pos * rot_vecs + cs_vec; + uint4 const cr = cs[0]; + uint4 const sr = cs[lane_xor]; + + bool const store = slot >= 0; + // Power-of-two block size: shift and mask, not division. + uint8_t* const dst = + cache + (static_cast(slot >> block_shift) * kv_s_block + + (slot & ((1 << block_shift) - 1)) * kv_s_token) * + kElemBytes; + constexpr int kChunkB = kHeadDim / 2 * kElemBytes; + constexpr int kV = kHeadDim * kElemBytes; + float const k_inv = __fdividef(1.0f, k_scale); + float const v_inv = __fdividef(1.0f, v_scale); + #ifdef VLLM_CTX_KV_BF16X2_TO_E4M3X2 + bool const unit_scale = k_scale == 1.0f && v_scale == 1.0f; + #else + constexpr bool unit_scale = false; + #endif + // Eight bf16 -> the cache row at byte offset `off`. + auto put = [&](int const off, uint4 const x, float const inv_scale) { + if constexpr (kFp8) { + uint2 const y = unit_scale ? bf16x8ToE4m3(x) : bf16x8ToE4m3(x, inv_scale); + if (store) *reinterpret_cast(dst + off) = y; + } else { + if (store) *reinterpret_cast(dst + off) = x; + } + }; + + // V first, so its registers free up before the K math. + #pragma unroll + for (int c = 0; c < 2; c++) put(kV + c * kChunkB, vr[c], v_inv); + + // K: Gemma RMSNorm over the 8 lanes of the row. + float x[2][kCtxElems]; + float sumsq = 0.0f; + #pragma unroll + for (int c = 0; c < 2; c++) { + uint32_t const kp[4] = {kr[c].x, kr[c].y, kr[c].z, kr[c].w}; + #pragma unroll + for (int j = 0; j < 4; j++) { + x[c][2 * j] = bf16lo(kp[j]); + x[c][2 * j + 1] = bf16hi(kp[j]); + sumsq = fmaf(x[c][2 * j], x[c][2 * j], sumsq); + sumsq = fmaf(x[c][2 * j + 1], x[c][2 * j + 1], sumsq); + } + } + #pragma unroll + for (int mask = kCtxLanesPerRow / 2; mask > 0; mask >>= 1) { + sumsq += __shfl_xor_sync(FINAL_MASK, sumsq, mask, 32); + } + float const rr = rsqrtApproxFtz(sumsq * (1.0f / kHeadDim) + eps); + // bf16-rounded norm output, packed. + uint32_t kn[2][4]; + #pragma unroll + for (int c = 0; c < 2; c++) { + #pragma unroll + for (int j = 0; j < 4; j++) { + kn[c][j] = packBf16(x[c][2 * j] * rr * g[c][2 * j], + x[c][2 * j + 1] * rr * g[c][2 * j + 1]); + } + } + // RoPE on chunk A; the partner chunk is half / 8 lanes away. + uint32_t const cp[4] = {cr.x, cr.y, cr.z, cr.w}; + uint32_t const sp[4] = {sr.x ^ sin_sign, sr.y ^ sin_sign, sr.z ^ sin_sign, + sr.w ^ sin_sign}; + #pragma unroll + for (int j = 0; j < 4; j++) { + uint32_t const pn = __shfl_xor_sync(FINAL_MASK, kn[0][j], lane_xor, 32); + uint32_t const rot = + packBf16(__fmaf_rn(bf16lo(kn[0][j]), bf16lo(cp[j]), + __fmul_rn(bf16lo(pn), bf16lo(sp[j]))), + __fmaf_rn(bf16hi(kn[0][j]), bf16hi(cp[j]), + __fmul_rn(bf16hi(pn), bf16hi(sp[j])))); + kn[0][j] = in_rope ? rot : kn[0][j]; + } + #pragma unroll + for (int c = 0; c < 2; c++) { + put(c * kChunkB, make_uint4(kn[c][0], kn[c][1], kn[c][2], kn[c][3]), k_inv); + } + + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); + #endif +} + +template +void launchFusedMiniMaxM3DSparkContextKV( + uint4 const* all_kv, uint4 const* k_norm_w, uint4 const* cos_sin_cache, + int64_t const* positions, int64_t const* slot_mapping, uint8_t* kv_cache, + int64_t const* cache_offsets, float const* scales, + int64_t const* scale_offsets, float const eps, int const rotary_dim, + int const num_tokens, int const num_layers, int const nkv, + int const block_shift, int const s_token, int const s_layer, + int const s_side, int const s_head, int const kv_s_block, + int const kv_s_head, int const kv_s_token, cudaStream_t stream) { + if (num_tokens == 0) return; + static int const sm_version = getSMVersion(); + cudaLaunchConfig_t config; + config.gridDim = dim3((num_tokens + kCtxRowsPerBlock - 1) / kCtxRowsPerBlock, + nkv, num_layers); + config.blockDim = dim3(kCtxBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attrs; + config.numAttrs = (sm_version >= 90) ? 1 : 0; + cudaLaunchKernelEx(&config, fusedMiniMaxM3DSparkContextKVKernel, all_kv, + k_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache, + cache_offsets, scales, scale_offsets, eps, rotary_dim, + num_tokens, num_layers, block_shift, s_token, s_layer, + s_side, s_head, kv_s_block, kv_s_head, kv_s_token); +} +#endif // !USE_ROCM + } // namespace minimax_m3_fused_ops } // namespace vllm @@ -948,8 +2187,9 @@ void launchFusedMiniMaxM3( static_cast(eps), 1.0f / static_cast(q_fp8_scale), \ static_cast(rotary_dim), num_tokens, nq, nkv, niq, \ static_cast(block_size), kv_s_block, kv_s_head, kv_s_token, \ - kv_s_dim, has_index, insert_kv, process_index, fp8_idx, kv_k_scale_ptr, \ - kv_v_scale_ptr, stream) + kv_s_dim, idx_s_block, has_index, insert_kv, process_index, fp8_idx, \ + nvfp4_idx, kv_k_scale_ptr, kv_v_scale_ptr, write_qkv, slots_fit_int32, \ + stream) #define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE) \ CALL_FUSED_MINIMAX_M3_IMPL(CACHE_T, KV_DTYPE, false) // clang-format on @@ -980,7 +2220,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert( q_fp8_out, // [N, nq*128] contiguous E4M3 double q_fp8_scale, std::optional kv_k_scale, // [] f32, NVFP4 only - std::optional kv_v_scale) { + std::optional kv_v_scale, bool write_qkv) { STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(), "qkv must be contiguous CUDA"); STD_TORCH_CHECK( @@ -1032,11 +2272,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert( " + num_index_heads + 1) * 128 for sparse, " "(num_heads + 2*num_kv_heads) * 128 for dense"); - // Only the sparse layer inserts here (dense lets the generic Attention layer - // own the KV write); there is no dense+insert kernel instantiation. - STD_TORCH_CHECK( - !insert_kv || has_index, - "insert mode (kv_cache) requires the index branch (sparse layer)"); + STD_TORCH_CHECK(!nvfp4_kv || has_index, + "nvfp4 kv_cache requires the index branch (sparse layer)"); STD_TORCH_CHECK(has_index || !skip_index_branch, "skip_index_branch requires sparse qkv rows"); if (process_index) { @@ -1084,14 +2321,18 @@ void fused_minimax_m3_qknorm_rope_kv_insert( kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte, "fp8 kv_cache must use uint8 storage"); } - // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3. + // Indexer index-K cache: independent dtype -- qkv dtype, fp8 e4m3, or + // uint8 packed NVFP4 pages. if (process_index) { STD_TORCH_CHECK( index_cache.has_value() && (index_cache->scalar_type() == qkv.scalar_type() || index_cache->scalar_type() == - torch::headeronly::ScalarType::Float8_e4m3fn), - "insert mode requires index_cache matching qkv dtype or fp8 e4m3"); + torch::headeronly::ScalarType::Float8_e4m3fn || + index_cache->scalar_type() == + torch::headeronly::ScalarType::Byte), + "insert mode requires index_cache matching qkv dtype, fp8 e4m3, " + "or uint8 (nvfp4)"); } STD_TORCH_CHECK(kv_cache->dim() == 4 && kv_cache->stride(3) == 1, "kv_cache must be [nb,nkv,bs,2*head_dim] with contiguous " @@ -1150,6 +2391,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert( STD_TORCH_CHECK(std::isfinite(q_fp8_scale) && q_fp8_scale > 0.0, "q_fp8_scale must be finite and positive"); } + STD_TORCH_CHECK(write_qkv || (!q_out.has_value() && q_fp8_out.has_value()), + "write_qkv=False writes q only to q_fp8_out: pass q_fp8_out " + "and no q_out"); if (index_q_out.has_value()) { STD_TORCH_CHECK(process_index, "index_q_out requires index branch processing"); @@ -1166,16 +2410,36 @@ void fused_minimax_m3_qknorm_rope_kv_insert( } // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while - // q/k/v + q_out stay qkv dtype. Both index outputs must agree. + // q/k/v + q_out stay qkv dtype. Both index outputs must agree. The NVFP4 + // index path keeps e4m3 index-Q and stores index-K as packed NVFP4 pages. auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn; + bool const nvfp4_idx = + process_index && insert_kv && + index_cache->scalar_type() == torch::headeronly::ScalarType::Byte; bool const fp8_idx = process_index && - ((index_cache.has_value() && index_cache->scalar_type() == kFp8) || + (nvfp4_idx || + (index_cache.has_value() && index_cache->scalar_type() == kFp8) || (index_q_out.has_value() && index_q_out->scalar_type() == kFp8)); - if (fp8_idx) { + int64_t idx_s_block = 0; + if (nvfp4_idx) { +#ifdef USE_ROCM + STD_TORCH_CHECK(false, "nvfp4 index_cache is not supported on ROCm"); +#endif + int64_t const page_width = kHeadDim / 2 + kHeadDim / 16; STD_TORCH_CHECK( - !index_cache.has_value() || index_cache->scalar_type() == kFp8, - "fp8 index path: index_cache must be fp8 e4m3"); + index_cache->dim() == 3 && index_cache->size(1) == block_size && + index_cache->size(2) == page_width && index_cache->stride(2) == 1 && + index_cache->stride(1) == page_width && + index_cache->stride(0) >= block_size * page_width, + "nvfp4 index_cache must be [nb, block_size, 72] uint8 with " + "contiguous pages"); + idx_s_block = index_cache->stride(0); + } + if (fp8_idx) { + STD_TORCH_CHECK(!index_cache.has_value() || nvfp4_idx || + index_cache->scalar_type() == kFp8, + "fp8 index path: index_cache must be fp8 e4m3"); STD_TORCH_CHECK( !index_q_out.has_value() || index_q_out->scalar_type() == kFp8, "fp8 index path: index_q_out must be fp8 e4m3"); @@ -1188,6 +2452,12 @@ void fused_minimax_m3_qknorm_rope_kv_insert( float const* kv_v_scale_ptr = nvfp4_kv ? kv_v_scale->const_data_ptr() : nullptr; + // The row-group kernel indexes cache slots in 32 bits. + bool const slots_fit_int32 = + !insert_kv || + (kv_cache->size(0) * block_size <= INT32_MAX && + (!process_index || index_cache->size(0) * block_size <= INT32_MAX)); + const torch::stable::accelerator::DeviceGuard device_guard( qkv.get_device_index()); auto stream = get_current_cuda_stream(qkv.get_device_index()); @@ -1216,3 +2486,122 @@ void fused_minimax_m3_qknorm_rope_kv_insert( #undef CALL_FUSED_MINIMAX_M3 #undef CALL_FUSED_MINIMAX_M3_IMPL + +// ──────────────────────────────────────────────────────────────────────────── +// DSpark context K/V torch op +// ──────────────────────────────────────────────────────────────────────────── +void fused_minimax_m3_dspark_context_kv( + torch::stable::Tensor const& all_kv, // bf16 [N, L, 2, nkv, 128] + torch::stable::Tensor const& k_norm_weight, // bf16 [L, 128] + torch::stable::Tensor const& cos_sin_cache, // bf16 [max_pos, rotary_dim] + torch::stable::Tensor const& positions, // [N] i64 + torch::stable::Tensor const& slot_mapping, // [N] i64 + torch::stable::Tensor& kv_cache, // layer 0's [nb, nkv, bs, 2*128] + torch::stable::Tensor const& cache_offsets, // [L] i64, elements + torch::stable::Tensor const& k_scale, // layer 0's [] f32 k scale + torch::stable::Tensor const& scale_offsets, // [2, L] i64, floats + int64_t rotary_dim, double eps, const std::string& kv_cache_dtype) { +#ifdef USE_ROCM + STD_TORCH_CHECK(false, "fused_minimax_m3_dspark_context_kv is CUDA-only"); +#else + using torch::headeronly::ScalarType; + int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim; + STD_TORCH_CHECK(all_kv.is_cuda() && all_kv.dim() == 5 && + all_kv.size(2) == 2 && all_kv.size(4) == kHeadDim && + all_kv.stride(4) == 1 && + all_kv.scalar_type() == ScalarType::BFloat16, + "all_kv must be bf16 CUDA [N, L, 2, nkv, 128], last dim " + "contiguous"); + for (int d = 0; d < 4; d++) { + STD_TORCH_CHECK(all_kv.stride(d) % 8 == 0, + "all_kv strides must be multiples of 8 elements"); + } + STD_TORCH_CHECK(reinterpret_cast(all_kv.data_ptr()) % 16 == 0, + "all_kv must be 16-byte aligned"); + int const num_tokens = static_cast(all_kv.size(0)); + int const num_layers = static_cast(all_kv.size(1)); + int const nkv = static_cast(all_kv.size(3)); + STD_TORCH_CHECK(k_norm_weight.is_contiguous() && + k_norm_weight.scalar_type() == ScalarType::BFloat16 && + k_norm_weight.numel() == num_layers * kHeadDim, + "k_norm_weight must be contiguous bf16 [L, 128]"); + STD_TORCH_CHECK(cos_sin_cache.is_contiguous() && + cos_sin_cache.scalar_type() == ScalarType::BFloat16 && + cos_sin_cache.dim() == 2 && + cos_sin_cache.size(1) == rotary_dim, + "cos_sin_cache must be contiguous bf16 [max_pos, " + "rotary_dim]"); + STD_TORCH_CHECK(rotary_dim == 16 || rotary_dim == 32 || rotary_dim == 64, + "rotary_dim must be 16, 32 or 64"); + for (auto const* t : {&positions, &slot_mapping}) { + STD_TORCH_CHECK(t->is_cuda() && t->is_contiguous() && + t->scalar_type() == ScalarType::Long && + t->numel() == num_tokens, + "positions and slot_mapping must be contiguous int64 [N]"); + } + STD_TORCH_CHECK(cache_offsets.is_cuda() && cache_offsets.is_contiguous() && + cache_offsets.scalar_type() == ScalarType::Long && + cache_offsets.numel() == num_layers, + "cache_offsets must be int64 [L]"); + STD_TORCH_CHECK(scale_offsets.is_cuda() && scale_offsets.is_contiguous() && + scale_offsets.scalar_type() == ScalarType::Long && + scale_offsets.numel() == 2 * num_layers, + "scale_offsets must be int64 [2, L]"); + STD_TORCH_CHECK( + k_scale.is_cuda() && k_scale.scalar_type() == ScalarType::Float, + "k_scale must be a float32 CUDA tensor"); + bool const fp8 = kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3"; + STD_TORCH_CHECK(fp8 || kv_cache_dtype == "auto", + "kv_cache_dtype must be auto, fp8 or fp8_e4m3"); + STD_TORCH_CHECK(fp8 ? kv_cache.scalar_type() == ScalarType::Byte + : kv_cache.scalar_type() == ScalarType::BFloat16, + "kv_cache must be uint8 (fp8) or bf16 (auto)"); + STD_TORCH_CHECK( + kv_cache.dim() == 4 && kv_cache.size(1) == nkv && + kv_cache.size(3) == 2 * kHeadDim && kv_cache.stride(3) == 1 && + kv_cache.stride(0) % 16 == 0 && kv_cache.stride(1) % 16 == 0 && + kv_cache.stride(2) % 16 == 0, + "kv_cache must be [nb, nkv, bs, 256] with a contiguous last dim"); + int64_t const block_size = kv_cache.size(2); + STD_TORCH_CHECK(block_size > 0 && (block_size & (block_size - 1)) == 0, + "kv_cache block size must be a power of two"); + int block_shift = 0; + while ((int64_t{1} << block_shift) < block_size) block_shift++; + int64_t constexpr kIntMax = std::numeric_limits::max(); + STD_TORCH_CHECK(all_kv.numel() / 8 <= kIntMax && + kv_cache.stride(0) <= kIntMax && + kv_cache.size(0) * block_size <= kIntMax, + "all_kv / kv_cache too large for 32-bit indexing"); + + const torch::stable::accelerator::DeviceGuard device_guard( + all_kv.get_device_index()); + auto stream = get_current_cuda_stream(all_kv.get_device_index()); + // all_kv strides in 16-byte vectors. + auto launch = [&](auto fp8_tag) { + constexpr bool kFp8 = decltype(fp8_tag)::value; + vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3DSparkContextKV( + reinterpret_cast(all_kv.data_ptr()), + reinterpret_cast(k_norm_weight.data_ptr()), + reinterpret_cast(cos_sin_cache.data_ptr()), + reinterpret_cast(positions.data_ptr()), + reinterpret_cast(slot_mapping.data_ptr()), + reinterpret_cast(kv_cache.data_ptr()), + reinterpret_cast(cache_offsets.data_ptr()), + reinterpret_cast(k_scale.data_ptr()), + reinterpret_cast(scale_offsets.data_ptr()), + static_cast(eps), static_cast(rotary_dim), num_tokens, + num_layers, nkv, block_shift, static_cast(all_kv.stride(0) / 8), + static_cast(all_kv.stride(1) / 8), + static_cast(all_kv.stride(2) / 8), + static_cast(all_kv.stride(3) / 8), + static_cast(kv_cache.stride(0)), + static_cast(kv_cache.stride(1)), + static_cast(kv_cache.stride(2)), stream); + }; + if (fp8) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } +#endif +} diff --git a/csrc/fused_mqa_topk/candidates.cuh b/csrc/fused_mqa_topk/candidates.cuh new file mode 100644 index 000000000000..30afdec28efd --- /dev/null +++ b/csrc/fused_mqa_topk/candidates.cuh @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// SPDX-License-Identifier: MIT +// Copyright (c) 2025 DeepSeek +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. +// +// Candidate records of the gated scan: a uint16 value (the low 16 bits of the +// fp24 score code) and an int32 word (key index in bits 0-19, the code's high +// 8 bits in bits 20-27). +#pragma once +#include +#include + +#include "common.cuh" +#include "constants.h" + +namespace fused_mqa_topk { +using CandidateValue = uint16_t; +constexpr uint32_t kCandidateIndexBits = 20; +constexpr uint32_t kCandidateIndexMask = (1u << kCandidateIndexBits) - 1u; +static_assert(kMaxKeys <= (1u << kCandidateIndexBits), + "key indices must fit the candidate index field"); + +CUTLASS_DEVICE uint32_t candidate_pack_index(const uint32_t payload, + const uint32_t kv_index) { + return (kv_index & kCandidateIndexMask) | + ((payload >> 16) << kCandidateIndexBits); +} + +CUTLASS_DEVICE uint32_t candidate_fp24_code(const float value) { + return ordered_code(value) >> 8; +} + +CUTLASS_DEVICE uint32_t candidate_load_score_code(const CandidateValue value, + const int32_t packed_idx) { + return (static_cast(packed_idx) >> kCandidateIndexBits) << 16 | + static_cast(value); +} + +CUTLASS_DEVICE int32_t candidate_decode_index(const int32_t packed_idx) { + return static_cast(static_cast(packed_idx) & + kCandidateIndexMask); +} + +// Stores a candidate as its fp24 code (low 16 bits in value_dst, high 8 bits +// in bits 20-27 of the index word) and key index. +CUTLASS_DEVICE void store_candidate_record(CandidateValue* value_dst, + int32_t* index_dst, + const uint32_t payload, + const uint32_t kv_index) { + __stcg(value_dst, static_cast(payload)); + __stcg(index_dst, + static_cast(candidate_pack_index(payload, kv_index))); +} + +CUTLASS_DEVICE void store_candidate(CandidateValue* value_dst, + int32_t* index_dst, const float bq, + const uint32_t kv_index) { + store_candidate_record(value_dst, index_dst, candidate_fp24_code(bq), + kv_index); +} + +} // namespace fused_mqa_topk diff --git a/csrc/fused_mqa_topk/common.cuh b/csrc/fused_mqa_topk/common.cuh new file mode 100644 index 000000000000..860bed7c3dcc --- /dev/null +++ b/csrc/fused_mqa_topk/common.cuh @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Small device helpers shared by the persistent kernel, its gate sample and +// the exact rescan. +#pragma once + +#include +#include + +#include + +#include "constants.h" + +namespace fused_mqa_topk { + +// Warp roles of both pipelines (persistent kernel and dense rescan): warps 0-7 +// are the math warps (two warpgroups of 4 query rows each), warps 8-11 the +// producer / UMMA / drain warps. +constexpr uint32_t kMathThreads = 256, kMathWgThreads = 128, + kMathWarps = kMathThreads / 32; +// Query rows of a block each math warpgroup scores. +constexpr uint32_t kWgRows = kBlockRows / 2; +// Register split: ptxas gives each of the 384 threads 168 registers +// (65536 / 384, rounded down to a multiple of 8); setmaxnreg moves them from +// the producer / UMMA / drain warpgroup to the two math warpgroups. +constexpr uint32_t kMathRegs = 224, kSpecRegs = 168 * 3 - 2 * kMathRegs; + +// Named barrier ids (cutlass::arch::NamedBarrier offsets them past the ids it +// reserves). kBarMath / kBarSelect: the 256 math threads; kBarDrain: the math +// threads and the two drain warps; kBarDenseMath: the dense rescan's math +// threads. +constexpr uint32_t kBarDenseMath = 0, kBarMath = 2, kBarDrain = 4, + kBarSelect = 5; +constexpr uint32_t kDrainSyncThreads = kMathThreads + 2 * 32; + +// Barrier of the 256 math threads (warps 0-7); the compiler barriers keep +// shared-memory accesses on their side of it. +__device__ __forceinline__ void math_wg_sync() { + asm volatile("" ::: "memory"); + cutlass::arch::NamedBarrier(kMathThreads, kBarMath).sync(); + asm volatile("" ::: "memory"); +} + +// Shared-memory histograms pad one word per 32 bins (bank conflicts): +// bin b lives at padded_bin(b), a histogram of n bins takes padded_bins(n). +__host__ __device__ constexpr uint32_t padded_bins(uint32_t n) { + return n + n / 32; +} +__device__ __forceinline__ uint32_t padded_bin(uint32_t b) { + return b + b / 32; +} + +// Mask of all 32 lanes for the warp-synchronous intrinsics. +constexpr uint32_t kFullMask = 0xffffffffu; + +// Inclusive prefix sum of v over the warp's lanes. +__device__ __forceinline__ uint32_t warp_inclusive_scan(uint32_t v, + uint32_t lane) { +#pragma unroll + for (uint32_t d = 1; d < 32; d *= 2) { + const uint32_t u = __shfl_up_sync(kFullMask, v, d); + if (lane >= d) v += u; + } + return v; +} + +__device__ __forceinline__ void mbar_inval(const void* b) { + asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"( + static_cast(__cvta_generic_to_shared(b))) + : "memory"); +} + +// Reads `x` with a volatile load, so the compiler re-reads it from (shared) +// memory instead of keeping it live in a register. +template +__device__ __forceinline__ T volatile_load(const T& x) { + return *reinterpret_cast(&x); +} + +// Order-preserving uint32 code of a float (higher code = larger value) and +// its inverse. +__device__ __forceinline__ uint32_t ordered_code(float x) { + const uint32_t u = __float_as_uint(x); + return (u & 0x80000000u) ? ~u : (u ^ 0x80000000u); +} +__device__ __forceinline__ float ordered_float(uint32_t code) { + return __uint_as_float((code & 0x80000000u) ? (code ^ 0x80000000u) : ~code); +} + +// Histogram bin of a score: its fp16 value (rounded down) as an +// order-preserving 16-bit code, >> Shift. half_bin_floor is the smallest value +// in bin `b`. +template +__device__ __forceinline__ uint32_t half_bin(float x) { + const uint32_t u = __half_as_ushort(__float2half_rd(x)); + return ((u & 0x8000u) ? ((~u) & 65535u) : (u ^ 0x8000u)) >> Shift; +} +template +__device__ __forceinline__ float half_bin_floor(uint32_t b) { + const uint32_t o = b << Shift, + u = (o & 0x8000u) ? (o ^ 0x8000u) : ((~o) & 65535u); + return __half2float(__ushort_as_half(u)); +} + +// Weighted head reduction of one row's 32 head scores s_h in two FFMA2 chains +// seeded with sum_0 / sum_1: sum_h w_h * |s_h| (kWithLinear = false: the +// persistent kernel adds the linear part on the tensor core) or +// sum_h w_h * (s_h + |s_h|) (kWithLinear = true); relu(s) = (s + |s|) / 2. +// Adapted from DeepGEMM's sm100_fp8_mqa_logits.cuh at 88965b078186. +template +__device__ __forceinline__ float reduce_heads(const float* s, const float* w, + float2 sum_0 = {}, + float2 sum_1 = {}) { + auto term = [&](uint32_t at, float2 sum) { + const auto abs_s = make_float2(fabsf(s[at]), fabsf(s[at + 1])); + const auto x = + kWithLinear ? __fadd2_rn(make_float2(s[at], s[at + 1]), abs_s) : abs_s; + return __ffma2_rn(x, make_float2(w[at], w[at + 1]), sum); + }; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { + sum_0 = term(j, sum_0); + sum_1 = term(j + 2, sum_1); + } + const auto sum = __fadd2_rn(sum_0, sum_1); + return sum.x + sum.y; +} + +} // namespace fused_mqa_topk diff --git a/csrc/fused_mqa_topk/constants.h b/csrc/fused_mqa_topk/constants.h new file mode 100644 index 000000000000..c831bbfcfa6e --- /dev/null +++ b/csrc/fused_mqa_topk/constants.h @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Shapes and workspace layout shared by the fused MQA-logits + top-k kernel +// and its host wrapper. The Python side mirrors the workspace sizes in +// vllm/model_executor/kernels/attention/dsa/fused_mqa_topk.py (the host wrapper +// checks them). +#pragma once + +#include + +namespace fused_mqa_topk { + +// Indexer shape: 32 heads of dimension 128. +constexpr uint32_t kNumHeads = 32, kHeadDim = 128; +// Query rows per block (the unit a CTA claims), threads per CTA, keys per +// scan tile. +constexpr uint32_t kBlockRows = 8, kNumThreads = 384, kTileKeys = 256; +// Largest key range per call. +constexpr uint32_t kMaxKeys = 1u << 20; +// Per-row workspace: the gated candidate slab (a row with more candidates is +// rescanned exactly) and the exact rescan's retained set. +constexpr uint32_t kCandidatesPerRow = 16384, kRetainedPerRow = 4096; +// Candidate / retained slots (of kBlockRows rows each) per persistent CTA: one +// block in selection, one claimed ahead. +constexpr uint32_t kSlotsPerCta = 2; + +// Top-k sizes the kernel is compiled for. +constexpr bool is_supported_topk(int64_t k) { + return k == 512 || k == 1024 || k == 2048; +} + +} // namespace fused_mqa_topk diff --git a/csrc/fused_mqa_topk/deep_gemm_compat.cuh b/csrc/fused_mqa_topk/deep_gemm_compat.cuh new file mode 100644 index 000000000000..ae13a12c3d57 --- /dev/null +++ b/csrc/fused_mqa_topk/deep_gemm_compat.cuh @@ -0,0 +1,180 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Adapted from DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) at +// 88965b078186ee7510ab9fc4f1d5ebc19adfa8d1, the last release before its SM100 +// MQA-logits kernels went MX-only. FP8 MQA-logits building blocks: the +// shared-storage layout with fp32 per-key scales and fp32 weights, the +// grid-stride contiguous-KV scheduler and the 2D/3D TMA load. Trimmed to what +// the fused kernel instantiates (FP8 only, no metadata schedule, no multicast, +// no paging). +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace fused_mqa_topk::dg { +namespace math = deep_gemm::math; + +// FP8 Q/K (one byte per element), fp32 per-key scales, fp32 weights: the +// DeepGEMM layout without its MX / FP4 variants. +template +struct MQALogitsSharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + DG_STATIC_ASSERT(sizeof(qk_dtype_t) == 1, "FP8 Q/K only"); + + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); + // Align to one 8-row Q/K swizzle tile (head_dim bytes per FP8 token). + static constexpr uint32_t kSwizzleAlignment = 8 * kHeadDim; + static constexpr uint32_t kNumQElementsPerStage = UMMA_N * kHeadDim; + static constexpr uint32_t kNumKVElementsPerStage = SPLIT_KV * kHeadDim; + // TMA destinations in shared memory must be 128-byte aligned. + static constexpr uint32_t kTmaAlignment = 128; + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightBytesPerStage = + math::constexpr_align(BLOCK_Q * kNumWeightBytesPerRow, kTmaAlignment); + static constexpr uint32_t kNumWeightElementsPerStage = + kNumWeightBytesPerStage / static_cast(sizeof(reduce_dtype_t)); + + DG_STATIC_ASSERT(kNumQElementsPerStage % kSwizzleAlignment == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(kNumKVElementsPerStage % kSwizzleAlignment == 0, + "Unaligned TMA swizzling"); + DG_STATIC_ASSERT(kSwizzleAlignment % 128 == 0, + "TMA destination must be 128-byte aligned"); + + alignas(kSwizzleAlignment) + qk_dtype_t smem_q[kNumQStages][kNumQElementsPerStage]; + alignas(kSwizzleAlignment) + qk_dtype_t smem_kv[kNumKVStages][kNumKVElementsPerStage]; + alignas(kTmaAlignment) float smem_sf_kv[kNumKVStages][SPLIT_KV]; + alignas(kTmaAlignment) + reduce_dtype_t smem_weights[kNumQStages][kNumWeightElementsPerStage]; + // Barriers require 8-byte alignment, already guaranteed by the preceding + // TMA-aligned arrays. + Barrier full_q_barriers[kNumQStages]; + Barrier empty_q_barriers[kNumQStages]; + Barrier full_kv_barriers[kNumKVStages]; + Barrier empty_kv_barriers[kNumKVStages]; + Barrier full_tmem_barriers[kNumTmemStages]; + Barrier empty_tmem_barriers[kNumTmemStages]; + uint32_t tmem_ptr_in_smem; +}; + +// Grid-stride contiguous-KV scheduler: SM `s` owns Q blocks `s`, `s + kNumSMs`, +// ..., each visiting the KV range [min k_start, max k_end) of its rows. +template +struct SM100MQALogitsScheduler { + uint32_t current_q_block_idx; + uint32_t num_q_blocks; + uint32_t num_q_tokens; + uint32_t num_kv_tokens; + const uint32_t* cu_seq_len_k_start; + const uint32_t* cu_seq_len_k_end; + + CUTLASS_DEVICE SM100MQALogitsScheduler(const uint32_t& sm_idx, + const uint32_t& num_q_tokens, + const uint32_t& num_kv_tokens, + const uint32_t* cu_seq_len_k_start, + const uint32_t* cu_seq_len_k_end) + : current_q_block_idx(sm_idx), + num_q_blocks(math::ceil_div(num_q_tokens, BLOCK_Q)), + num_q_tokens(num_q_tokens), + num_kv_tokens(num_kv_tokens), + cu_seq_len_k_start(cu_seq_len_k_start), + cu_seq_len_k_end(cu_seq_len_k_end) { + DG_STATIC_ASSERT(kNumSMs > 0, "Invalid SM count"); + } + + // NOTE: Out-of-range rows are clamped to the last valid row + CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx, + uint32_t& kv_token_base, + uint32_t& num_kv_splits) { + if (current_q_block_idx >= num_q_blocks) return false; + q_block_idx = current_q_block_idx; + current_q_block_idx += kNumSMs; + uint32_t start = cute::numeric_limits::max(); + uint32_t end = cute::numeric_limits::min(); +#pragma unroll 8 + for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++token_idx) { + const auto row_idx = + cute::min(q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1); + start = cute::min(start, + cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens)); + end = cute::max(end, cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens)); + } + kv_token_base = start / 4 * 4; + num_kv_splits = math::ceil_div(end - kv_token_base, SPLIT_KV); + return true; + } + + CUTLASS_DEVICE uint32_t + get_q_tma_token_base(const uint32_t& q_block_idx) const { + return q_block_idx * BLOCK_Q; + } + + CUTLASS_DEVICE static uint32_t get_kv_tma_offset( + const uint32_t& kv_token_base, const uint32_t& kv_split_idx) { + return kv_token_base + kv_split_idx * SPLIT_KV; + } + + CUTLASS_DEVICE static uint32_t get_logits_row(const uint32_t& q_block_idx, + const uint32_t& token_idx) { + return q_block_idx * BLOCK_Q + token_idx; + } + + CUTLASS_DEVICE static uint32_t get_logits_col( + const uint32_t& kv_token_base, const uint32_t& kv_split_idx, + const uint32_t& math_thread_idx) { + return kv_token_base + kv_split_idx * SPLIT_KV + math_thread_idx; + } +}; + +// Single-CTA TMA load of a BLOCK_OUTER x BLOCK_INNER box (3D adds a batch +// coordinate). +template +CUTLASS_DEVICE void tma_copy( + void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr, + dtype_t* smem_ptr, const uint32_t& inner_idx, const uint32_t& outer_idx, + const uint32_t& batch_idx = 0) { + constexpr auto kCacheHint = + static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL); + constexpr uint32_t kPackFactor = deep_gemm::get_smem_pack_factor(); + DG_STATIC_ASSERT(kPackFactor == 1 or sizeof(dtype_t) == 1, + "Packing expects a 1-byte storage type"); + constexpr uint32_t BLOCK_INNER_ATOM_STORAGE = + kSwizzleMode == 0 ? BLOCK_INNER / kPackFactor + : kSwizzleMode / sizeof(dtype_t); + constexpr uint32_t BLOCK_INNER_ATOM_LOGICAL = + BLOCK_INNER_ATOM_STORAGE * kPackFactor; + DG_STATIC_ASSERT(BLOCK_INNER % BLOCK_INNER_ATOM_LOGICAL == 0, + "TMA inner block must contain whole atoms"); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++i) { + if constexpr (kIs3DTMA) { + cute::SM90_TMA_LOAD_3D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx, batch_idx); + } else { + cute::SM90_TMA_LOAD_2D::copy( + desc_ptr, reinterpret_cast(barrier_ptr), kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx); + } + } +} + +} // namespace fused_mqa_topk::dg diff --git a/csrc/fused_mqa_topk/dense_scan.cuh b/csrc/fused_mqa_topk/dense_scan.cuh new file mode 100644 index 000000000000..395cb84f183e --- /dev/null +++ b/csrc/fused_mqa_topk/dense_scan.cuh @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Adapted from DeepGEMM (https://github.com/deepseek-ai/DeepGEMM) at +// 88965b078186ee7510ab9fc4f1d5ebc19adfa8d1: +// deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh. +// +// Dense FP8 MQA scoring of one 8-row block over its whole KV range: DeepGEMM's +// SM100 MQA logits pipeline (contiguous KV, 32 heads, D = 128) with the logits +// store replaced by a per-score policy. Used by the persistent kernel's exact +// recovery rescan (rescan.cuh). +// +// Roles (384 threads): warps 0-7 reduce the 32 heads (4 rows per warpgroup), +// warp 8 loads Q and the weights, warp 9 the KV tiles and their scales, warp 10 +// issues the UMMAs, warp 11 is idle. +// +// Policy interface (all called by the math warps): +// kWeightScale: score = (sum_h kWeightScale * w_h * (s_h + |s_h|)) * scale_kv +// begin_block(sched, q_block, query_base): after the block's Q / weights +// landed +// begin_split(split): before the split's scores +// score(i, kv_offset, kv_lane, split, score): row query_base + i of the +// warpgroup +// end_split(split, num_splits, kv_base): after the split's scores +// end_block(sched, q_block): before the block's Q stage is released +#pragma once + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common.cuh" +#include "constants.h" +#include "deep_gemm_compat.cuh" + +namespace fused_mqa_topk::dense { +using namespace deep_gemm; + +constexpr uint32_t BLOCK_Q = kBlockRows, SPLIT_KV = kTileKeys; +constexpr uint32_t kNumKVStages = 2, kNumTmemStages = 2; +constexpr uint32_t UMMA_M = 128, UMMA_N = BLOCK_Q * kNumHeads, UMMA_K = 32; +// The rescan scores exactly one 8-row block: the scheduler is built with that +// block as its `sm_idx`, and its grid stride (kNumSMs) exceeds any block count +// (rows < 2^26), so next_q_block yields that block only. A smaller stride would +// also score block + stride, outside the rescan's row bounds and output rows. +constexpr uint32_t kSingleBlockStride = 1u << 31; +using Scheduler = + dg::SM100MQALogitsScheduler; +using SharedStorage = + dg::MQALogitsSharedStorage; + +// The policy's state lives in dynamic shared memory right after the pipeline +// storage. Barrier words are zeroed before init: the persistent kernel reuses +// this memory between rescans. +template +CUTLASS_DEVICE void score_block(const Scheduler& block, Policy& policy, + const cute::TmaDescriptor& tensor_map_q, + const cute::TmaDescriptor& tensor_map_kv, + const cute::TmaDescriptor& tensor_map_sf_kv, + const cute::TmaDescriptor& tensor_map_weights) { + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr uint32_t kSpecWarpStart = 8; + constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = BLOCK_Q * kNumHeads * kHeadDim; + constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = SPLIT_KV * kHeadDim; + constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = SPLIT_KV * sizeof(float); + constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = + BLOCK_Q * kNumHeads * sizeof(float); + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols(); + + if (warp_idx == kSpecWarpStart) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_weights); + cute::prefetch_tma_descriptor(&tensor_map_kv); + cute::prefetch_tma_descriptor(&tensor_map_sf_kv); + } + + extern __shared__ __align__(SharedStorage::kSwizzleAlignment) + uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + if (warp_idx == kSpecWarpStart + 1 and cute::elect_one_sync()) { + const auto init = [](auto& barrier, uint32_t count) { + *reinterpret_cast(&barrier) = 0; + barrier.init(count); + }; + init(smem.full_q_barriers[0], 1); + init(smem.empty_q_barriers[0], kMathThreads + 32); // + the UMMA warp +#pragma unroll + for (uint32_t i = 0; i < kNumKVStages; ++i) { + init(smem.full_kv_barriers[i], 1); + init(smem.empty_kv_barriers[i], kMathThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumTmemStages; ++i) { + init(smem.full_tmem_barriers[i], 1); + init(smem.empty_tmem_barriers[i], kMathThreads); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + if (warp_idx == kSpecWarpStart + 2) + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); + __syncthreads(); + + RingPipeline<1> q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + + uint32_t q_block_idx, kv_base, num_kv_splits; + if (warp_idx == kSpecWarpStart) { + cutlass::arch::warpgroup_reg_dealloc(); + if (cute::elect_one_sync()) { + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.empty_q_barriers[q_stage_idx].wait(q_phase ^ 1); + const uint32_t q_token_base = + scheduler.get_q_tma_token_base(q_block_idx); + dg::tma_copy( + &tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); + dg::tma_copy( + &tensor_map_weights, &smem.full_q_barriers[q_stage_idx], + smem.smem_weights[q_stage_idx], 0, q_token_base); + smem.full_q_barriers[q_stage_idx].arrive_and_expect_tx( + SMEM_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE); + } + } + __syncwarp(); + } else if (warp_idx == kSpecWarpStart + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + if (cute::elect_one_sync()) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); + const uint32_t kv_tma_offset = + scheduler.get_kv_tma_offset(kv_base, kv_split_idx); + dg::tma_copy( + &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); + dg::tma_copy( + &tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], + smem.smem_sf_kv[kv_stage_idx], kv_tma_offset, 0); + smem.full_kv_barriers[kv_stage_idx].arrive_and_expect_tx( + SMEM_KV_SIZE_PER_STAGE + SMEM_SF_KV_SIZE_PER_STAGE); + } + __syncwarp(); + } + } + } else if (warp_idx == kSpecWarpStart + 2) { + cutlass::arch::warpgroup_reg_dealloc(); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); + auto instr_desc = cute::UMMA::make_instr_desc< + cutlass::float_e4m3_t, cutlass::float_e4m3_t, float, UMMA_M, UMMA_N, + cute::UMMA::Major::K, cute::UMMA::Major::K>(); + auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); + if (cute::elect_one_sync()) { + // One UMMA per math warpgroup: KV rows [128 i, 128 i + 128) x all 256 + // Q rows. +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + const uint32_t tmem_addr = tmem_stage_idx * UMMA_N; + smem.empty_tmem_barriers[tmem_stage_idx].wait(tmem_phase ^ 1); + ptx::tcgen05_after_thread_sync(); +#pragma unroll + for (uint32_t k = 0; k < kHeadDim / UMMA_K; ++k) { + auto a_desc = mma::sm100::make_umma_desc( + smem.smem_kv[kv_stage_idx], i * UMMA_M, k * UMMA_K); + auto b_desc = mma::sm100::make_umma_desc( + smem.smem_q[q_stage_idx], 0, k * UMMA_K); + ptx::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, tmem_addr, k, + runtime_instr_desc); + } + ptx::umma_arrive_no_elect(smem.full_tmem_barriers[tmem_stage_idx]); + } + } + __syncwarp(); + } + smem.empty_q_barriers[q_stage_idx].arrive(); + } + } else if (warp_idx == kSpecWarpStart + 3) { + cutlass::arch::warpgroup_reg_dealloc(); + } else { + cutlass::arch::warpgroup_reg_alloc(); + const uint32_t math_thread_idx = warp_idx * 32 + lane_idx; + const uint32_t query_base = (warp_idx / 4) * 4; + float weights[4][kNumHeads]; + Scheduler scheduler = block; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; ++j) + weights[i][j] = Policy::kWeightScale * + ptx::ld_shared(smem.smem_weights[q_stage_idx] + + (query_base + i) * kNumHeads + j); + policy.begin_block(scheduler, q_block_idx, query_base); + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; + ++kv_split_idx) { + policy.begin_split(kv_split_idx); + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); +#pragma unroll + for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) { + const uint32_t kv_lane = + (math_thread_idx % kMathWgThreads) + kv_half * kMathWgThreads; + const uint32_t kv_offset = + scheduler.get_logits_col(kv_base, kv_split_idx, kv_lane); + const float scale_kv = + ptx::ld_shared(smem.smem_sf_kv[kv_stage_idx] + kv_lane); + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + smem.full_tmem_barriers[tmem_stage_idx].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + // Both KV half-MMAs have committed before the shared stage is freed. + if (kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + smem.empty_kv_barriers[kv_stage_idx].arrive(); + } + float accum_pair[2][kNumHeads]; +#pragma unroll + for (uint32_t pair = 0; pair < 4; pair += 2) { + // Two query rows at a time keep the live accumulator registers low. + ptx::tmem_load_32dp32b<2 * kNumHeads>( + tmem_stage_idx * UMMA_N + (query_base + pair) * kNumHeads, + reinterpret_cast(accum_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair == 2) { + ptx::tcgen05_before_thread_sync(); + smem.empty_tmem_barriers[tmem_stage_idx].arrive(); + } +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) + policy.score( + pair + r, kv_offset, kv_lane, kv_split_idx, + reduce_heads(accum_pair[r], weights[pair + r]) * + scale_kv); + } + } + policy.end_split(kv_split_idx, num_kv_splits, kv_base); + } + policy.end_block(scheduler, q_block_idx); + cutlass::arch::fence_view_async_shared(); + smem.empty_q_barriers[q_stage_idx].arrive(); + } + cutlass::arch::NamedBarrier(kMathThreads, kBarDenseMath).sync(); + if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } +} + +// Invalidates the pipeline's barriers so the next score_block can re-initialise +// the same shared memory (mbarrier.inval before reuse). +CUTLASS_DEVICE void invalidate_barriers(SharedStorage& smem) { + for (auto& b : smem.full_q_barriers) mbar_inval(&b); + for (auto& b : smem.empty_q_barriers) mbar_inval(&b); + for (auto& b : smem.full_kv_barriers) mbar_inval(&b); + for (auto& b : smem.empty_kv_barriers) mbar_inval(&b); + for (auto& b : smem.full_tmem_barriers) mbar_inval(&b); + for (auto& b : smem.empty_tmem_barriers) mbar_inval(&b); +} + +} // namespace fused_mqa_topk::dense diff --git a/csrc/fused_mqa_topk/fused_mqa_topk.cu b/csrc/fused_mqa_topk/fused_mqa_topk.cu new file mode 100644 index 000000000000..652453494fbe --- /dev/null +++ b/csrc/fused_mqa_topk/fused_mqa_topk.cu @@ -0,0 +1,394 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Single-launch fused FP8 MQA-logits + top-k (k = 512, 1024, 2048) prefill +// selection (SM100 family): one persistent kernel (persistent_kernel.cuh, one +// CTA per SM; sampled gate, candidate emission and selection, exact rescan of +// out-of-range blocks) scores every visible key of a row and selects its top-k +// keys without materializing the [rows, keys] logits matrix. +// This file is the libtorch-stable host wrapper. All scratch memory is passed +// in by the caller (no allocation here), so the op is CUDA-graph friendly. + +#include "../torch_utils.h" + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "constants.h" +#include "persistent_kernel.cuh" + +namespace { + +using fused_mqa_topk::kBlockRows; +using fused_mqa_topk::kCandidatesPerRow; +using fused_mqa_topk::kHeadDim; +using fused_mqa_topk::kMaxKeys; +using fused_mqa_topk::kNumHeads; +using fused_mqa_topk::kNumThreads; +using fused_mqa_topk::kRetainedPerRow; +using fused_mqa_topk::kSlotsPerCta; +using fused_mqa_topk::kTileKeys; +using torch::headeronly::ScalarType; +namespace persist = fused_mqa_topk::persist; + +void check_cuda(cudaError_t err) { + STD_TORCH_CHECK(err == cudaSuccess, cudaGetErrorString(err)); +} + +// cuTensorMapEncodeTiled from the driver, resolved once through the runtime +// (thread-safe static initialisation; no link-time libcuda dependency). +PFN_cuTensorMapEncodeTiled_v12000 encode_tiled() { + static const PFN_cuTensorMapEncodeTiled_v12000 fn = [] { + void* ptr = nullptr; + cudaDriverEntryPointQueryResult status; + const cudaError_t err = cudaGetDriverEntryPointByVersion( + "cuTensorMapEncodeTiled", &ptr, 12000, cudaEnableDefault, &status); + STD_TORCH_CHECK(err == cudaSuccess && status == cudaDriverEntryPointSuccess, + "failed to resolve cuTensorMapEncodeTiled"); + return reinterpret_cast(ptr); + }(); + return fn; +} + +// SM count of `device`, queried once per device. +int sm_count(int device) { + static std::mutex mutex; + static std::map cache; + const std::lock_guard lock(mutex); + auto [it, inserted] = cache.try_emplace(device, 0); + if (inserted) { + const cudaError_t err = cudaDeviceGetAttribute( + &it->second, cudaDevAttrMultiProcessorCount, device); + if (err != cudaSuccess) cache.erase(it); + check_cuda(err); + } + return it->second; +} + +// Adapted from DeepGEMM's make_tma_2d_desc +// (csrc/jit_kernels/impls/runtime_utils.hpp at 88965b078186) +CUtensorMap make_2d(void* ptr, CUtensorMapDataType dt, int elem_size, + int gmem_inner, int gmem_outer, int smem_inner, + int smem_outer, long gmem_outer_stride, int swizzle_mode) { + if (swizzle_mode != 0) smem_inner = swizzle_mode / elem_size; + CUtensorMap tm; + const cuuint64_t gdims[2] = {(cuuint64_t)gmem_inner, (cuuint64_t)gmem_outer}; + const cuuint32_t sdims[2] = {(cuuint32_t)smem_inner, (cuuint32_t)smem_outer}; + const cuuint64_t gstrides[1] = {(cuuint64_t)(gmem_outer_stride * elem_size)}; + const cuuint32_t estrides[2] = {1, 1}; + CUtensorMapSwizzle swizzle = swizzle_mode == 128 ? CU_TENSOR_MAP_SWIZZLE_128B + : swizzle_mode == 64 ? CU_TENSOR_MAP_SWIZZLE_64B + : swizzle_mode == 32 + ? CU_TENSOR_MAP_SWIZZLE_32B + : CU_TENSOR_MAP_SWIZZLE_NONE; + CUresult r = encode_tiled()(&tm, dt, 2, ptr, gdims, gstrides, sdims, estrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + STD_TORCH_CHECK(r == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", (int)r); + return tm; +} + +// KV [keys,128] bytes viewed as {128, groups, stride} with strides +// {128*stride, 128}: box {128, 256, 1} at (0, 0, o) loads keys o, o+stride, ... +// with 128B swizzle. The view spans groups * stride keys, which can exceed +// `keys` by < stride / 2; that is safe only because the kernel loads offset +// o = stride / 2, whose last key ((groups - 1) * stride + stride / 2) is +// always < keys (launch_config sizes `groups` for that). +CUtensorMap make_strided_kv(void* ptr, uint32_t groups, uint32_t stride) { + CUtensorMap tm; + constexpr uint64_t kRowBytes = kHeadDim; + const cuuint64_t gdims[3] = {kRowBytes, groups, stride}; + const cuuint64_t gstrides[2] = {kRowBytes * stride, kRowBytes}; + const cuuint32_t box[3] = {kRowBytes, kTileKeys, 1}; + const cuuint32_t estrides[3] = {1, 1, 1}; + CUresult r = encode_tiled()( + &tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, ptr, gdims, gstrides, box, + estrides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + STD_TORCH_CHECK(r == CUDA_SUCCESS, + "strided KV cuTensorMapEncodeTiled failed: ", (int)r); + return tm; +} + +// -------- launch configuration -------- +// Keys per row above which the long-context instantiations (Interval 64, 128) +// and gate settings apply. +constexpr uint32_t kLongContextKeys = 131072; + +using PersistentKernel = + decltype(&persist::fused_mqa_topk_fp8_prefill_kernel<16, 2048>); + +// vLLM dispatches chunks of >= 32K keys per row here (fused_mqa_topk_min_keys); +// smaller key counts are supported but take the same <= 32K configuration. +template +PersistentKernel persistent_kernel(uint32_t keys) { + return keys <= 32768 ? persist::fused_mqa_topk_fp8_prefill_kernel<16, TopK> + : keys <= kLongContextKeys + ? persist::fused_mqa_topk_fp8_prefill_kernel<32, TopK> + : keys <= 524288 + ? persist::fused_mqa_topk_fp8_prefill_kernel<64, TopK> + : persist::fused_mqa_topk_fp8_prefill_kernel<128, TopK>; +} + +struct LaunchConfig { + PersistentKernel kernel; + int smem; // dynamic shared memory bytes + uint32_t phase_bytes; // selection scratch bytes past the Q / weights region + // Gate sample: sample_tiles tiles of 256 keys on a uniform grid (stride + // sample_stride). + uint32_t sample_target, sample_stride, sample_groups, sample_tiles; +}; + +// Dynamic shared memory of `kernel` on `device`: everything past its static +// shared memory up to the device's opt-in limit (232448 bytes on SM100; the +// staged selection uses it all), at least the exact rescan's layout. Sets +// the opt-in attribute; both are computed once per device and kernel. +cudaError_t kernel_smem(int device, PersistentKernel kernel, int& smem) { + // The exact rescan reuses the same dynamic shared memory with its own layout. + constexpr int kRescanSmem = int(sizeof(fused_mqa_topk::dense::SharedStorage) + + sizeof(persist::RescanShared)); + constexpr int kMinSmem = int( + sizeof(persist::KernelSmem) + + std::max(sizeof(persist::HandoffRing), sizeof(persist::SampleShared))); + static std::mutex mutex; + static std::map, int> cache; + const std::lock_guard lock(mutex); + if (const auto it = cache.find({device, kernel}); it != cache.end()) { + smem = it->second; + return cudaSuccess; + } + int max_smem = 0; + if (const cudaError_t e = cudaDeviceGetAttribute( + &max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + e != cudaSuccess) + return e; + cudaFuncAttributes attr; + if (const cudaError_t e = + cudaFuncGetAttributes(&attr, reinterpret_cast(kernel)); + e != cudaSuccess) + return e; + const int bytes = std::max(kRescanSmem, max_smem - int(attr.sharedSizeBytes)); + if (bytes < kMinSmem) return cudaErrorInvalidConfiguration; + if (const cudaError_t e = cudaFuncSetAttribute( + reinterpret_cast(kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + e != cudaSuccess) + return e; + cache.emplace(std::make_pair(device, kernel), bytes); + smem = bytes; + return cudaSuccess; +} + +// Picks the kernel instantiation, its shared memory and the gate sample for +// `keys` keys and top-`topk` rows on `device` (the current device). +cudaError_t launch_config(int device, uint32_t keys, uint32_t topk, + LaunchConfig& c) { + // The gate aims for `target` candidates per row and sizes the sample so its + // quantile rank is `rank` (count error ~1/sqrt(rank)); a row whose count + // falls below the top-k is rescanned exactly. Up to 128K keys the candidate + // stream dominates, so smaller top-ks aim lower with a coarser sample. Above + // 128K the sample is clamped to 16-32 tiles (rank ~32 at 1M keys): a sample + // tile costs about two scan tiles, so a larger sample costs more than a + // smaller candidate set saves. Tuned on B300 real captures. + struct Gate { + uint32_t target, rank; + }; + constexpr Gate kGates[3][2] = { + // <= 128K > 128K + {{1024, 32}, {1536, 64}}, // top-512 + {{1536, 48}, {2048, 48}}, // top-1024 + {{3072, 64}, {4096, 64}}, // top-2048 + }; + const bool long_ctx = keys > kLongContextKeys; + const Gate gate = + kGates[topk == 512 ? 0 : (topk == 1024 ? 1 : 2)][long_ctx ? 1 : 0]; + c.sample_target = gate.target; + const uint32_t rank = gate.rank; + const uint64_t tile = kTileKeys; + const uint32_t rank_tiles = std::max( + 1u, uint32_t((uint64_t(rank) * keys + tile * c.sample_target - 1) / + (tile * c.sample_target))); + c.sample_tiles = + long_ctx ? std::min(32u, std::max(16u, rank_tiles)) : rank_tiles; + c.sample_stride = std::max(1u, uint32_t((keys + tile * c.sample_tiles - 1) / + (tile * c.sample_tiles))); + c.sample_groups = + (keys - c.sample_stride / 2 + c.sample_stride - 1) / c.sample_stride; + c.kernel = topk == 512 ? persistent_kernel<512>(keys) + : topk == 1024 ? persistent_kernel<1024>(keys) + : persistent_kernel<2048>(keys); + if (const cudaError_t e = kernel_smem(device, c.kernel, c.smem); + e != cudaSuccess) + return e; + c.phase_bytes = uint32_t(c.smem) - persist::kSelBase; + return cudaSuccess; +} + +void check_tensor(const torch::stable::Tensor& t, int32_t device, + ScalarType dtype, const char* name) { + STD_TORCH_CHECK(t.is_cuda() && t.get_device_index() == device, name, + " must be on the same CUDA device as q"); + STD_TORCH_CHECK(t.scalar_type() == dtype, name, " has the wrong dtype"); + STD_TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +// TMA tensor maps need 16-byte aligned base addresses. +void check_aligned(const torch::stable::Tensor& t, const char* name) { + STD_TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +} // namespace + +// Inputs: q [rows,32,128] fp8e4m3, k [keys,128] fp8e4m3, scales [keys] fp32, +// weights [rows,32] fp32, starts/ends [rows] int32 (the visible key range +// [start,end) of each row), with 0 < rows < 2^26 and 0 < keys <= 1048576 (the +// last 8-row block may be partial). starts/ends are expected in [0, keys] with +// start <= end; they are read as uint32_t and clamped to keys, so a negative +// value means `keys` and an end before its start means an empty row. +// +// Output: out [rows,K] int32 with K (the top-k) 512, 1024 or 2048: key indices +// relative to the row's start (key - start, the top_k_per_row_prefill +// convention), unordered; -1 padded when fewer than K keys are visible. Rows +// that see at most K keys get them all, in order. Exact up to ties: the gated +// path ranks fp24 score codes (the low 8 mantissa bits dropped), the exact +// rescan ranks fp32 scores. +// +// All tensors are contiguous; q, k, scales, weights and out are 16-byte aligned +// (TMA / vector stores). +// +// Workspace (caller-owned, reused across calls), with B = ceil(rows / 8) and +// S >= 2 slots per CTA of the persistent grid (one CTA per SM): +// counter [2] int32 (zero between calls; the kernel leaves it zeroed, so a +// workspace must not be shared by launches that can run concurrently, e.g. +// on two streams), +// values [8S,16384] uint16 (low 16 bits of the fp24 candidate codes), +// indices [8S,16384] int32 (key index in bits 0-19, the code's high 8 bits +// in bits 20-27), +// retained [8S,4096] fp32, retained_indices [8S,4096] int32, +// counts [>=8B] int32, defer [>=B] int32. +void fused_mqa_topk_prefill( + torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor scales, torch::stable::Tensor weights, + torch::stable::Tensor starts, torch::stable::Tensor ends, + torch::stable::Tensor out, torch::stable::Tensor counter, + torch::stable::Tensor values, torch::stable::Tensor indices, + torch::stable::Tensor retained, torch::stable::Tensor retained_indices, + torch::stable::Tensor counts, torch::stable::Tensor defer) { + STD_TORCH_CHECK(q.is_cuda(), "q must be a CUDA tensor"); + const int32_t dev = q.get_device_index(); + check_tensor(q, dev, ScalarType::Float8_e4m3fn, "q"); + check_tensor(k, dev, ScalarType::Float8_e4m3fn, "k"); + check_tensor(scales, dev, ScalarType::Float, "scales"); + check_tensor(weights, dev, ScalarType::Float, "weights"); + check_tensor(starts, dev, ScalarType::Int, "starts"); + check_tensor(ends, dev, ScalarType::Int, "ends"); + check_tensor(out, dev, ScalarType::Int, "out"); + check_tensor(counter, dev, ScalarType::Int, "counter"); + check_tensor(values, dev, ScalarType::UInt16, "values"); + check_tensor(indices, dev, ScalarType::Int, "indices"); + check_tensor(retained, dev, ScalarType::Float, "retained"); + check_tensor(retained_indices, dev, ScalarType::Int, "retained_indices"); + check_tensor(counts, dev, ScalarType::Int, "counts"); + check_tensor(defer, dev, ScalarType::Int, "defer"); + + STD_TORCH_CHECK( + q.dim() == 3 && q.size(1) == kNumHeads && q.size(2) == kHeadDim, + "q must be [rows, 32, 128]"); + STD_TORCH_CHECK(k.dim() == 2 && k.size(1) == kHeadDim, + "k must be [keys, 128]"); + const int64_t rows = q.size(0), keys = k.size(0); + // TMA row coordinates (rows * 32 Q rows) are int32. + STD_TORCH_CHECK(rows > 0 && rows < (int64_t(1) << 26), + "rows must be in [1, 2^26)"); + STD_TORCH_CHECK(keys > 0 && keys <= kMaxKeys, "keys must be in [1, 1048576]"); + const int64_t blocks = (rows + kBlockRows - 1) / kBlockRows; + STD_TORCH_CHECK(scales.numel() == keys, "scales must be [keys]"); + STD_TORCH_CHECK(weights.dim() == 2 && weights.size(0) == rows && + weights.size(1) == kNumHeads, + "weights must be [rows, 32]"); + STD_TORCH_CHECK(starts.numel() == rows && ends.numel() == rows, + "starts/ends must be [rows]"); + STD_TORCH_CHECK(out.dim() == 2 && out.size(0) == rows && + fused_mqa_topk::is_supported_topk(out.size(1)), + "out must be [rows, K] with K in {512, 1024, 2048}"); + check_aligned(q, "q"); + check_aligned(k, "k"); + check_aligned(scales, "scales"); + check_aligned(weights, "weights"); + check_aligned(out, "out"); + + torch::stable::accelerator::DeviceGuard guard(dev); + const cudaStream_t stream = get_current_cuda_stream(dev); + const uint32_t urows = static_cast(rows), + ukeys = static_cast(keys), + ublocks = static_cast(blocks), + topk = static_cast(out.size(1)); + const auto* starts_p = static_cast(starts.data_ptr()); + const auto* ends_p = static_cast(ends.data_ptr()); + auto* out_p = static_cast(out.data_ptr()); + // Q: rows * 32 head rows of 128 bytes, a block's 256 per box; K: keys rows of + // 128 bytes, a tile's 256 per box (both 128B-swizzled); per-key scales: a + // tile's 256; weights: a block's 8 rows of 32. + const int head_rows = int(rows) * int(kNumHeads); + auto tq = make_2d(q.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, kHeadDim, + head_rows, kHeadDim, kBlockRows * kNumHeads, kHeadDim, 128); + auto tk = make_2d(k.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_UINT8, 1, kHeadDim, + int(keys), kHeadDim, kTileKeys, kHeadDim, 128); + auto ts = make_2d(scales.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, + int(keys), 1, kTileKeys, 1, 0, 0); + auto tw = make_2d(weights.data_ptr(), CU_TENSOR_MAP_DATA_TYPE_FLOAT32, 4, + kNumHeads, int(rows), kNumHeads, kBlockRows, kNumHeads, 0); + + STD_TORCH_CHECK(counter.numel() == 2, "counter must be [2]"); + constexpr int64_t kCtaRows = kSlotsPerCta * kBlockRows; + STD_TORCH_CHECK(values.dim() == 2 && values.size(1) == kCandidatesPerRow && + values.size(0) % kCtaRows == 0 && indices.dim() == 2 && + indices.size(0) == values.size(0) && + indices.size(1) == kCandidatesPerRow, + "values/indices must be [8 * slots, 16384], slots even"); + STD_TORCH_CHECK(values.size(0) > 0, "the candidate workspace is empty"); + STD_TORCH_CHECK(retained.dim() == 2 && retained.size(1) == kRetainedPerRow && + retained.size(0) == values.size(0) && + retained_indices.dim() == 2 && + retained_indices.size(0) == retained.size(0) && + retained_indices.size(1) == kRetainedPerRow, + "retained/retained_indices must be [8 * slots, 4096]"); + STD_TORCH_CHECK( + counts.numel() >= blocks * kBlockRows && defer.numel() >= blocks, + "per-row workspace too small"); + // One persistent CTA per SM (capped by the blocks and the workspace slots); + // CTA c owns candidate slots 2c and 2c + 1. + const uint32_t grid = static_cast(std::min( + {int64_t(sm_count(dev)), blocks, values.size(0) / kCtaRows})); + LaunchConfig c; + const cudaError_t config_err = launch_config(dev, ukeys, topk, c); + STD_TORCH_CHECK(config_err != cudaErrorInvalidConfiguration, + "fused_mqa_topk_prefill needs more shared memory per block " + "than this GPU allows"); + check_cuda(config_err); + auto tks = make_strided_kv(k.data_ptr(), c.sample_groups, c.sample_stride); + const fused_mqa_topk::rescan::RescanBuffers rescan{ + static_cast(retained.data_ptr()), + static_cast(retained_indices.data_ptr())}; + c.kernel<<>>( + urows, ukeys, starts_p, ends_p, out_p, + static_cast(values.data_ptr()), + static_cast(indices.data_ptr()), + static_cast(counts.data_ptr()), rescan, c.phase_bytes, + c.sample_target, c.sample_stride, c.sample_groups, c.sample_tiles, + static_cast(scales.data_ptr()), tks, tq, ts, tk, tw, + ublocks, static_cast(counter.data_ptr()), + static_cast(defer.data_ptr())); + check_cuda(cudaGetLastError()); +} diff --git a/csrc/fused_mqa_topk/persistent_epoch.inc b/csrc/fused_mqa_topk/persistent_epoch.inc new file mode 100644 index 000000000000..1af6a1d2088b --- /dev/null +++ b/csrc/fused_mqa_topk/persistent_epoch.inc @@ -0,0 +1,742 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Body of one persistent epoch (included twice by persistent_kernel.cuh: the +// first epoch is straight-line code so the hot loops are not nested in an outer +// loop). +if (warp_idx == 9 && cute::elect_one_sync()) { + if (epoch) + for (uint32_t i = 0; i < BarIdx::total; ++i) mbar_inval(bar_words + i); + for (uint32_t i = 0; i < BarIdx::total; ++i) + reinterpret_cast(bar_words)[i] = 0; + full_q.init(1); + empty_q.init(kMathThreads + 32); // math threads + the UMMA warp + for (uint32_t i = 0; i < kKVStages; ++i) { + full_kv[i].init(kFullKvArrivals); + empty_kv[i].init(kMathThreads); + } + for (uint32_t i = 0; i < BarIdx::kTmemSlots; ++i) { + full_tmem[i].init(1); + empty_tmem[i].init(kMathWgThreads); + } + full_lin->init(32); + // Every reader lane arrives on claim_empty. + for (uint32_t i = 0; i < 2; ++i) { + claim_full[i].init(32); + claim_empty[i].init(kClaimConsumers * 32); + } + sel_done.init(kMathWarps); // one lane per math warp + for (uint32_t i = 0; i < 2 * kMathWarps; ++i) { + h_full[i].init(32); + h_empty[i].init(32); + } + cutlass::arch::fence_barrier_init(); + unsafe_rows = 0; +} +__syncwarp(); +if (warp_idx == 10) cute::TMEM::Allocator1Sm().allocate(512, &tmem_ptr); +__syncthreads(); + +if (warp_idx == 8) { + // -------- claimer + Q/weights producer + drain of WG1 rows -------- + cutlass::arch::warpgroup_reg_dealloc(); + // n_blk: blocks this CTA processed (all-short blocks excluded). + uint32_t n_blk = 0, drain_interval = 0; + for (uint32_t k = 0;; ++k) { + if (lane == 0) { + empty_q.wait((n_blk & 1) ^ 1); + } + __syncwarp(); + // Every lane waits: lanes 0-7 write the entry. + claim_empty[k & 1].wait(((k >> 1) & 1) ^ 1); + uint32_t b = 0; + if (lane == 0) + b = atomicOr(&pause_req, 0u) ? kPause : atomicAdd(counter, 1u); + b = __shfl_sync(kFullMask, b, 0); + ClaimEntry& e = entries[k & 1]; + if (b >= nblocks) { + if (lane == 0) { + const uint32_t kind = b == kPause ? kPause : kTerm; + e.vblock = kind; + epoch_kind = kind; + } + __syncwarp(); + claim_full[k & 1].arrive(); // every writer lane arrives + break; + } + // Blocks are claimed heaviest (last rows, longest context) first. + const uint32_t vblock = nblocks - 1 - b, row_base = vblock * kBlockRows; + uint32_t st = 0, en = 0; + if (lane < kBlockRows) { + const uint32_t r = row_base + lane; + if (r < rows) { + st = min(starts_g[r], keys); + en = max(min(ends_g[r], keys), st); + } + e.starts[lane] = st; + e.ends[lane] = en; + } + const bool all_short = + __ballot_sync(kFullMask, lane < kBlockRows && en - st > topk) == 0; + const bool gated = + __ballot_sync(kFullMask, + lane < kBlockRows && en - st > sample_target) != 0; + // KV range: rows past the end are clamped to the last valid row (DeepGEMM + // scheduler). + const uint32_t src = min(row_base + (lane & 7), rows - 1) - row_base; + uint32_t lo = __shfl_sync(kFullMask, st, src), + hi = __shfl_sync(kFullMask, en, src); +#pragma unroll + for (int d = 4; d >= 1; d >>= 1) { + lo = min(lo, __shfl_xor_sync(kFullMask, lo, d)); + hi = max(hi, __shfl_xor_sync(kFullMask, hi, d)); + } + const uint32_t kv_base = lo / 4 * 4, + splits = (hi - kv_base + kTileKeys - 1) / kTileKeys; + // CTA-private slot pair: block n of this CTA uses slot 2 * cta + (n & 1); + // at most two blocks (one in selection, one claimed ahead) are in flight, + // so no claim atomics. + const uint32_t slot = kSlotsPerCta * blockIdx.x + (n_blk & 1); + if (!all_short && lane < kBlockRows) row_counts[row_base + lane] = 0; + if (lane == 0) { + e.vblock = vblock; + e.all_short = all_short; + e.kv_base = kv_base; + e.splits = splits; + e.samples = gated ? sample_tiles : 0u; + e.slot = slot; + } + __syncwarp(); + claim_full[k & 1].arrive(); // every writer lane arrives + if (all_short) continue; + if (lane == 0) { + dg::tma_copy(&tq, &full_q, sm.q, 0, + row_base * kNumHeads); + dg::tma_copy(&tw, &full_q, sm.weights, 0, + row_base); + full_q.arrive_and_expect_tx(kQBytes + kWBytes); + } + __syncwarp(); + ++n_blk; + const int64_t row_delta = slot_row_delta(slot, row_base); + for (uint32_t base = 0; base < splits; base += Interval, ++drain_interval) { + for (uint32_t p = kMathWarps / 2; p < kMathWarps; ++p) + drain_handoff(hdata, h_full, h_empty, p, drain_interval, row_base + 4, + row_counts, base + Interval >= splits, + [&](uint32_t row, uint32_t slot2, uint32_t record, + uint32_t batch_base, uint32_t source) { + store_ring_record(vals, indices, row, row_delta, slot2, + record, kv_base, batch_base, + source % kMathWgThreads); + }); + } + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); + } +} else if (warp_idx == 9) { + // -------- KV producer: sample tiles then scan tiles, one ring -------- + cutlass::arch::warpgroup_reg_dealloc(); + RingPipeline kv_pipe; + uint32_t n_blk = 0; + for (uint32_t k = 0;; ++k) { + ClaimHeader claim; + if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break; + if (claim.all_short) continue; + const uint32_t kv_base = claim.kv_base, splits = claim.splits, + samples = claim.samples; + { + // The whole warp stages the strided sample scales into sf_kv[stage] + // (prefetched one tile ahead in registers); lane 0 issues the KV TMA. + if (lane == 0 && n_blk > 0) sel_done.wait((n_blk - 1) & 1); + __syncwarp(); + const uint32_t center = sample_stride / 2; + // Per-key scales of one sample tile, kScalesPerLane per lane. + constexpr uint32_t kScalesPerLane = kTileKeys / 32; + float pf[kScalesPerLane]; +#pragma unroll + for (uint32_t j = 0; j < kScalesPerLane; ++j) { + const uint32_t g = j * 32 + lane; + pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center) + : 0.0f; + } +#pragma unroll 1 + for (uint32_t t = 0; t < samples; ++t) { + // Only lane 0's ring state is maintained across blocks (the scan loop + // runs on lane 0). + CUTE_TIE_DECL(kv_pipe.advance(), s0, ph); + const uint32_t s = __shfl_sync(kFullMask, s0, 0); + // Every writing lane acquires and releases the stage itself. + empty_kv[s].wait(__shfl_sync(kFullMask, ph, 0) ^ 1); +#pragma unroll + for (uint32_t j = 0; j < kScalesPerLane; ++j) + sm.sf_kv[s][j * 32 + lane] = pf[j]; + // Generic writes before later TMA writes of this stage. + cutlass::arch::fence_view_async_shared(); + if (lane == 0) { + full_kv[s].expect_transaction(kKVBytes); + dg::tma_copy( + &tks, &full_kv[s], sm.kv[s], 0, t * kTileKeys, center); + } + full_kv[s].arrive(); +#pragma unroll + for (uint32_t j = 0; j < kScalesPerLane; ++j) { + const uint32_t g = (t + 1) * kTileKeys + j * 32 + lane; + pf[j] = g < sample_groups ? __ldg(scales + g * sample_stride + center) + : 0.0f; + } + } + } + if (lane == 0) { +#pragma unroll 1 + for (uint32_t t = 0; t < splits; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), s, ph); + empty_kv[s].wait(ph ^ 1); + const uint32_t off = kv_base + t * kTileKeys; + dg::tma_copy(&tk, &full_kv[s], sm.kv[s], 0, + off); + dg::tma_copy(&ts, &full_kv[s], sm.sf_kv[s], off, 0); + mbar_arrive_expect_tx(&full_kv[s], kKVBytes + kScaleBytes, + kFullKvArrivals); + } + } + __syncwarp(); + ++n_blk; + } +} else if (warp_idx == 10) { + // -------- UMMA issuer -------- + cutlass::arch::warpgroup_reg_dealloc(); + RingPipeline<1> q_pipe; + RingPipeline kv_pipe; + RingPipeline tmem_pipe; + auto instr_desc = + cute::UMMA::make_instr_desc(); + auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + for (uint32_t k = 0;; ++k) { + ClaimHeader claim; + if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break; + if (claim.all_short) continue; + const uint32_t splits = claim.splits, samples = claim.samples; + CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase); + (void)q_stage; + full_q.wait(q_phase); + full_lin->wait(q_phase); // this block's linear-term rows (warp 11) + ptx::tcgen05_after_thread_sync(); + const uint32_t total = samples + splits; +#pragma unroll 1 + for (uint32_t t = 0; t < total; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), s, ph); + full_kv[s].wait(ph); + if (cute::elect_one_sync()) { + // Per KV half, one N=144 MMA per math warpgroup into its own TMEM slot + // (ring of 3; WG0 takes even, WG1 odd positions of the issue sequence). +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { +#pragma unroll + for (uint32_t g = 0; g < 2; ++g) { + CUTE_TIE_DECL(tmem_pipe.advance(), ts_idx, ts_ph); + const uint32_t tmem_addr = ts_idx * kSlotCols; + empty_tmem[ts_idx].wait(ts_ph ^ 1); + ptx::tcgen05_after_thread_sync(); + cutlass::float_e4m3_t* b_base = + g == 0 ? reinterpret_cast(sm.lin_lo) + : sm.q + 4 * kNumHeads * kHeadDim; // WG1's rows +#pragma unroll + for (uint32_t kk = 0; kk < 4; ++kk) { + auto a_desc = + mma::sm100::make_umma_desc( + sm.kv[s], i * 128, kk * 32); + auto b_desc = + mma::sm100::make_umma_desc( + b_base, 0, kk * 32); + ptx::SM100_MMA_F8F6F4_SS::fma(a_desc, b_desc, tmem_addr, kk, + runtime_instr_desc); + } + ptx::umma_arrive_no_elect(full_tmem[ts_idx]); + } + } + } + __syncwarp(); + } + empty_q.arrive(); + } +} else if (warp_idx == 11) { + // -------- linear-term rows + drain of WG0 rows -------- + cutlass::arch::warpgroup_reg_dealloc(); + uint32_t drain_interval = 0, n_lin = 0; + for (uint32_t k = 0;; ++k) { + ClaimHeader claim; + if (!consume_claim(k, claim_full, claim_empty, entries, claim)) break; + if (claim.all_short) continue; + const uint32_t vblock = claim.vblock, kv_base = claim.kv_base, + splits = claim.splits, slot = claim.slot; + // Q of block k lands while the math warps still select block k-1, so this + // is off the critical path except for a CTA's first block. + full_q.wait(n_lin & 1); + ++n_lin; + compute_lin_rows(sm, lane); + // Generic smem writes -> UMMA (async proxy). + cutlass::arch::fence_view_async_shared(); + full_lin->arrive(); // every lane arrives (count 32) + const uint32_t row_base = vblock * kBlockRows; + const int64_t row_delta = slot_row_delta(slot, row_base); + for (uint32_t base = 0; base < splits; base += Interval, ++drain_interval) { + for (uint32_t p = 0; p < kMathWarps / 2; ++p) + drain_handoff(hdata, h_full, h_empty, p, drain_interval, row_base, + row_counts, base + Interval >= splits, + [&](uint32_t row, uint32_t slot2, uint32_t record, + uint32_t batch_base, uint32_t source) { + store_ring_record(vals, indices, row, row_delta, slot2, + record, kv_base, batch_base, + source % kMathWgThreads); + }); + } + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).arrive(); + } +} else { + // -------- math warps 0-7 -------- + cutlass::arch::warpgroup_reg_alloc(); + const uint32_t math_tid = threadIdx.x; + const uint32_t query_base = (warp_idx / 4) * kWgRows; // 4 warps per WG + RingPipeline<1> q_pipe; + RingPipeline kv_pipe; + RingPipeline tmem_pipe; + // Each warpgroup consumes every second TMEM slot position (WG1 starts at 1). + constexpr uint32_t kTmemStep = 2; + if (warp_idx >= 4) tmem_pipe.advance(1); + const uint32_t heads_col = warp_idx < 4 ? 16u : 0u; + const uint32_t lin_col = warp_idx < 4 ? 0u : 128u; + uint32_t produce_interval = 0; + for (uint32_t k = 0;; ++k) { + claim_full[k & 1].wait((k >> 1) & 1); + const ClaimEntry& e = entries[k & 1]; + const uint32_t vblock = e.vblock; + if (vblock >= kPause) break; + const uint32_t row_base = vblock * kBlockRows; + if (e.all_short) { + for (uint32_t r = 0; r < kBlockRows && row_base + r < rows; ++r) + write_all_keys(out + (uint64_t)(row_base + r) * topk, topk, + e.ends[r] - e.starts[r], math_tid, kMathThreads); + __syncwarp(); + claim_empty[k & 1].arrive(); + continue; + } + // ---- block setup: Q, weights, sample histogram ---- + { + CUTE_TIE_DECL(q_pipe.advance(), q_stage, q_phase); + (void)q_stage; + full_q.wait(q_phase); + full_lin->wait(q_phase); + } + static_assert(sizeof(SampleShared) % 16 == 0); + for (uint32_t j = math_tid; j < sizeof(SampleShared) / 16; + j += kMathThreads) + reinterpret_cast(sample_smem.hist)[j] = make_uint4(0, 0, 0, 0); + float weights[kWgRows][kNumHeads]; + float row_s[kWgRows]; // pow2 row scale of the linear term + uint32_t seq_k_start[kWgRows], seq_k_end[kWgRows]; +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + seq_k_start[i] = e.starts[query_base + i]; + seq_k_end[i] = e.ends[query_base + i]; + // Weights in units of 1 / s_r (exact: s_r is a power of two). + row_s[i] = sm.lin_scale[query_base + i]; + const float inv_s = 1.0f / row_s[i]; +#pragma unroll + for (uint32_t j = 0; j < kNumHeads; j += 4) { + const float4 w4 = *reinterpret_cast( + sm.weights + (query_base + i) * kNumHeads + j); + weights[i][j] = w4.x * 0.5f * inv_s; + weights[i][j + 1] = w4.y * 0.5f * inv_s; + weights[i][j + 2] = w4.z * 0.5f * inv_s; + weights[i][j + 3] = w4.w * 0.5f * inv_s; + } + } + { + // Weights that are subnormal, or become subnormal or infinite once + // scaled, take the exact recovery path. + const float w = sm.weights[warp_idx * kNumHeads + lane]; // lane = head + // 2^-125: the smallest magnitude whose half is still a normal float. + constexpr float kMinScalable = 0x1p-125f; + bool bad_w = w != 0.0f && fabsf(w) < kMinScalable; + const float wr = w * 0.5f * (1.0f / sm.lin_scale[warp_idx]); + bad_w |= (wr != 0.0f && fabsf(wr) < kMinScalable) || isinf(wr); + const uint32_t bad = __ballot_sync(kFullMask, bad_w); + if (lane == 0 && bad) atomicOr(&unsafe_rows, 1u << warp_idx); + } + math_wg_sync(); + // ---- sample tiles ---- + // The first tile's per-row max is the histogram reference (see + // SampleShared): tile 0 keeps its scores until the reference is known. + float ref[kWgRows] = {}; +#pragma unroll 1 + for (uint32_t t = 0; t < e.samples; ++t) { + CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); + full_kv[kv_stage].wait(kv_phase); + uint32_t first_code[kWgRows] = {}; +#pragma unroll + for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) { + const uint32_t kv_lane = + (math_tid % kMathWgThreads) + kv_half * kMathWgThreads; + const uint32_t group = t * kTileKeys + kv_lane; + const uint32_t kv_offset = group * sample_stride + sample_stride / 2; + const bool in_view = group < sample_groups; + // 0 outside the view (staged by the KV producer). + const float scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + kv_lane); + CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase); + full_tmem[tmem_stage].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + if (kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + empty_kv[kv_stage].arrive(); + } + // One query row (32 TMEM columns) at a time keeps the sample loop's + // live registers low enough to carry the scan weights through it. +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + float acc[kNumHeads]; + float lin[4]; + ptx::tmem_load_32dp32b<32>( + tmem_stage * kSlotCols + heads_col + i * kNumHeads, + reinterpret_cast(acc)); + ptx::tmem_load_32dp32b<4>(tmem_stage * kSlotCols + lin_col + i * 4, + reinterpret_cast(lin)); + cutlass::arch::fence_view_async_tmem_load(); + if (i == kWgRows - 1) { + ptx::tcgen05_before_thread_sync(); + empty_tmem[tmem_stage].arrive(); + } + const float bq = reduce_heads_with_lin(acc, weights[i], lin) * + (row_s[i] * scale_kv); + const bool valid = in_view && kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq); + if (t == 0) { + sample_smem.first[kv_half][i][math_tid] = + valid ? bq : -CUDART_INF_F; + if (valid) first_code[i] = max(first_code[i], ordered_code(bq)); + } else if (valid) + add_sample(sample_smem.hist[query_base + i], bq - ref[i]); + } + } + if (t == 0) { +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + const uint32_t code = __reduce_max_sync(kFullMask, first_code[i]); + if (lane == 0) atomicMax(&sample_smem.ref_code[query_base + i], code); + } + math_wg_sync(); +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + ref[i] = sample_ref(sample_smem.ref_code[query_base + i]); +#pragma unroll + for (uint32_t h = 0; h < 2; ++h) { + const float x = sample_smem.first[h][i][math_tid]; + if (isfinite(x)) + add_sample(sample_smem.hist[query_base + i], x - ref[i]); + } + } + } + } + math_wg_sync(); + finish_sample_histogram( + sample_smem.hist[warp_idx], sample_ref(sample_smem.ref_code[warp_idx]), + gate_s + warp_idx, e.ends[warp_idx] - e.starts[warp_idx], sample_target, + lane); + math_wg_sync(); + // The block index is re-read from the claim entry (volatile) instead of + // kept live across the sample loop. + const uint32_t scan_row_base = volatile_load(e.vblock) * kBlockRows; + + // ---- scan tiles ---- + const uint32_t kv_base = e.kv_base, splits = e.splits; + const int64_t cand_row_delta = slot_row_delta(e.slot, scan_row_base); + float gate_score[kWgRows]; +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + gate_score[i] = gate_s[query_base + i]; + if ((volatile_load(unsafe_rows) >> (query_base + i)) & 1u) { + // Unsafe weights: no key passes (+inf gate), and a count above the + // candidate cap sends the block to the exact rescan. + gate_score[i] = CUDART_INF_F; + if (math_tid % kMathWgThreads == 0) + atomicAdd(row_counts + scan_row_base + query_base + i, + static_cast(kCandidatesPerRow) + 1); + } + // Compare S * scale_kv against gate / row_s (row_s is a power of two -> + // exact). + gate_score[i] *= 1.0f / row_s[i]; + } + uint32_t lane_counts[kWgRows] = {}; + uint32_t flush_base = 0; +#pragma unroll 1 + for (uint32_t kv_split_idx = 0; kv_split_idx < splits; ++kv_split_idx) { + if ((kv_split_idx & (Interval - 1)) == 0) { + h_empty[(produce_interval & 1) * kMathWarps + warp_idx].wait( + ((produce_interval >> 1) & 1) ^ 1); + __syncwarp(); + } + CUTE_TIE_DECL(kv_pipe.advance(), kv_stage, kv_phase); + full_kv[kv_stage].wait(kv_phase); +#pragma unroll + for (uint32_t kv_half = 0; kv_half < 2; ++kv_half) { + const uint32_t kv_lane = + (math_tid % kMathWgThreads) + kv_half * kMathWgThreads; + const uint32_t kv_offset = kv_base + kv_split_idx * kTileKeys + kv_lane; + const float scale_kv = ptx::ld_shared(sm.sf_kv[kv_stage] + kv_lane); + CUTE_TIE_DECL(tmem_pipe.advance(kTmemStep), tmem_stage, tmem_phase); + full_tmem[tmem_stage].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + if (kv_half == 1) { + cutlass::arch::fence_view_async_shared(); + empty_kv[kv_stage].arrive(); + } + float accum_pair[2][kNumHeads]; + float lin_pair[2][4]; + float + bp4[kWgRows]; // row scores of the slot, in units of the row scale +#pragma unroll + for (uint32_t pair = 0; pair < kWgRows; pair += 2) { + const uint32_t slot_base = tmem_stage * kSlotCols; + ptx::tmem_load_32dp32b<2 * kNumHeads>( + slot_base + heads_col + pair * kNumHeads, + reinterpret_cast(accum_pair)); + ptx::tmem_load_32dp32b<8>(slot_base + lin_col + pair * 4, + reinterpret_cast(lin_pair)); + cutlass::arch::fence_view_async_tmem_load(); + if (pair + 2 >= kWgRows) { + ptx::tcgen05_before_thread_sync(); + empty_tmem[tmem_stage].arrive(); + } + // The linear term (sum_t term_t . k, in units of the row scale; terms + // 2, 3 in units of s * 2^-8) seeds the sums. +#pragma unroll + for (uint32_t r = 0; r < 2; ++r) + bp4[pair + r] = reduce_heads_with_lin( + accum_pair[r], weights[pair + r], lin_pair[r]) * + scale_kv; + } + // One warp vote per slot: most keys fail every gate, so the common path + // is 4 compares + 1 vote + 1 uniform branch. + const bool any_pass = + (bp4[0] > gate_score[0]) | (bp4[1] > gate_score[1]) | + (bp4[2] > gate_score[2]) | (bp4[3] > gate_score[3]); + if (__any_sync(kFullMask, any_pass)) [[unlikely]] { +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + const float bp = bp4[i]; + // Gate first; range and finiteness only for passing keys. + if (bp > gate_score[i]) { + const float bq = bp * row_s[i]; + const uint32_t row = scan_row_base + query_base + i; + if (kv_offset >= seq_k_start[i] && kv_offset < seq_k_end[i] && + isfinite(bq)) { + const uint32_t count = lane_counts[i]; + if (count < kRingSlots) { + const uint32_t code = candidate_fp24_code(bq); + hdata + .records[produce_interval & 1][warp_idx][i][count][lane] = + (code << 8) | ((kv_split_idx - flush_base) * 2 + kv_half); + lane_counts[i] = count + 1; + } else { + const int slot = atomicAdd(row_counts + row, 1); + if (slot < static_cast(kCandidatesPerRow)) { + const uint64_t off = + candidate_offset(row, cand_row_delta, slot); + store_candidate(vals + off, indices + off, bq, kv_offset); + } + } + } + } + } + } + } + if (((kv_split_idx + 1) & (Interval - 1)) == 0 && + kv_split_idx + 1 != splits) { +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + hdata.num_records[produce_interval & 1][warp_idx][i][lane] = + lane_counts[i]; + lane_counts[i] = 0; + } + if (lane == 0) hdata.bases[produce_interval & 1][warp_idx] = flush_base; + __syncwarp(); + h_full[(produce_interval & 1) * kMathWarps + warp_idx].arrive(); + ++produce_interval; + flush_base = kv_split_idx + 1; + } + } + { + // Final interval: every math warp stores its own ring records (8 warps in + // parallel instead of 2 drain warps after the scan); the drain warps only + // cycle the barriers. + const uint32_t bank = produce_interval & 1, + rbase = scan_row_base + query_base; + static_assert(sizeof(hdata.records[0][0]) == + kWgRows * kRingSlots * 32 * sizeof(uint32_t), + "ring record layout"); + self_drain_last(&hdata.records[bank][warp_idx][0][0][0], lane_counts[0], + lane_counts[1], lane_counts[2], lane_counts[3], + row_counts, rbase, kv_base + (math_tid % kMathWgThreads), + flush_base, cand_row_delta, vals, indices); +#pragma unroll + for (uint32_t row = 0; row < kWgRows; ++row) lane_counts[row] = 0; + __syncwarp(); + h_full[bank * kMathWarps + warp_idx].arrive(); + ++produce_interval; + } + cutlass::arch::fence_view_async_shared(); + empty_q.arrive(); + // Drains of this block done. + cutlass::arch::NamedBarrier(kDrainSyncThreads, kBarDrain).sync(); + // ---- selection ---- + // Re-read as well (not kept live across the scan). + const uint32_t sel_vblock = volatile_load(e.vblock), + sel_row_base = sel_vblock * kBlockRows; + __threadfence_block(); + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + const ClaimEntry& sel_entry = entries[k & 1]; + if (math_tid == 0) { + uint32_t recover = 0, short_rows = 0; + row_offset[0] = 0; + for (uint32_t r = 0; r < kBlockRows; ++r) { + const bool is_short = sel_entry.ends[r] - sel_entry.starts[r] <= topk; + short_rows |= uint32_t(is_short) << r; + const int c = __ldcg(row_counts + sel_row_base + r); + recover |= !is_short && (c < int(topk) || c > int(kCandidatesPerRow)); + // Rows start 16-byte aligned for select_staged_row's vector loads. + row_offset[r + 1] = + row_offset[r] + (is_short ? 0 : (uint32_t(max(c, 0)) + 3u) & ~3u); + } + // 32-bit byte arithmetic (sizeof is a 64-bit size_t). + constexpr uint32_t kWordBytes = sizeof(uint32_t); + // Stage all 8 rows' candidates at once (1), in two rounds of 4 (2), or, + // when they exceed the scratch, rescan exactly. + const uint32_t half = max(row_offset[kWgRows], + row_offset[kBlockRows] - row_offset[kWgRows]); + const uint32_t staged = + recover ? 0u + : ((kBlockRows * selection::kBins + row_offset[kBlockRows]) * + kWordBytes <= + phase_bytes + ? 1u + : ((kWgRows * selection::kBins + half) * kWordBytes <= + phase_bytes + ? 2u + : 0u)); + if (staged == 0) recover = 1; + if (recover) { + defer_next[sel_vblock] = int(defer_head); + defer_head = sel_vblock; + atomicExch(&pause_req, 1u); + } + sel_recover = recover; + sel_staged = staged; + sel_short = short_rows; + unsafe_rows = 0; + } + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + const uint32_t recover = sel_recover, short_rows = sel_short; + const auto write_short = [&](uint32_t r, uint32_t t, uint32_t stride) { + if (r >= rows) return; + write_all_keys( + out + (uint64_t)r * topk, topk, + sel_entry.ends[r - sel_row_base] - sel_entry.starts[r - sel_row_base], + t, stride); + }; + if (!recover) { + const int64_t row_delta = slot_row_delta(sel_entry.slot, sel_row_base); + const uint32_t per_round = sel_staged == 1 ? kBlockRows : kWgRows; + for (uint32_t round = 0; round < kBlockRows / per_round; ++round) { + if (warp_idx < per_round) { + const uint32_t local = round * per_round + warp_idx, + first = round * per_round; + uint32_t* hist = reinterpret_cast(smem_buffer + kSelBase) + + warp_idx * selection::kBins; + uint32_t* codes = + reinterpret_cast(smem_buffer + kSelBase) + + per_round * selection::kBins + + (row_offset[local] - row_offset[first]); + const uint32_t r = sel_row_base + local; + const uint64_t base = candidate_offset(r, row_delta, 0); + if ((short_rows >> local) & 1u) + write_short(r, lane, 32); + else { + const uint32_t cnt = __ldcg(row_counts + r); + int* o = out + (uint64_t)r * topk; + if (!selection::select_staged_row(codes, cnt, vals + base, + indices + base, o, topk, hist)) + selection::select_row(codes, cnt, vals + base, indices + base, o, + topk, hist); + } + } + cutlass::arch::NamedBarrier(kMathThreads, kBarSelect).sync(); + } + } + // Every row is final here: the last selection round ends with a barrier, + // and a recovered block skipped the rounds right after the barrier above. + // Short rows were written relative already. + if (!recover) + make_relative(out, topk, rows, sel_row_base, starts_g, math_tid, + kMathThreads, [&](uint32_t local) { + return sel_entry.ends[local] - sel_entry.starts[local] <= + topk; + }); + // The selection scratch aliases kv[] and sf_kv[]. Order its generic + // accesses before the next block's TMA writes, which have no proxy fence + // of their own when that block has no sample tiles. + cutlass::arch::fence_view_async_shared(); + __syncwarp(); + if (lane == 0) sel_done.arrive(); + claim_empty[k & 1].arrive(); + } +} + +// -------- exact rescan of the blocks queued by the selection -------- +// The rescan's first TMA loads overwrite the epoch's generic shared memory. +cutlass::arch::fence_view_async_shared(); +__syncthreads(); +if (warp_idx == 0) cute::TMEM::Allocator1Sm().free(0, 512); +if (defer_head != kTerm) { + // The rescan expects the 224 / 56 register split of the math warps. + __syncthreads(); + for (uint32_t cur = defer_head; cur != kTerm;) { + const uint32_t vblock = cur, row_base = vblock * kBlockRows; + if (threadIdx.x < kBlockRows) { + const uint32_t r = row_base + threadIdx.x; + row_bounds[0][threadIdx.x] = r < rows ? min(starts_g[r], keys) : 0u; + row_bounds[1][threadIdx.x] = + r < rows ? max(min(ends_g[r], keys), row_bounds[0][threadIdx.x]) : 0u; + } + __syncthreads(); + const uint32_t* const starts = row_bounds[0] - row_base; + const uint32_t* const ends = row_bounds[1] - row_base; + // No block is in flight during the rescan: use the CTA's first slot. + const int64_t row_delta = + slot_row_delta(kSlotsPerCta * blockIdx.x, row_base); + const RescanBuffers block_buffers{ + rescan.retained + row_delta * kRetainedPerRow, + rescan.retained_indices + row_delta * kRetainedPerRow}; + rescan::rescan_block(dense::Scheduler(vblock, rows, keys, starts, ends), + block_buffers, out, topk, tq, tk, ts, tw); + __syncthreads(); + if (threadIdx.x == 0) + dense::invalidate_barriers( + *reinterpret_cast(smem_buffer)); + __syncthreads(); + // The rescan writes absolute indices for every row of the block. + make_relative(out, topk, rows, row_base, starts_g, threadIdx.x, kNumThreads, + [](uint32_t) { return false; }); + cur = uint32_t(__ldcg(defer_next + vblock)); + // Order the generic mbarrier.inval above before the TMA writes of the + // next rescan block or the next epoch. + cutlass::arch::fence_view_async_shared(); + __syncthreads(); + } + // The loop ran at least once and ended with a barrier. + if (threadIdx.x == 0) { + defer_head = kTerm; + pause_req = 0; + } +} // deferred blocks +more = epoch_kind != kTerm; +__syncthreads(); diff --git a/csrc/fused_mqa_topk/persistent_kernel.cuh b/csrc/fused_mqa_topk/persistent_kernel.cuh new file mode 100644 index 000000000000..86063b08fbac --- /dev/null +++ b/csrc/fused_mqa_topk/persistent_kernel.cuh @@ -0,0 +1,592 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Persistent fused FP8 MQA logits + per-row top-K (K = 512, 1024, 2048; SM100), +// DeepGEMM/FA4-style: one CTA per SM, barriers and TMEM initialised once, every +// warp role runs its own loop over blocks claimed (heaviest first) from a +// global counter, and the sample and scan tiles of a block flow through ONE +// continuous Q/KV/TMEM pipeline whose stage/phase state is carried across +// blocks. +// +// Roles (384 threads): +// warp 8 claimer + Q/weights TMA + drain of math WG1 rows +// warp 9 KV TMA (sample tiles via the strided map, then the scan tiles) +// warp 10 UMMA issuer +// warp 11 linear-term rows + drain of math WG0 rows +// warps 0-7 math: sample histogram -> gate -> scan/emit -> selection +// A block is rescanned exactly (rescan.cuh) when a row's candidate count +// leaves [K, kCandidatesPerRow], a row's weights are unsafe to scale, or the +// block's candidates exceed the selection scratch (~11K+ per row): the role +// loops pause, the CTA rescans its queued blocks, and the loops resume with +// freshly initialised barriers. +#pragma once +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "candidates.cuh" +#include "common.cuh" +#include "constants.h" +#include "deep_gemm_compat.cuh" +#include "rescan.cuh" +#include "select.cuh" + +namespace fused_mqa_topk::persist { +using namespace deep_gemm; +using rescan::RescanBuffers; +using rescan::RescanShared; +using TBar = cutlass::arch::ClusterTransactionBarrier; + +constexpr uint32_t kTerm = 0xffffffffu, kPause = 0xfffffffeu; +// Bytes of a block's Q, of one KV tile and its per-key scales, of a block's +// weights. +constexpr uint32_t kQBytes = kBlockRows * kNumHeads * kHeadDim, + kKVBytes = kTileKeys * kHeadDim, + kScaleBytes = kTileKeys * sizeof(float), + kWBytes = kBlockRows * kNumHeads * sizeof(float); +// KV (and per-key scale) pipeline stages. +constexpr uint32_t kKVStages = 3; +// Linear-term B rows per math warpgroup: 4 query rows x 4 FP8 terms. +constexpr uint32_t kLinRows = 16; +// Terms 2 and 3 of the linear term are stored in units of s * 2^-8. +constexpr float kLinLoScale = 0x1p-8f; +// TMEM: one UMMA per math warpgroup and KV half writes 4 rows x 32 heads plus +// the kLinRows linear-term columns (N = 144). Slots are 160 columns apart (144 +// rounded up to a multiple of 32); the 3 slots take 480 of the 512 allocated +// columns. +constexpr uint32_t kUmmaN = 4 * kNumHeads + kLinRows, kSlotCols = 160; +static_assert(kSlotCols >= kUmmaN && kSlotCols % 32 == 0); + +// Rebalanced scan: the linear-term rows of each math warpgroup bracket Q, so +// WG0's N=144 UMMA reads B = [lin WG0 | Q rows 0..127] (D = [lin | heads]) and +// WG1's reads B = [Q rows 128..255 | lin WG1] (D = [heads | lin]). +struct KernelSmem { + alignas(1024) uint8_t lin_lo[kLinRows * kHeadDim]; + cutlass::float_e4m3_t q[kQBytes]; + uint8_t lin_hi[kLinRows * kHeadDim]; + alignas(128) float weights[kBlockRows * kNumHeads]; + float lin_scale[kBlockRows]; + alignas(1024) cutlass::float_e4m3_t kv[kKVStages][kKVBytes]; + alignas(128) float sf_kv[kKVStages][kTileKeys]; +}; +static_assert(offsetof(KernelSmem, q) == kLinRows * kHeadDim && + offsetof(KernelSmem, lin_hi) == kLinRows * kHeadDim + kQBytes, + "lin rows must bracket Q"); +// Selection scratch starts after lin + Q + lin + weights + row scales, so the +// next block's Q and linear term land during this block's selection. +constexpr uint32_t kSelBase = offsetof(KernelSmem, kv); + +// Byte offset of byte `col` of row `row` in a 128B-swizzled K-major tile of +// 128-byte rows (the UMMA B layout of Q and the linear-term rows). +__device__ __forceinline__ uint32_t swizzle128_offset(uint32_t row, + uint32_t col) { + return row * 128u + (((col >> 4) ^ (row & 7u)) << 4) + (col & 15u); +} + +// q'_r = sum_h (w_rh / 2) q_rh, split into 4 FP8 e4m3 terms: terms 0,1 scale +// s_r, terms 2,3 s_r * 2^-8 (lin_scale[r]); term t of row r is B row (r % 4) * +// 4 + t of lin_lo (rows 0-3) / lin_hi (4-7), 128B-swizzled K-major like Q. One +// warp; lane l owns dims 4l..4l+3. +__device__ __forceinline__ void compute_lin_rows(KernelSmem& sm, + uint32_t lane) { + const uint8_t* qbytes = reinterpret_cast(sm.q); +#pragma unroll 1 + for (uint32_t r = 0; r < kBlockRows; ++r) { + float acc[4] = {0.f, 0.f, 0.f, 0.f}; +#pragma unroll 8 + for (uint32_t h = 0; h < kNumHeads; ++h) { + const uint32_t off = swizzle128_offset(r * kNumHeads + h, lane * 4u); + const uint32_t packed = + ptx::ld_shared(reinterpret_cast(qbytes + off)); + const float wv = ptx::ld_shared(sm.weights + r * kNumHeads + h) * 0.5f; + const __half2_raw lo = __nv_cvt_fp8x2_to_halfraw2( + static_cast<__nv_fp8x2_storage_t>(packed & 0xffffu), __NV_E4M3); + const __half2_raw hi = __nv_cvt_fp8x2_to_halfraw2( + static_cast<__nv_fp8x2_storage_t>(packed >> 16), __NV_E4M3); + const float2 f0 = __half22float2(*reinterpret_cast(&lo)); + const float2 f1 = __half22float2(*reinterpret_cast(&hi)); + acc[0] = fmaf(wv, f0.x, acc[0]); + acc[1] = fmaf(wv, f0.y, acc[1]); + acc[2] = fmaf(wv, f1.x, acc[2]); + acc[3] = fmaf(wv, f1.y, acc[3]); + } + float m = fmaxf(fmaxf(fabsf(acc[0]), fabsf(acc[1])), + fmaxf(fabsf(acc[2]), fabsf(acc[3]))); +#pragma unroll + for (uint32_t o = 16; o > 0; o >>= 1) + m = fmaxf(m, __shfl_xor_sync(kFullMask, m, o)); + int e = -120; // pow2 scale s >= m / 448, clamped so 1/s stays finite + if (m > 0.f && isfinite(m)) { + frexpf(m * (1.0f / 448.0f), &e); + e = max(e, -120); + } + const float s = ldexpf(1.0f, e), inv_s = ldexpf(1.0f, -e); + float res[4]; +#pragma unroll + for (uint32_t dd = 0; dd < 4; ++dd) res[dd] = acc[dd] * inv_s; + uint8_t* lin_base = r < kWgRows ? sm.lin_lo : sm.lin_hi; +#pragma unroll + for (uint32_t tt = 0; tt < 4; ++tt) { + const __nv_fp8x2_storage_t c01 = __nv_cvt_float2_to_fp8x2( + make_float2(res[0], res[1]), __NV_SATFINITE, __NV_E4M3); + const __nv_fp8x2_storage_t c23 = __nv_cvt_float2_to_fp8x2( + make_float2(res[2], res[3]), __NV_SATFINITE, __NV_E4M3); + const __half2_raw d01 = __nv_cvt_fp8x2_to_halfraw2(c01, __NV_E4M3); + const __half2_raw d23 = __nv_cvt_fp8x2_to_halfraw2(c23, __NV_E4M3); + const float2 g01 = + __half22float2(*reinterpret_cast(&d01)); + const float2 g23 = + __half22float2(*reinterpret_cast(&d23)); + res[0] -= g01.x; + res[1] -= g01.y; + res[2] -= g23.x; + res[3] -= g23.y; + // Terms 2, 3 in units of s * 2^-8: the residual after two terms is <= + // 0.5, so x256 <= 128 (no saturation), and the e4m3 subnormal floor drops + // from s * 2^-10 to s * 2^-18 per dim. + if (tt == 1) { + res[0] *= 1.0f / kLinLoScale; + res[1] *= 1.0f / kLinLoScale; + res[2] *= 1.0f / kLinLoScale; + res[3] *= 1.0f / kLinLoScale; + } + const uint32_t off = + swizzle128_offset((r % kWgRows) * 4u + tt, lane * 4u); + ptx::st_shared( + reinterpret_cast(lin_base + off), + static_cast(c01) | (static_cast(c23) << 16)); + } + if (lane == 0) sm.lin_scale[r] = s; + } +} + +// One row's weighted head reduction seeded with its linear-term columns +// (terms 2, 3 in units of s * 2^-8). +__device__ __forceinline__ float reduce_heads_with_lin(const float* s, + const float* w, + const float* lin) { + return reduce_heads(s, w, make_float2(lin[0], lin[1]), + __fmul2_rn(make_float2(lin[2], lin[3]), + make_float2(kLinLoScale, kLinLoScale))); +} + +// Gate sample state: each math warp histograms its row's sampled scores +// relative to the row's reference `ref` (its max over the first sample tile): +// half_bin<5>(x - ref), kSampleBins bins, 32 per octave of the distance, so +// the bins near the row's top stay fine whatever the scores' common offset +// (absolute bins are ~2 wide at |x| ~ 100, too coarse for top tails that +// narrow). +constexpr uint32_t kSampleShift = 5, kSampleBins = 65536 >> kSampleShift; +struct alignas(16) SampleShared { + uint32_t hist[kBlockRows][padded_bins(kSampleBins)]; + // Ordered code of each row's max over the first sample tile (0: none). + uint32_t ref_code[kBlockRows]; + // The first tile's scores, held until the reference is known (in registers + // they cost the hot loops spills). + float first[2][kWgRows][kMathThreads]; +}; + +__device__ __forceinline__ void add_sample(uint32_t* hist, float distance) { + atomicAdd(hist + padded_bin(half_bin(distance)), 1u); +} +__device__ __forceinline__ uint32_t sample_count(const uint32_t* hist, + uint32_t bin) { + return hist[padded_bin(bin)]; +} +// The row reference from its ordered code; 0 when the row had no valid sample +// in the first tile (then the bins are absolute). +__device__ __forceinline__ float sample_ref(uint32_t code) { + const float ref = ordered_float(code); + return isfinite(ref) ? ref : 0.0f; +} + +// The gate aims kGateSlack16 / 16 above the target count, so that few rows fall +// below the top-k (and into the exact rescan) when the sampled estimate runs +// low. On the B300 captures 17/16 cuts the rescanned blocks by ~85% at about +// the candidate counts of the coarser absolute bins. +constexpr uint32_t kGateSlack16 = 17; + +// Gate: the k-th largest sample (k = count * target / visible, count = the +// histogram total) of one row from its histogram; one warp per row. -inf when +// the row sees at most `target` keys or has no samples. +__device__ __noinline__ inline void finish_sample_histogram( + const uint32_t* hist, float ref, float* gate, uint32_t visible, + uint32_t target, uint32_t lane) { + if (visible <= target) { + if (lane == 0) *gate = -CUDART_INF_F; + __syncwarp(); + return; + } + // Lane l owns kChunksPerLane 32-bin chunks, descending with the lane: chunk + // kChunks - 1 - kChunksPerLane * l - j for j < kChunksPerLane. + constexpr uint32_t kChunks = kSampleBins / 32, kChunksPerLane = kChunks / 32; + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < kChunksPerLane; ++j) { + const uint32_t c = kChunks - 1 - kChunksPerLane * lane - j; +#pragma unroll 8 + for (uint32_t b = 0; b < 32; ++b) sum += hist[padded_bin(c * 32 + b)]; + } + const uint32_t incl = warp_inclusive_scan(sum, lane); + const uint32_t tot = __shfl_sync(kFullMask, incl, 31); + if (tot == 0) { + if (lane == 0) *gate = -CUDART_INF_F; + __syncwarp(); + return; + } + const uint32_t rank = max( + 1u, min(tot, (tot * target * kGateSlack16 / 16 + visible - 1) / visible)); + const uint32_t w = __ffs(__ballot_sync(kFullMask, incl >= rank)) - 1; + uint32_t above = __shfl_sync(kFullMask, incl - sum, w), chosen = 0; +#pragma unroll 1 + for (uint32_t cc = 0; cc < kChunksPerLane; ++cc) { + const uint32_t c = kChunks - 1 - kChunksPerLane * w - cc; + // Lane 0 = highest bin of the chunk. + const uint32_t v = + warp_inclusive_scan(sample_count(hist, c * 32 + 31 - lane), lane); + const uint32_t mask = __ballot_sync(kFullMask, above + v >= rank); + if (mask) { + chosen = c * 32 + 31 - (__ffs(mask) - 1); + break; + } + above += __shfl_sync(kFullMask, v, 31); + } + if (lane == 0) *gate = ref + half_bin_floor(chosen); + __syncwarp(); +} + +// Candidate records each lane of a math warp buffers per row and handoff +// interval before it falls back to direct global stores. +constexpr uint32_t kRingSlots = 8; +// The drain's prefix sum packs two 16-bit per-row counts into one word. +static_assert(32 * kRingSlots < 65536, + "per-warp record count must fit 16 bits"); + +// Handoff ring: [bank][math warp][row][slot][lane] records, per-lane record +// counts and the interval's first scan tile. +struct alignas(16) HandoffRing { + uint32_t records[2][kMathWarps][kWgRows][kRingSlots][32]; + uint32_t num_records[2][kMathWarps][kWgRows][32]; + uint32_t bases[2][kMathWarps]; +}; + +// samples: gate sample tiles of the block (0 when no row sees more keys than +// the gate target, so every gate is -inf). +struct ClaimHeader { + uint32_t vblock, all_short, kv_base, splits, slot, samples; +}; +struct ClaimEntry : ClaimHeader { + uint32_t starts[kBlockRows], ends[kBlockRows]; +}; + +// Claim k in a producer role (warps 9-11): waits for the entry, copies its +// header and releases the entry (every lane arrives). False once the claimer +// ends the epoch (kTerm / kPause). +__device__ __forceinline__ bool consume_claim(uint32_t k, TBar* claim_full, + TBar* claim_empty, + const ClaimEntry* entries, + ClaimHeader& claim) { + claim_full[k & 1].wait((k >> 1) & 1); + claim = entries[k & 1]; + __syncwarp(); + if (claim.vblock >= kPause) return false; + claim_empty[k & 1].arrive(); + return true; +} + +// Barrier words (static shared, never aliased with data). +struct BarIdx { + // TMEM slots, shared by the two math warpgroups (WG0 even, WG1 odd + // positions). + static constexpr uint32_t kTmemSlots = 3; + static constexpr uint32_t full_q = 0, empty_q = 1, full_kv = 2, + empty_kv = 2 + kKVStages, + full_tmem = 2 + 2 * kKVStages, + empty_tmem = full_tmem + kTmemSlots, + claim_full = empty_tmem + kTmemSlots, + claim_empty = claim_full + 2, + sel_done = claim_empty + 2, h_full = sel_done + 1, + h_empty = h_full + 2 * kMathWarps, + full_lin = h_empty + 2 * kMathWarps, + total = full_lin + 1; +}; + +// Expect `bytes` of TMA transactions and arrive `count` times on one barrier. +__device__ __forceinline__ void mbar_arrive_expect_tx(const void* b, + uint32_t bytes, + uint32_t count) { + const uint32_t a = static_cast(__cvta_generic_to_shared(b)); + asm volatile( + "mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;" ::"r"(a), + "r"(bytes) + : "memory"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;" ::"r"(a), + "r"(count) + : "memory"); +} + +// Candidate slabs: a block's 8 rows live in its CTA slot's 8 slab rows, so row +// r maps to slab row r + slot_row_delta(slot, row_base). +__device__ __forceinline__ int64_t slot_row_delta(uint32_t slot, + uint32_t row_base) { + return int64_t(slot) * kBlockRows - int64_t(row_base); +} +// Offset of entry `index` of row `row` in the candidate slabs. +__device__ __forceinline__ uint64_t candidate_offset(uint32_t row, + int64_t row_delta, + uint32_t index) { + return static_cast(int64_t(row) + row_delta) * kCandidatesPerRow + + index; +} +// Stores ring record `record` (fp24 code << 8 | tile offset * 2 + KV half, +// tile offset relative to the interval's first tile `interval_base`) of KV +// lane `kv_lane` as entry `index` of row `row`. +__device__ __forceinline__ void store_ring_record( + uint16_t* vals, int32_t* indices, uint32_t row, int64_t row_delta, + uint32_t index, uint32_t record, uint32_t kv_base, uint32_t interval_base, + uint32_t kv_lane) { + const uint32_t key = + kv_base + (interval_base * 2 + (record & 255u)) * 128 + kv_lane; + const uint64_t off = candidate_offset(row, row_delta, index); + store_candidate_record(vals + off, indices + off, record >> 8, key); +} + +// Output row `out` of a row that sees at most top-k keys: its `visible` keys +// in order (0 .. visible - 1, relative to the row start), then -1 padding. +// Threads t, t + stride, ... write the slots. +__device__ __forceinline__ void write_all_keys(int* out, uint32_t topk, + uint32_t visible, uint32_t t, + uint32_t stride) { + for (uint32_t j = t; j < topk; j += stride) + out[j] = j < visible ? int(j) : -1; +} + +// Rewrites the output rows of the block at row_base in place from absolute key +// indices to indices relative to the row's start (the top_k_per_row_prefill +// convention; -1 padding stays). Rows that start at 0 and rows for which +// skip(local row) holds are left as they are. Threads t, t + stride, ... each +// rewrite 4 slots. +template +__device__ __forceinline__ void make_relative(int* output, uint32_t topk, + uint32_t rows, uint32_t row_base, + const uint32_t* starts_g, + uint32_t t, uint32_t stride, + Skip skip) { + const uint32_t quads = topk / 4; + for (uint32_t i = t; i < kBlockRows * quads; i += stride) { + const uint32_t r = row_base + i / quads; + if (r >= rows || skip(i / quads)) continue; + const int s = int(starts_g[r]); + if (s == 0) continue; + int4* p = reinterpret_cast(output + (uint64_t)r * topk) + i % quads; + int4 v = *p; + v.x = v.x < 0 ? -1 : v.x - s; + v.y = v.y < 0 ? -1 : v.y - s; + v.z = v.z < 0 ? -1 : v.z - s; + v.w = v.w < 0 ? -1 : v.w - s; + *p = v; + } +} + +// Stores one warp's ring records of a handoff interval in its rows' candidate +// slabs. recs = [4 rows][kRingSlots records][32 lanes]; lane l holds n[row] +// records of each row. Each row's slots are reserved with one atomic on +// row_counts[row_base + row]; records past kCandidatesPerRow are dropped (the +// row's count then exceeds the cap, so its block is rescanned exactly). +// store(row, slot, record) writes one record. +// Written for a warpgroup's 4 rows (two 16-bit counts per word in the scan). +static_assert(kWgRows == 4, "store_ring_records packs 4 row counts"); +template +__device__ __forceinline__ void store_ring_records(const uint32_t* recs, + const uint32_t (&n)[4], + int* row_counts, + uint32_t row_base, + Store store) { + const uint32_t lane = threadIdx.x & 31; + uint32_t my_total = 0; +#pragma unroll + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t total = __reduce_add_sync(kFullMask, n[row]); + if (lane == row) my_total = total; + } + uint32_t my_start = 0; + if (lane < 4 && my_total) + my_start = atomicAdd(row_counts + row_base + lane, my_total); + // Two 16-bit per-row counts per word: one scan covers two rows. + uint32_t p01 = n[0] | (n[1] << 16), p23 = n[2] | (n[3] << 16); +#pragma unroll + for (int dd = 1; dd < 32; dd *= 2) { + const uint32_t a = __shfl_up_sync(kFullMask, p01, dd), + b = __shfl_up_sync(kFullMask, p23, dd); + if (lane >= dd) { + p01 += a; + p23 += b; + } + } + const uint32_t prefix[4] = {p01 & 65535u, p01 >> 16, p23 & 65535u, p23 >> 16}; +#pragma unroll + for (uint32_t row = 0; row < 4; ++row) { + const uint32_t start = __shfl_sync(kFullMask, my_start, row); + for (uint32_t j = 0; j < n[row]; ++j) { + const uint32_t slot = start + prefix[row] - n[row] + j; + const uint32_t record = recs[(row * kRingSlots + j) * 32 + lane]; + if (slot < kCandidatesPerRow) store(row_base + row, slot, record); + } + } +} + +// Self-drain, out of line: stores this warp's ring records of the last +// handoff interval. recs = records[bank][warp]. +__device__ __noinline__ inline void self_drain_last( + const uint32_t* __restrict__ recs, uint32_t n0, uint32_t n1, uint32_t n2, + uint32_t n3, int* __restrict__ row_counts, uint32_t rbase, uint32_t lane_kv, + uint32_t flush_base, int64_t row_delta, uint16_t* __restrict__ vals, + int32_t* __restrict__ indices) { + const uint32_t n[4] = {n0, n1, n2, n3}; + store_ring_records(recs, n, row_counts, rbase, + [&](uint32_t row, uint32_t slot, uint32_t record) { + store_ring_record(vals, indices, row, row_delta, slot, + record, lane_kv, flush_base, 0); + }); + __syncwarp(); +} + +// Drain warp: waits for producer warp `producer`'s handoff interval, stores its +// records (unless the producer self-drained it) and frees the ring bank. +template +__device__ __forceinline__ void drain_handoff( + HandoffRing& d, TBar* hfull, TBar* hempty, int producer, uint32_t interval, + uint32_t row_base, int* row_counts, bool skip, Store store) { + const int lane = threadIdx.x & 31, bank = interval & 1; + hfull[bank * kMathWarps + producer].wait((interval >> 1) & 1); + __syncwarp(); + if (skip) { // self-drained by the producer + hempty[bank * kMathWarps + producer].arrive(); + return; + } + const uint32_t base = d.bases[bank][producer]; + uint32_t n[4]; +#pragma unroll + for (uint32_t row = 0; row < 4; ++row) + n[row] = d.num_records[bank][producer][row][lane]; + store_ring_records(&d.records[bank][producer][0][0][0], n, row_counts, + row_base, + [&](uint32_t row, uint32_t slot, uint32_t record) { + store(row, slot, record, base, producer * 32 + lane); + }); + __syncwarp(); + hempty[bank * kMathWarps + producer].arrive(); +} + +// Interval: scan tiles per candidate handoff to the drain warps. +template +__global__ +__launch_bounds__(kNumThreads, 1) void fused_mqa_topk_fp8_prefill_kernel( + uint32_t rows, uint32_t keys, const uint32_t* starts_g, + const uint32_t* ends_g, int* out, uint16_t* vals, int* indices, + int* row_counts, RescanBuffers rescan, uint32_t phase_bytes, + uint32_t sample_target, uint32_t sample_stride, uint32_t sample_groups, + uint32_t sample_tiles, const float* scales, + const __grid_constant__ cute::TmaDescriptor tks, + const __grid_constant__ cute::TmaDescriptor tq, + const __grid_constant__ cute::TmaDescriptor ts, + const __grid_constant__ cute::TmaDescriptor tk, + const __grid_constant__ cute::TmaDescriptor tw, uint32_t nblocks, + uint32_t* counter, int* defer_next) { + // Compile-time: a runtime top-k costs ~0.5% in the scan (codegen). + constexpr uint32_t topk = TopK; + static_assert(BarIdx::kTmemSlots * kSlotCols <= 512, + "TMEM slots exceed 512 columns"); + // A ring record's low byte holds (tile - interval start) * 2 + half, and the + // handoff points use `& (Interval - 1)`. + static_assert((Interval & (Interval - 1)) == 0 && Interval <= 128, + "Interval must be a power of two <= 128"); + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& sm = *reinterpret_cast(smem_buffer); + uint8_t* const region = smem_buffer + sizeof(KernelSmem); + auto& sample_smem = *reinterpret_cast(region); + auto& hdata = *reinterpret_cast(region); + + __shared__ alignas(8) uint64_t bar_words[BarIdx::total]; + __shared__ uint32_t tmem_ptr; + __shared__ ClaimEntry entries[2]; + __shared__ float gate_s[kBlockRows]; + __shared__ uint32_t sel_recover, sel_staged, sel_short, + row_offset[kBlockRows + 1]; + __shared__ uint32_t defer_head, unsafe_rows, pause_req, epoch_kind; + + TBar* const bars = reinterpret_cast(bar_words); + TBar& full_q = bars[BarIdx::full_q]; + TBar& empty_q = bars[BarIdx::empty_q]; + TBar* const full_kv = bars + BarIdx::full_kv; + // The 32 producer lanes stage the sample scales, so each arrives on full_kv. + constexpr uint32_t kFullKvArrivals = 32; + TBar* const empty_kv = bars + BarIdx::empty_kv; + TBar* const full_tmem = bars + BarIdx::full_tmem; + TBar* const empty_tmem = bars + BarIdx::empty_tmem; + TBar* const claim_full = bars + BarIdx::claim_full; + TBar* const claim_empty = bars + BarIdx::claim_empty; + TBar& sel_done = bars[BarIdx::sel_done]; + TBar* const h_full = bars + BarIdx::h_full; + TBar* const h_empty = bars + BarIdx::h_empty; + TBar* const full_lin = bars + BarIdx::full_lin; + + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane = ptx::get_lane_idx(); + constexpr uint32_t kClaimConsumers = 11; // warps 9, 10, 11, 0-7 + + if (warp_idx == 8 && lane == 0) { + cute::prefetch_tma_descriptor(&tq); + cute::prefetch_tma_descriptor(&tw); + cute::prefetch_tma_descriptor(&tk); + cute::prefetch_tma_descriptor(&ts); + cute::prefetch_tma_descriptor(&tks); + } + if (warp_idx == 9 && cute::elect_one_sync()) { + defer_head = kTerm; + pause_req = 0; + } + // Epochs: the persistent role loops run until the work runs out (kTerm) or a + // block needs the exact rescan (kPause); rare blocks are then rescanned + // exactly (rescan.cuh) and the loops resume with freshly initialised + // barriers. The first epoch is emitted outside any loop (nesting the role + // loops in an epoch loop costs ~1.5 percent in the scan: uniform address math + // is rematerialised every tile). Later epochs only run after a recovery + // pause. + __shared__ uint32_t row_bounds[2][kBlockRows]; + bool more = false; + { + constexpr uint32_t epoch = 0; +#include "persistent_epoch.inc" + } +#pragma unroll 1 + for (uint32_t epoch = 1; more; ++epoch) { +#include "persistent_epoch.inc" + } + // counter = {next block, finished CTAs}: the last CTA to finish (every claim + // is done by then) zeroes both, so the caller's workspace is ready for the + // next launch / graph replay. + if (threadIdx.x == 0) { + __threadfence(); + if (atomicAdd(counter + 1, 1u) == gridDim.x - 1) { + // Acquire: every other CTA's last claim on counter happens before the + // reset below. + __threadfence(); + atomicExch(counter, 0u); + atomicExch(counter + 1, 0u); + } + } +} + +} // namespace fused_mqa_topk::persist diff --git a/csrc/fused_mqa_topk/rescan.cuh b/csrc/fused_mqa_topk/rescan.cuh new file mode 100644 index 000000000000..80d9dc57370f --- /dev/null +++ b/csrc/fused_mqa_topk/rescan.cuh @@ -0,0 +1,350 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Exact recovery for the blocks the persistent kernel cannot select from its +// gated candidates (see persistent_kernel.cuh): the block is rescanned densely, +// and every 1024-key segment of fp32 scores is merged into a running per-row +// top set (kept <= kRetainedPerRow, compacted to K). +#pragma once +#include +#include +#include +#include +#include + +#include "common.cuh" +#include "constants.h" +#include "dense_scan.cuh" + +namespace fused_mqa_topk::rescan { +using dense::Scheduler; + +// Keys whose scores are merged into the retained sets at a time. +constexpr uint32_t kSegmentKeys = 1024; +// Per-row retained top set of the exact rescan (scores and key indices). +struct RescanBuffers { + float* retained; + int* retained_indices; +}; +// Per-row segment histograms: half_bin<6> of a score (1024 bins); the radix +// levels of compact use 256 bins. +constexpr uint32_t kSegmentBins = 1024, kRadixBins = 256; +struct alignas(16) RescanShared { + float ring[2][kBlockRows][kSegmentKeys]; + uint32_t hist[kBlockRows][padded_bins(kSegmentBins)], gates[kBlockRows], + kept[kBlockRows]; + uint32_t work[padded_bins(kRadixBins)], prefix, rank, emit_counts[kMathWarps]; + __device__ void init(uint32_t tid) { + for (uint32_t j = tid; j < sizeof(hist) / sizeof(uint32_t); + j += kNumThreads) + reinterpret_cast(hist)[j] = 0; + if (tid < kBlockRows) { + gates[tid] = __float_as_uint(-CUDART_INF_F); + kept[tid] = 0; + } + } +}; +// What compact_retained's emit pass keeps. +enum class Emit { kAboveGate, kAboveCode, kTies }; + +__device__ inline void compact_retained(RescanShared& shared, + RescanBuffers buffers, uint32_t topk, + uint32_t local, uint32_t row, + uint32_t tid, + bool compact_to_k = false) { + constexpr uint32_t kThreads = kMathThreads, + kItems = kRetainedPerRow / kThreads; + uint32_t lane = tid & 31, count = shared.kept[local]; + uint64_t base = (uint64_t)row * kRetainedPerRow; + float values[kItems]; + int indices[kItems]; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + uint32_t j = tid + i * kThreads; + values[i] = j < count ? buffers.retained[base + j] : -CUDART_INF_F; + indices[i] = j < count ? buffers.retained_indices[base + j] : -1; + } + // Emits the values >= gate, the codes above threshold, or the codes equal + // to it (ties, up to K) to the retained set from base_offset on. + auto emit = [&](float gate, uint32_t threshold, Emit mode, + uint32_t base_offset) { + const auto passes = [&](float v) { + return isfinite(v) && + (mode == Emit::kAboveGate + ? v >= gate + : (mode == Emit::kAboveCode ? ordered_code(v) > threshold + : ordered_code(v) == threshold)); + }; + uint32_t total = 0; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) total += passes(values[i]); + uint32_t warp_total = __reduce_add_sync(kFullMask, total); + if (lane == 0) shared.emit_counts[tid / 32] = warp_total; + math_wg_sync(); + uint32_t start = base_offset; + for (uint32_t w = 0; w < tid / 32; ++w) start += shared.emit_counts[w]; + uint32_t all = 0; + for (uint32_t w = 0; w < kMathWarps; ++w) all += shared.emit_counts[w]; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + const float v = values[i]; + const bool pass = passes(v); + uint32_t mask = __ballot_sync(kFullMask, pass), + pos = start + __popc(mask & ((1u << lane) - 1)); + if (pass && (mode != Emit::kTies || pos < topk)) { + buffers.retained[base + pos] = v; + buffers.retained_indices[base + pos] = indices[i]; + } + start += __popc(mask); + } + math_wg_sync(); + return all; + }; + if (!compact_to_k) { + if (tid == 0) { + shared.rank = 0; + } + math_wg_sync(); + float gate = __uint_as_float(shared.gates[local]); + uint32_t kept = 0; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) + kept += isfinite(values[i]) && values[i] >= gate; + kept = __reduce_add_sync(kFullMask, kept); + if (lane == 0) atomicAdd(&shared.rank, kept); + math_wg_sync(); + kept = shared.rank; + if (kept >= topk && kept <= kRetainedPerRow - kSegmentKeys) { + emit(gate, 0, Emit::kAboveGate, 0); + math_wg_sync(); + if (tid == 0) shared.kept[local] = kept; + math_wg_sync(); + return; + } + } + if (tid == 0) { + shared.prefix = 0; + shared.rank = topk; + } + math_wg_sync(); + for (int shift = 24; shift >= 0; shift -= 8) { + for (uint32_t j = tid; j < padded_bins(kRadixBins); j += kThreads) + shared.work[j] = 0; + math_wg_sync(); + uint32_t prefix = shared.prefix; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + uint32_t code = ordered_code(values[i]); + bool valid = isfinite(values[i]) && + (shift == 24 || (code >> (shift + 8)) == prefix); + uint32_t mask = __ballot_sync(kFullMask, valid); + if (valid) { + uint32_t bin = (code >> shift) & 255u, + peers = __match_any_sync(mask, bin); + if (lane == uint32_t(__ffs(peers) - 1)) + atomicAdd(shared.work + padded_bin(bin), __popc(peers)); + } + } + math_wg_sync(); + if (tid < 32) { + uint32_t sum = 0; + for (uint32_t j = 0; j < 8; ++j) { + uint32_t bin = 255 - lane * 8 - j; + sum += shared.work[padded_bin(bin)]; + } + const uint32_t inclusive = warp_inclusive_scan(sum, lane); + uint32_t before = inclusive - sum, rank = shared.rank, + crossing = + __ballot_sync(kFullMask, before < rank && inclusive >= rank), + chosen_lane = __ffs(crossing) - 1, chosen = 0, above = before; + if (lane == chosen_lane) { + for (uint32_t j = 0; j < 8; ++j) { + uint32_t bin = 255 - lane * 8 - j; + if (above + shared.work[padded_bin(bin)] >= rank) { + chosen = bin; + break; + } + above += shared.work[padded_bin(bin)]; + } + } + chosen = __shfl_sync(kFullMask, chosen, chosen_lane); + above = __shfl_sync(kFullMask, above, chosen_lane); + __syncwarp(); // every lane read shared.rank above + if (lane == 0) { + shared.prefix = (prefix << 8) | chosen; + shared.rank = rank - above; + } + } + math_wg_sync(); + } + uint32_t threshold = shared.prefix; + uint32_t greater = emit(0, threshold, Emit::kAboveCode, 0); + emit(0, threshold, Emit::kTies, greater); + if (tid == 0) { + shared.kept[local] = topk; + float gate = + fmaxf(__uint_as_float(shared.gates[local]), ordered_float(threshold)); + atomicExch(shared.gates + local, __float_as_uint(gate)); + } + math_wg_sync(); +} + +__device__ inline void merge_segment(RescanShared& shared, + RescanBuffers buffers, uint32_t topk, + uint32_t local, uint32_t row, + uint32_t segment, uint32_t first, + uint32_t n, uint32_t lane) { + constexpr uint32_t kItems = kSegmentKeys / 32; + const uint32_t bank = segment & 1; + float values[kItems], gate = __uint_as_float(shared.gates[local]); + uint32_t count = 0; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + uint32_t j = lane + i * 32; + float v = j < n ? shared.ring[bank][local][j] : -CUDART_INF_F; + if (v <= gate) v = -CUDART_INF_F; + values[i] = v; + bool valid = isfinite(v); + count += valid; + uint32_t mask = __ballot_sync(kFullMask, valid); + if (valid) { + uint32_t bin = half_bin<6>(v), peers = __match_any_sync(mask, bin); + if (lane == uint32_t(__ffs(peers) - 1)) + atomicAdd(&shared.hist[local][padded_bin(bin)], __popc(peers)); + } + } + count = __reduce_add_sync(kFullMask, count); + __syncwarp(); + if (count) { + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < 32; ++j) + sum += shared.hist[local][padded_bin((31 - lane) * 32 + j)]; + uint32_t inclusive = warp_inclusive_scan(sum, lane); + uint32_t before = inclusive - sum, + mask = + __ballot_sync(kFullMask, before < topk && inclusive >= topk); + if (mask) { + uint32_t chosen = __ffs(mask) - 1, group = 31 - chosen, + rank = topk - __shfl_sync(kFullMask, before, chosen), + bin = group * 32 + 31 - lane; + sum = shared.hist[local][padded_bin(bin)]; + inclusive = warp_inclusive_scan(sum, lane); + before = inclusive - sum; + chosen = + __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - + 1; + bin = __shfl_sync(kFullMask, bin, chosen); + gate = fmaxf(gate, half_bin_floor<6>(bin)); + if (lane == 0) atomicExch(shared.gates + local, __float_as_uint(gate)); + } + uint32_t written = shared.kept[local]; + uint64_t base = (uint64_t)row * kRetainedPerRow; +#pragma unroll + for (uint32_t i = 0; i < kItems; ++i) { + float v = values[i]; + bool pass = isfinite(v) && v >= gate; + uint32_t mask = __ballot_sync(kFullMask, pass), + pos = written + __popc(mask & ((1u << lane) - 1)); + if (pass) { + buffers.retained[base + pos] = v; + buffers.retained_indices[base + pos] = first + lane + i * 32; + } + written += __popc(mask); + } + __syncwarp(); + if (lane == 0) shared.kept[local] = written; + __syncwarp(); + } +} + +// Streams each row's scores (above its running gate) into the segment ring and +// merges every full segment into the row's retained top set; the final set is +// written to the output. +struct RescanPolicy { + static constexpr float kWeightScale = 0.5f; + static constexpr uint32_t kSplitsPerSegment = kSegmentKeys / kTileKeys; + RescanShared& shared; + RescanBuffers buffers; + int* out; // [rows, topk] + uint32_t topk, query_base, row_base, segment = 0; + uint32_t seq_k_start[kWgRows], seq_k_end[kWgRows]; + float gate[kWgRows]; + CUTLASS_DEVICE void begin_block(const Scheduler& sched, uint32_t q_block, + uint32_t base) { + query_base = base; + row_base = sched.get_logits_row(q_block, 0); +#pragma unroll + for (uint32_t i = 0; i < kWgRows; ++i) { + const uint32_t row = sched.get_logits_row(q_block, base + i); + gate[i] = -CUDART_INF_F; + seq_k_start[i] = min(sched.cu_seq_len_k_start[row], sched.num_kv_tokens); + seq_k_end[i] = min(sched.cu_seq_len_k_end[row], sched.num_kv_tokens); + } + } + CUTLASS_DEVICE void begin_split(uint32_t split) { + if (split % kSplitsPerSegment == 0) { + __syncwarp(); + for (uint32_t i = 0; i < kWgRows; ++i) + gate[i] = __uint_as_float(atomicAdd(shared.gates + query_base + i, 0u)); + } + } + CUTLASS_DEVICE void score(uint32_t i, uint32_t kv_offset, uint32_t kv_lane, + uint32_t split, float bq) { + const bool valid = kv_offset >= seq_k_start[i] && + kv_offset < seq_k_end[i] && isfinite(bq) && bq > gate[i]; + shared.ring[segment & 1][query_base + i] + [(split % kSplitsPerSegment) * kTileKeys + kv_lane] = + valid ? bq : -CUDART_INF_F; + } + CUTLASS_DEVICE void end_split(uint32_t split, uint32_t num_splits, + uint32_t kv_base) { + if ((split + 1) % kSplitsPerSegment != 0 && split + 1 != num_splits) return; + __threadfence_block(); + math_wg_sync(); + const uint32_t tid = threadIdx.x; + for (uint32_t local = 0; local < kBlockRows; ++local) + if (shared.kept[local] > kRetainedPerRow - kSegmentKeys) + compact_retained(shared, buffers, topk, local, row_base + local, tid); + math_wg_sync(); + const uint32_t local = tid / 32; + merge_segment(shared, buffers, topk, local, row_base + local, segment, + kv_base + (split / kSplitsPerSegment) * kSegmentKeys, + (split % kSplitsPerSegment + 1) * kTileKeys, tid % 32); + math_wg_sync(); + if (split + 1 == num_splits) + for (uint32_t l = 0; l < kBlockRows; ++l) + if (shared.kept[l] > topk) + compact_retained(shared, buffers, topk, l, row_base + l, tid, true); + math_wg_sync(); + ++segment; + } + CUTLASS_DEVICE void end_block(const Scheduler& sched, uint32_t q_block) { + const uint32_t local = threadIdx.x / 32, lane = threadIdx.x % 32; + const uint32_t row = sched.get_logits_row(q_block, local), + count = shared.kept[local]; + if (row < sched.num_q_tokens) + for (uint32_t j = lane; j < topk; j += 32) + out[(uint64_t)row * topk + j] = + j < count + ? buffers.retained_indices[(uint64_t)row * kRetainedPerRow + j] + : -1; + } +}; + +// Rescans one block exactly and writes its output rows of `out`; `buffers` +// rows are addressed by absolute row (callers offset them). Out of line: the +// rare rescan must not add register pressure to the caller's hot loops. +__device__ __noinline__ inline void rescan_block( + const Scheduler& block, RescanBuffers buffers, int* out, uint32_t topk, + const cute::TmaDescriptor& tq, const cute::TmaDescriptor& tk, + const cute::TmaDescriptor& ts, const cute::TmaDescriptor& tw) { + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& shared = *reinterpret_cast(smem_buffer + + sizeof(dense::SharedStorage)); + shared.init(threadIdx.x); + RescanPolicy policy{shared, buffers, out, topk}; + dense::score_block(block, policy, tq, tk, ts, tw); +} + +} // namespace fused_mqa_topk::rescan diff --git a/csrc/fused_mqa_topk/select.cuh b/csrc/fused_mqa_topk/select.cuh new file mode 100644 index 000000000000..759ec7fce46a --- /dev/null +++ b/csrc/fused_mqa_topk/select.cuh @@ -0,0 +1,452 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Per-row top-K selection of the persistent kernel's candidates, one warp per +// row (K = topk, the row's top-k: 512, 1024 or 2048, a runtime value). +// +// select_staged_row: bracket the K-th largest fp24 code with sampled codes, +// place the codes above the bracket directly, and resolve the bracket with +// histogram levels on register-resident codes. +// +// select_row: the fallback of select_staged_row when its sampled bracket +// misses. Candidate codes (higher code = higher score, staged once in shared +// memory) span a narrow range (gate .. row max), so each level histograms the +// range-relative code into at most kBins bins: collisions are rare, plain +// shared atomics suffice and each lane keeps kUnroll independent elements in +// flight. Levels repeat inside the boundary bin until it is one code wide; +// codes above it are emitted, then ties until K. +#pragma once +#include + +#include "candidates.cuh" +#include "common.cuh" +#include "constants.h" + +namespace fused_mqa_topk::selection { +constexpr uint32_t kBinsLog2 = 10, kBins = 1u << kBinsLog2; +// Elements each lane keeps in flight. +constexpr uint32_t kUnroll = 8; + +// Bin (descending) where the cumulative count first reaches `rank`; +// `above` = elements in higher bins. Lane l owns Bins / 32 bins, descending +// from bin Bins - 1 - l * Bins / 32. +template +__device__ __forceinline__ uint32_t find_bin_n(const uint32_t* hist, + uint32_t rank, uint32_t& above) { + const uint32_t lane = threadIdx.x & 31; + constexpr uint32_t kPerLane = Bins / 32; + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < kPerLane; ++j) + sum += hist[Bins - 1 - lane * kPerLane - j]; + const uint32_t inclusive = warp_inclusive_scan(sum, lane); + const uint32_t before = inclusive - sum; + const uint32_t winner = + __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - 1; + uint32_t chosen = 0, acc = before; + if (lane == winner) { + for (uint32_t j = 0; j < kPerLane; ++j) { + const uint32_t bin = Bins - 1 - lane * kPerLane - j, h = hist[bin]; + if (acc + h >= rank) { + chosen = bin; + break; + } + acc += h; + } + } + above = __shfl_sync(kFullMask, acc, winner); + return __shfl_sync(kFullMask, chosen, winner); +} + +// Requires count >= topk (callers recover rows with fewer candidates). +__device__ inline void select_row(uint32_t* codes, uint32_t count, + const uint16_t* values, const int32_t* packed, + int32_t* out, uint32_t topk, uint32_t* hist) { + const uint32_t lane = threadIdx.x & 31; + uint32_t lo = 0xffffffffu, hi = 0; + uint32_t j = lane; + for (; j + 32 * (kUnroll - 1) < count; j += 32 * kUnroll) { + uint16_t v[kUnroll]; + int32_t p[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + v[u] = values[j + 32 * u]; + p[u] = packed[j + 32 * u]; + } +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + const uint32_t c = candidate_load_score_code(v[u], p[u]); + codes[j + 32 * u] = c; + lo = min(lo, c); + hi = max(hi, c); + } + } + for (; j < count; j += 32) { + const uint32_t c = candidate_load_score_code(values[j], packed[j]); + codes[j] = c; + lo = min(lo, c); + hi = max(hi, c); + } + lo = __reduce_min_sync(kFullMask, lo); + hi = __reduce_max_sync(kFullMask, hi); + __syncwarp(); + + // Current bucket: codes in [lo, lo + width). rank = elements still needed + // from it. + uint32_t width = hi - lo + 1, rank = topk; +#pragma unroll 1 + while (width > 1) { + const uint32_t spread = (width - 1) >> kBinsLog2; + const uint32_t shift = spread ? 32 - __clz(spread) : 0; + for (uint32_t b = lane; b < kBins; b += 32) hist[b] = 0; + __syncwarp(); + const uint32_t top = lo + width - 1; + uint32_t i = lane; + for (; i + 32 * (kUnroll - 1) < count; i += 32 * kUnroll) { + uint32_t c[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) c[u] = codes[i + 32 * u]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) + if (c[u] >= lo && c[u] <= top) + atomicAdd(hist + ((c[u] - lo) >> shift), 1u); + } + for (; i < count; i += 32) { + const uint32_t c = codes[i]; + if (c >= lo && c <= top) atomicAdd(hist + ((c - lo) >> shift), 1u); + } + __syncwarp(); + uint32_t above; + const uint32_t chosen = find_bin_n(hist, rank, above); + rank -= above; + lo += chosen << shift; + width = min(1u << shift, top - lo + 1); + __syncwarp(); + } + + // lo is now the K-th largest code: emit codes above it, then ties until K. + const uint32_t threshold = lo, lt = (1u << lane) - 1u; + uint32_t pos = 0; +#pragma unroll 1 + for (uint32_t tie = 0; tie < 2; ++tie) { + for (uint32_t j0 = 0; j0 < count && pos < topk; j0 += 32 * kUnroll) { + bool pass[kUnroll]; + int32_t p[kUnroll]; +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + const uint32_t e = j0 + 32 * u + lane; + const uint32_t c = e < count ? codes[e] : 0; + pass[u] = e < count && (tie ? c == threshold : c > threshold); + p[u] = pass[u] ? packed[e] : 0; + } +#pragma unroll + for (uint32_t u = 0; u < kUnroll; ++u) { + const uint32_t mask = __ballot_sync(kFullMask, pass[u]), + at = pos + __popc(mask & lt); + if (pass[u] && at < topk) out[at] = candidate_decode_index(p[u]); + pos += __popc(mask); + } + } + } +} + +// Register samples per lane for the sampled bracket of select_staged_row. +constexpr uint32_t kSamplesPerLane = 16; +// Largest bracket (elements between the two sampled bounds) that +// select_staged_row resolves in registers; wider brackets fall back to +// select_row. +constexpr uint32_t kMidCap = 1024; +// Bins of the histogram levels that resolve the exact threshold inside a +// bracket (select_staged_row). +constexpr uint32_t kMidBins = 512; + +// Bracket element ids are stored as uint16 (midbuf). +static_assert(kCandidatesPerRow <= 65536, "candidate ids must fit 16 bits"); +// Per-warp scratch of select_staged_row is kBins words: radix_bracket's 768, +// or the midbuf (kMidCap uint16) plus the kMidBins-bin histogram. +static_assert(768 <= kBins && kMidCap / 2 + kMidBins <= kBins, + "selection scratch exceeds kBins words"); + +// r-th largest (1-based) of the kNumSamples * 32 register samples (kNumSamples +// per lane) for two ranks at once +// (r_hi <= r_lo), via a 256-bin hi8 level and two 256-bin lo8 levels. h8: 768 +// words of scratch. +template +__device__ __forceinline__ void radix_bracket( + const uint32_t (&smp)[kNumSamples], uint32_t r_hi, uint32_t r_lo, + uint32_t* h8, uint32_t& h_hi, uint32_t& h_lo) { + const uint32_t lane = threadIdx.x & 31; +#pragma unroll + for (uint32_t b = lane; b < 192; b += 32) + reinterpret_cast(h8)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (uint32_t u = 0; u < kNumSamples; ++u) atomicAdd(h8 + (smp[u] >> 8), 1u); + __syncwarp(); + uint32_t ab_hi, ab_lo; + const uint32_t b_hi = find_bin_n<256>(h8, r_hi, ab_hi); + const uint32_t b_lo = find_bin_n<256>(h8, r_lo, ab_lo); + uint32_t* const lh = h8 + 256; + uint32_t* const ll = h8 + 512; +#pragma unroll + for (uint32_t u = 0; u < kNumSamples; ++u) { + const uint32_t b = smp[u] >> 8; + if (b == b_hi) atomicAdd(lh + (smp[u] & 255u), 1u); + if (b == b_lo) atomicAdd(ll + (smp[u] & 255u), 1u); + } + __syncwarp(); + uint32_t d0, d1; + h_hi = b_hi << 8 | find_bin_n<256>(lh, r_hi - ab_hi, d0); + h_lo = b_lo << 8 | find_bin_n<256>(ll, r_lo - ab_lo, d1); + __syncwarp(); // h8 is reused as the output stage +} + +// Brackets the topk-th largest of the `count` staged codes between the hi16 +// codes of sample ranks expect -+ margin, where the kPerLane * 32 samples sit +// at evenly spaced quantiles and expect is the topk-th largest's expected rank +// among them. +template +__device__ __forceinline__ void sampled_bracket(const uint32_t* codes, + uint32_t count, uint32_t topk, + uint32_t margin, uint32_t* hist, + uint32_t& h_hi, + uint32_t& h_lo) { + constexpr uint32_t kSamples = kPerLane * 32; + const uint32_t lane = threadIdx.x & 31; + uint32_t smp[kPerLane]; +#pragma unroll + for (uint32_t u = 0; u < kPerLane; ++u) + smp[u] = codes[((u * 32 + lane) * 2 + 1) * count / (2 * kSamples)] >> 8; + const uint32_t expect = (topk * kSamples + count / 2) / count; + const uint32_t r_hi = expect > margin ? expect - margin : 1, + r_lo = min(expect + margin, kSamples); + radix_bracket(smp, r_hi, r_lo, hist, h_hi, h_lo); +} + +// find_bin_n with each lane's kPer contiguous bins loaded as 16-byte vectors +// and searched in registers (needs a 16-byte aligned histogram and kPer % 4 == +// 0; find_bin_n reads the bins one at a time and has neither requirement). +template +__device__ __forceinline__ uint32_t find_bin_v(const uint32_t* hist, + uint32_t rank, uint32_t& above) { + const uint32_t lane = threadIdx.x & 31; + constexpr uint32_t kPer = Bins / 32, kVec = kPer / 4; + static_assert(kPer % 4 == 0); + // lane l owns bins [Bins - (l + 1) * kPer, Bins - l * kPer), highest first. + const uint4* hv = + reinterpret_cast(hist + Bins - (lane + 1) * kPer); + uint32_t h[kPer]; +#pragma unroll + for (uint32_t v = 0; v < kVec; ++v) { + const uint4 x = hv[v]; + h[4 * v] = x.x; + h[4 * v + 1] = x.y; + h[4 * v + 2] = x.z; + h[4 * v + 3] = x.w; + } + uint32_t sum = 0; +#pragma unroll + for (uint32_t j = 0; j < kPer; ++j) sum += h[j]; + const uint32_t inclusive = warp_inclusive_scan(sum, lane); + const uint32_t before = inclusive - sum; + const uint32_t winner = + __ffs(__ballot_sync(kFullMask, before < rank && inclusive >= rank)) - 1; + // Descending walk inside the lane: bin index Bins - lane * kPer - 1 - j <-> + // h[kPer - 1 - j]. + uint32_t acc = before, chosen = 0, found = 0, acc_at = before; +#pragma unroll + for (uint32_t j = 0; j < kPer; ++j) { + const uint32_t x = h[kPer - 1 - j]; + const bool hit = !found && acc + x >= rank; + if (hit) { + chosen = Bins - lane * kPer - 1 - j; + acc_at = acc; + found = 1; + } + acc += x; + } + above = __shfl_sync(kFullMask, acc_at, winner); + return __shfl_sync(kFullMask, chosen, winner); +} + +// One row's global-memory candidates, staged to shared memory as +// fp24 codes (codes must be 16-byte aligned, i.e. row offsets rounded up to 4). +// Register budget: ptxas compiles the whole kernel to 168 regs, so every pass +// keeps <= ~32 regs of arrays live (more spills into the scan loop). +// Passes: stage (16-byte global loads of values / packed indices, two chunks of +// 8 in flight per lane, to fp24 codes); hi16 bracket (radix levels over 512 +// shared-memory samples); count above / inside the bracket (16-byte shared +// loads); collect (codes above the bracket to out, bracket element ids to +// midbuf); histogram levels on register codes; emit. +// Returns false before writing out; the caller's select_row then restages and +// runs the levels. +__device__ __noinline__ inline bool select_staged_row( + uint32_t* __restrict__ codes, uint32_t count, + const uint16_t* __restrict__ values, const int32_t* __restrict__ packed, + int32_t* __restrict__ out, uint32_t topk, uint32_t* hist) { + const uint32_t lane = threadIdx.x & 31; + if (count <= topk) return false; + const uint32_t chunks = (count + 7) / 8; + constexpr uint32_t kChunksInFlight = 2; + const uint4* v4p = reinterpret_cast(values); + const uint4* p4p = reinterpret_cast(packed); +#pragma unroll 1 + for (uint32_t c0i = lane; c0i < chunks; c0i += 32 * kChunksInFlight) { + uint4 vv[kChunksInFlight], pp0[kChunksInFlight], pp1[kChunksInFlight]; +#pragma unroll + for (uint32_t u = 0; u < kChunksInFlight; ++u) { + const uint32_t c = min(c0i + 32 * u, chunks - 1); + vv[u] = __ldcg(v4p + c); + pp0[u] = __ldcg(p4p + 2 * c); + pp1[u] = __ldcg(p4p + 2 * c + 1); + } +#pragma unroll + for (uint32_t u = 0; u < kChunksInFlight; ++u) { + const uint32_t c = c0i + 32 * u; + if (c < chunks) { + const uint32_t vw[4] = {vv[u].x, vv[u].y, vv[u].z, vv[u].w}; + const uint32_t pw[8] = {pp0[u].x, pp0[u].y, pp0[u].z, pp0[u].w, + pp1[u].x, pp1[u].y, pp1[u].z, pp1[u].w}; + uint32_t cw[8]; +#pragma unroll + for (uint32_t k = 0; k < 8; ++k) + cw[k] = candidate_load_score_code( + uint16_t(vw[k / 2] >> (16 * (k & 1))), int32_t(pw[k])); + if (8 * c + 8 <= count) { + reinterpret_cast(codes)[2 * c] = + make_uint4(cw[0], cw[1], cw[2], cw[3]); + reinterpret_cast(codes)[2 * c + 1] = + make_uint4(cw[4], cw[5], cw[6], cw[7]); + } else { +#pragma unroll + for (uint32_t k = 0; k < 8; ++k) + if (8 * c + k < count) codes[8 * c + k] = cw[k]; + } + } + } + } + __syncwarp(); + // The bracket is the code range between the sampled ranks expect -+ margin + // (expect = the K-th largest's expected rank among the samples). Rows up to + // kFewCandidates use 512 samples, +-40 ranks. Long-context rows (4-6K + // candidates) would get a bracket of ~16% of the row that way, over kMidCap + // for 10-20% of the rows at 512K-1M: 1024 samples, +-48 (~3 sd), give ~9%. + constexpr uint32_t kFewCandidates = 3584, kMargin = 40, kWideMargin = 48; + uint32_t h_hi, h_lo; + if (count <= kFewCandidates) + sampled_bracket(codes, count, topk, kMargin, hist, h_hi, + h_lo); + else + sampled_bracket<2 * kSamplesPerLane>(codes, count, topk, kWideMargin, hist, + h_hi, h_lo); + const uint4* c4p = reinterpret_cast(codes); + const uint32_t quads = (count + 3) / 4; + const uint32_t lt = (1u << lane) - 1u; + uint16_t* midbuf = reinterpret_cast(hist); + uint32_t* h512 = hist + kMidCap / 2; + // Collect with warp-ballot placement: passing lanes of one step write + // consecutive slots, so the global stores of the above elements are + // coalesced. Writes are bounded; the + // bracket check follows (out may be partially written on failure: the + // fallback select_row rewrites all K slots). + uint32_t above = 0, mid = 0; + constexpr uint32_t kQuadsInFlight = 4; +#pragma unroll 1 + for (uint32_t q0 = 0; q0 < quads; q0 += 32 * kQuadsInFlight) { + uint4 pp[kQuadsInFlight], cc[kQuadsInFlight]; +#pragma unroll + for (uint32_t u = 0; u < kQuadsInFlight; ++u) { + const uint32_t q = min(q0 + 32 * u + lane, quads - 1); + pp[u] = __ldcg(p4p + q); + cc[u] = c4p[q]; + } +#pragma unroll + for (uint32_t u = 0; u < kQuadsInFlight; ++u) { + const uint32_t q = q0 + 32 * u + lane; + const uint32_t w[4] = {cc[u].x, cc[u].y, cc[u].z, cc[u].w}; + const uint32_t p[4] = {pp[u].x, pp[u].y, pp[u].z, pp[u].w}; +#pragma unroll + for (uint32_t k = 0; k < 4; ++k) { + const uint32_t e = 4 * q + k, h = w[k] >> 8; + const bool valid = e < count; + const bool a = valid && h > h_hi, m = valid && !a && h >= h_lo; + const uint32_t ma = __ballot_sync(kFullMask, a), + mm = __ballot_sync(kFullMask, m); + const uint32_t at = above + __popc(ma & lt), mt = mid + __popc(mm & lt); + if (a && at < topk) out[at] = candidate_decode_index(int32_t(p[k])); + if (m && mt < kMidCap) midbuf[mt] = uint16_t(e); + above += __popc(ma); + mid += __popc(mm); + } + } + } + __syncwarp(); + if (above >= topk || above + mid < topk || mid > kMidCap) return false; + // Bracket codes per lane (register-resident). + constexpr uint32_t kMidPerLane = kMidCap / 32; + uint32_t mc[kMidPerLane]; +#pragma unroll + for (uint32_t u = 0; u < kMidPerLane; ++u) { + const uint32_t i = lane + 32 * u; + mc[u] = 0; + if (i < mid) mc[u] = codes[midbuf[i]]; + } + uint32_t rank = topk - above, base = h_lo << 8, top = h_hi << 8 | 0xffu; +#pragma unroll 1 + while (top > base) { + const uint32_t width = top - base + 1, spread = (width - 1) / kMidBins; + const uint32_t shift = spread ? 32 - __clz(spread) : 0; +#pragma unroll + for (uint32_t b = lane; b < kMidBins / 4; b += 32) + reinterpret_cast(h512)[b] = make_uint4(0, 0, 0, 0); + __syncwarp(); +#pragma unroll + for (uint32_t u = 0; u < kMidPerLane; ++u) { + if (u * 32 >= mid) break; + if (lane + 32 * u < mid && mc[u] >= base && mc[u] <= top) + atomicAdd(h512 + ((mc[u] - base) >> shift), 1u); + } + __syncwarp(); + uint32_t above_bin; + const uint32_t chosen = find_bin_v(h512, rank, above_bin); + rank -= above_bin; + base += chosen << shift; + top = min(base + (1u << shift) - 1, top); + __syncwarp(); + } + const uint32_t th = base; + { + // Ballot placement (coalesced stores): greater codes fill [above, K - + // rank), ties [K - rank, K). + uint32_t pg = above, pt = topk - rank; +#pragma unroll + for (uint32_t h = 0; h < 2; ++h) { + if (h * 16 * 32 >= mid) break; + int32_t pk[16]; +#pragma unroll + for (uint32_t j = 0; j < 16; ++j) { + const uint32_t u = h * 16 + j; + pk[j] = (lane + 32 * u < mid && mc[u] >= th) + ? __ldcg(packed + midbuf[lane + 32 * u]) + : 0; + } +#pragma unroll + for (uint32_t j = 0; j < 16; ++j) { + const uint32_t u = h * 16 + j; + if (u * 32 >= mid) break; + const bool v = lane + 32 * u < mid; + const bool g = v && mc[u] > th, tq = v && mc[u] == th; + const uint32_t mg = __ballot_sync(kFullMask, g), + mt = __ballot_sync(kFullMask, tq); + const uint32_t ag = pg + __popc(mg & lt), at = pt + __popc(mt & lt); + if (g) out[ag] = candidate_decode_index(pk[j]); + if (tq && at < topk) out[at] = candidate_decode_index(pk[j]); + pg += __popc(mg); + pt += __popc(mt); + } + } + } + return true; +} + +} // namespace fused_mqa_topk::selection diff --git a/csrc/hisparse_kernels.cu b/csrc/hisparse_kernels.cu index da3857d1e973..cae27cb44a48 100644 --- a/csrc/hisparse_kernels.cu +++ b/csrc/hisparse_kernels.cu @@ -21,6 +21,7 @@ #include #include #include +#include #include #include @@ -155,13 +156,25 @@ __device__ __forceinline__ int warp_inclusive_scan(int32_t* s_data, int lane_id, return __shfl_sync(0xffffffff, val, 31); } +// Non-negative x only. Block sizes are powers of two in practice, and integer +// division would otherwise dominate the resolver's per-position instructions. +__device__ __forceinline__ int32_t block_div(int32_t x, int32_t block_size) { + return (block_size & (block_size - 1)) == 0 ? x >> (__ffs(block_size) - 1) + : x / block_size; +} + +__device__ __forceinline__ int32_t block_mod(int32_t x, int32_t block_size) { + return (block_size & (block_size - 1)) == 0 ? x & (block_size - 1) + : x % block_size; +} + __device__ __forceinline__ int64_t get_physical_hot_row(const int32_t* hot_block_table, int32_t row, int64_t table_stride, int32_t block_size, int32_t slot) { const int32_t block = hot_block_table[static_cast(row) * table_stride + - slot / block_size]; - return static_cast(block) * block_size + slot % block_size; + block_div(slot, block_size)]; + return static_cast(block) * block_size + block_mod(slot, block_size); } __device__ __forceinline__ void store_hot_index( @@ -173,8 +186,8 @@ __device__ __forceinline__ void store_hot_index( attention_indices[index] = physical_row < 0 ? -1 - : (physical_row / hot_block_size) * attention_block_stride + - physical_row % hot_block_size; + : block_div(physical_row, hot_block_size) * attention_block_stride + + block_mod(physical_row, hot_block_size); } } @@ -185,9 +198,11 @@ constexpr int32_t kEntryMiss = 1 << 29; constexpr int32_t kEntryInvalid = 1 << 28; constexpr int32_t kEntrySlotMask = 0xFFFF; constexpr int kResidencyCounters = 5; -// Per-request table entries per referenced top-k position (load <= 1/4 -// until the hash_size cap). +// Per-request table entries per referenced top-k position, within the +// device's shared-memory budget. constexpr int kHashLoadInverse = 4; +// Top-k positions per thread the resolver loads before its request lookups. +constexpr int kPrefetchedTopK = 2; // Translate a request-relative top-k position to its resident GPU row (-1 when // the page is not resident) and returns its host row (-1 when not host-backed). @@ -204,7 +219,7 @@ __device__ __forceinline__ int32_t translate_topk_entry( resident_row = -1; if (source_block_table != nullptr) { const int32_t source_block = - token_index >= 0 ? token_index / source_block_size : -1; + token_index >= 0 ? block_div(token_index, source_block_size) : -1; if (request_row >= 0 && request_row < source_num_reqs && source_block >= 0 && source_block < source_num_blocks) { const int32_t physical_block = @@ -212,14 +227,14 @@ __device__ __forceinline__ int32_t translate_topk_entry( source_bt_stride + source_block]; g = physical_block > 0 ? physical_block * source_block_size + - token_index % source_block_size + block_mod(token_index, source_block_size) : -1; } else { g = -1; } if (resident_block_table != nullptr) { const int32_t resident_block = - token_index >= 0 ? token_index / resident_block_size : -1; + token_index >= 0 ? block_div(token_index, resident_block_size) : -1; if (state_row >= 0 && state_row < resident_num_reqs && resident_block >= 0 && resident_block < resident_num_blocks) { const int32_t physical_block = @@ -228,7 +243,7 @@ __device__ __forceinline__ int32_t translate_topk_entry( resident_block]; if (physical_block != resident_null_block && physical_block >= 0) { resident_row = physical_block * resident_block_size + - token_index % resident_block_size; + block_mod(token_index, resident_block_size); } } } @@ -236,6 +251,12 @@ __device__ __forceinline__ int32_t translate_topk_entry( return g >= host_rows ? -1 : g; } +// Phase 1 parks host rows in hot_indices as values below -1, distinct from +// resident rows (>= 0) and invalid positions (-1). +__device__ __forceinline__ int32_t encode_host_row(const int32_t g) { + return -2 - g; +} + // Write the hot index of one reference to a resolved union entry and tally it. __device__ __forceinline__ void write_entry_index( const int32_t entry, const int row, const int i, const int32_t top_k, @@ -285,11 +306,15 @@ __device__ __forceinline__ void write_entry_index( // s_hash_vals[hash_size] first position, then slot | kEntry* tag // s_chunk_off[nbc + 1] prefix sums for hit, then miss, compaction // s_evict_off[nbc + 1] prefix sums for evictable compaction -// s_counters[5] distinct hits, union size, per-row valid count, -// misses, per-entry hits -// s_done[nbc] positions phase 3 resolved, one word per chunk +// s_counters[5] distinct hits, union size, per-entry misses, +// host positions (phase 1) then misses, +// per-entry hits +// s_done[nbc] hot slots holding a union row (phase 2), then +// positions phase 3 resolved, one word per chunk // s_lru_out[hot_size] int16, compacted slots: [hits fwd | evict bwd] -// Valid global ids must be unique within each row. +// Valid global ids must be unique within each row. A nonzero kBlockSize fixes +// every block size at compile time. +template __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( const int32_t* __restrict__ hot_block_table, // [max_rows, hot_blocks] const int32_t* __restrict__ global_indices, // global or request-relative @@ -310,79 +335,98 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( const int32_t* __restrict__ request_state_indices, // [num_requests] or // nullptr const int32_t request_state_count, const int64_t host_rows, - const int64_t hot_table_stride, const int32_t hot_block_size, + const int64_t hot_table_stride, const int32_t hot_block_size_arg, const int32_t top_k, const int32_t hot_size, const int32_t hash_size, - const int64_t region_stride, const int64_t attention_block_stride, - const int64_t source_bt_stride, const int32_t source_num_reqs, - const int32_t source_num_blocks, const int32_t source_block_size, - const int64_t resident_bt_stride, const int32_t resident_num_reqs, - const int32_t resident_num_blocks, const int32_t resident_block_size, - const int32_t resident_null_block, const int64_t input_row_stride, - const int64_t attention_row_stride, const int64_t valid_count_stride) { + const int32_t max_union_rows, const int64_t region_stride, + const int64_t attention_block_stride, const int64_t source_bt_stride, + const int32_t source_num_reqs, const int32_t source_num_blocks, + const int32_t source_block_size_arg, const int64_t resident_bt_stride, + const int32_t resident_num_reqs, const int32_t resident_num_blocks, + const int32_t resident_block_size_arg, const int32_t resident_null_block, + const int64_t input_row_stride, const int64_t attention_row_stride, + const int64_t valid_count_stride, + const int32_t* __restrict__ num_valid_rows) { // [1] or nullptr + const int32_t hot_block_size = + kBlockSize > 0 ? kBlockSize : hot_block_size_arg; + const int32_t source_block_size = + kBlockSize > 0 ? kBlockSize : source_block_size_arg; + const int32_t resident_block_size = + kBlockSize > 0 ? kBlockSize : resident_block_size_arg; const int NUM_WARPS = blockDim.x / kWarpSize; const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize; - // The first row of each contiguous run of a request's rows leads it. FULL - // CUDA-graph padding rows map to request 0 after the real rows, forming a - // second run, but carry no valid top-k: that run returns before phase 2 - // touches the request's shared state. + const int tid = threadIdx.x; + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1; const int first_row = blockIdx.x; - const int request_row = - request_ids != nullptr ? request_ids[first_row] : first_row; - if (request_ids != nullptr && first_row > 0 && - request_ids[first_row - 1] == request_row) { - return; - } - int end_row = first_row + 1; - if (request_ids != nullptr) { - while (end_row < static_cast(gridDim.x) && - request_ids[end_row] == request_row) { - ++end_row; - } + + // Issued before the request lookups resolve so that an all-resident step + // pays few dependent memory round trips. + int32_t first_topk[kPrefetchedTopK]; +#pragma unroll + for (int k = 0; k < kPrefetchedTopK; ++k) { + const int i = tid + k * blockDim.x; + first_topk[k] = i < top_k ? global_indices[static_cast(first_row) * + input_row_stride + + i] + : -1; } + // FULL CUDA-graph padding rows past num_valid_rows belong to no request. + const int valid_rows = num_valid_rows != nullptr + ? min(*num_valid_rows, static_cast(gridDim.x)) + : static_cast(gridDim.x); + const auto request_of = [&](const int row) -> int32_t { + if (request_ids == nullptr) return row; + const int32_t request = request_ids[row]; + return row < valid_rows ? request : -1; + }; + const int request_row = request_of(first_row); const int state_row = request_state_indices != nullptr && request_row >= 0 && request_row < request_state_count ? request_state_indices[request_row] : (request_state_indices == nullptr ? request_row : -1); - // V2 publishes -1 for CUDA-graph padding rows. + // V2 publishes -1 for CUDA-graph padding rows. They share no request state, + // so every block masks its own row. if (state_row < 0) { - for (int row = first_row; row < end_row; ++row) { - for (int i = threadIdx.x; i < top_k; i += blockDim.x) { - const int64_t index = static_cast(row) * top_k + i; - hot_indices[index] = -1; - if (attention_indices != nullptr) { - attention_indices[static_cast(row) * attention_row_stride + - i] = -1; - } - if (resolved_global_indices != nullptr) { - resolved_global_indices[index] = -1; - } - } - if (valid_counts != nullptr && threadIdx.x == 0) { - valid_counts[static_cast(row) * valid_count_stride] = 0; + for (int i = tid; i < top_k; i += blockDim.x) { + const int64_t index = static_cast(first_row) * top_k + i; + hot_indices[index] = -1; + if (attention_indices != nullptr) { + attention_indices[static_cast(first_row) * + attention_row_stride + + i] = -1; } - if (swap_counts != nullptr && threadIdx.x == 0) { - swap_counts[row] = 0; + if (resolved_global_indices != nullptr) { + resolved_global_indices[index] = -1; } } + if (valid_counts != nullptr && tid == 0) { + valid_counts[static_cast(first_row) * valid_count_stride] = 0; + } + if (swap_counts != nullptr && tid == 0) { + swap_counts[first_row] = 0; + } return; } - const int tid = threadIdx.x; - const int warp_id = tid / kWarpSize; - const int lane_id = tid % kWarpSize; - const unsigned int lanes_before = ((unsigned int)1 << lane_id) - 1; + // The first row of each contiguous run of a request's rows leads it. + if (request_ids != nullptr && first_row > 0 && + request_of(first_row - 1) == request_row) { + return; + } + int end_row = first_row + 1; + if (request_ids != nullptr) { + while (end_row < static_cast(gridDim.x) && + request_of(end_row) == request_row) { + ++end_row; + } + } int32_t* row_dgi = device_global_indices + static_cast(state_row) * region_stride; int16_t* row_lru = lru_slots + static_cast(state_row) * hot_size; - // A union within max_union_rows (hash_size - 1) and within the rows' top-k - // leaves an empty slot, so every probe sequence ends. Use a slice sized to - // this request's rows so single-row requests initialize a small table. - const int table_size = - min(hash_size, kHashLoadInverse * (end_row - first_row) * top_k); - extern __shared__ char smem_raw[]; int32_t* s_hash_keys = reinterpret_cast(smem_raw); int32_t* s_hash_vals = s_hash_keys + hash_size; @@ -393,10 +437,6 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( reinterpret_cast(s_counters + kResidencyCounters); int16_t* s_lru_out = reinterpret_cast(s_done + num_buffer_chunks); - for (int i = tid; i < table_size; i += blockDim.x) { - s_hash_keys[i] = kHashEmpty; - s_hash_vals[i] = INT32_MAX; - } for (int i = tid; i < num_buffer_chunks + 1; i += blockDim.x) { s_chunk_off[i] = 0; s_evict_off[i] = 0; @@ -406,129 +446,195 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( } __syncthreads(); - // Phase 1: translate every row's request-relative positions, resolve - // resident rows directly, and collect the union of host rows. - for (int row = first_row; row < end_row; ++row) { - const int32_t* row_topk = - global_indices + static_cast(row) * input_row_stride; + // Phase 1: translate every row's request-relative positions in one pass, + // resolve resident rows directly, and count the host-backed positions. + // s_chunk_off holds per-row valid counts until phase 2. + int32_t* s_row_valid = s_chunk_off; + const int num_positions = (end_row - first_row) * top_k; + const auto resolve_position = [&](const int position, + const int32_t token_index) { + const int row = first_row + position / top_k; + const int i = position % top_k; int32_t* row_out = hot_indices + static_cast(row) * top_k; int32_t* row_attention = attention_indices != nullptr ? attention_indices + static_cast(row) * attention_row_stride : nullptr; - int32_t* row_miss = miss_mask != nullptr - ? miss_mask + static_cast(row) * top_k - : nullptr; - for (int i = tid; i < top_k; i += blockDim.x) { - int32_t resident_row; - const int32_t g = translate_topk_entry( - row_topk[i], request_row, state_row, source_block_table, - resident_block_table, host_rows, source_bt_stride, source_num_reqs, - source_num_blocks, source_block_size, resident_bt_stride, - resident_num_reqs, resident_num_blocks, resident_block_size, - resident_null_block, resident_row); - if (resolved_global_indices != nullptr) { - resolved_global_indices[static_cast(row) * top_k + i] = g; + int32_t resident_row; + const int32_t g = translate_topk_entry( + token_index, request_row, state_row, source_block_table, + resident_block_table, host_rows, source_bt_stride, source_num_reqs, + source_num_blocks, source_block_size, resident_bt_stride, + resident_num_reqs, resident_num_blocks, resident_block_size, + resident_null_block, resident_row); + if (resolved_global_indices != nullptr) { + resolved_global_indices[static_cast(row) * top_k + i] = g; + } + if (resident_row >= 0 || g >= 0) + atomicAdd(&s_row_valid[row - first_row], 1); + if (miss_mask != nullptr) + miss_mask[static_cast(row) * top_k + i] = 0; + if (resident_row >= 0) { + store_hot_index(row_out, row_attention, i, resident_row, hot_block_size, + attention_block_stride); + } else if (g < 0) { + store_hot_index(row_out, row_attention, i, -1, hot_block_size, + attention_block_stride); + } else { + // Stash the host row in its output slot; later phases overwrite it. + row_out[i] = encode_host_row(g); + atomicAdd(&s_counters[3], 1); + } + }; +#pragma unroll + for (int k = 0; k < kPrefetchedTopK; ++k) { + const int position = tid + k * blockDim.x; + if (position < num_positions) { + resolve_position(position, + position < top_k + ? first_topk[k] + : global_indices[static_cast( + first_row + position / top_k) * + input_row_stride + + position % top_k]); + } + } + for (int position = tid + kPrefetchedTopK * blockDim.x; + position < num_positions; position += blockDim.x) { + resolve_position( + position, + global_indices[static_cast(first_row + position / top_k) * + input_row_stride + + position % top_k]); + } + __syncthreads(); + for (int row = first_row + tid; row < end_row; row += blockDim.x) { + if (valid_counts != nullptr) { + valid_counts[static_cast(row) * valid_count_stride] = + s_row_valid[row - first_row]; + } + s_row_valid[row - first_row] = 0; + } + // Fully resident rows need only request-relative page translation. Avoid + // building the table or touching the hot LRU when no row can consult it. + const int host_positions = s_counters[3]; + if (host_positions == 0) { + for (int row = first_row + tid; row < end_row; row += blockDim.x) { + if (swap_counts != nullptr) swap_counts[row] = 0; + } + return; + } + + // Use a slice sized to this request's rows so single-row requests + // initialize a small table. Legal unions always leave an empty slot. + const int table_size = + min(hash_size, kHashLoadInverse * (end_row - first_row) * top_k); + for (int i = tid; i < table_size; i += blockDim.x) { + s_hash_keys[i] = kHashEmpty; + s_hash_vals[i] = INT32_MAX; + } + __syncthreads(); + + // Phase 1b: collect the union of host rows. + for (int position = tid; position < num_positions; position += blockDim.x) { + const int32_t stashed = + hot_indices[static_cast(first_row) * top_k + position]; + if (stashed > -2) continue; + const int32_t g = encode_host_row(stashed); + int h = hash_slot(g, table_size); + bool inserted = false; + for (int probe = 0; probe < table_size; ++probe) { + const int32_t old = atomicCAS(&s_hash_keys[h], kHashEmpty, g); + if (old == kHashEmpty) { + atomicAdd(&s_counters[1], 1); + inserted = true; + break; } - if (resident_row >= 0 || g >= 0) atomicAdd(&s_counters[2], 1); - if (row_miss != nullptr) row_miss[i] = 0; - if (resident_row >= 0) { - store_hot_index(row_out, row_attention, i, resident_row, hot_block_size, - attention_block_stride); - } else if (g < 0) { - store_hot_index(row_out, row_attention, i, -1, hot_block_size, - attention_block_stride); - } else { - int h = hash_slot(g, table_size); - bool inserted = false; - for (int probe = 0; probe < table_size; ++probe) { - const int32_t old = atomicCAS(&s_hash_keys[h], kHashEmpty, g); - if (old == kHashEmpty) { - atomicAdd(&s_counters[1], 1); - inserted = true; - break; - } - if (old == g) { - inserted = true; - break; - } - h = next_slot(h, table_size); - } - if (inserted) { - atomicMin(&s_hash_vals[h], (row - first_row) * top_k + i); - } else { - // A full table: push the union count past the cap so the check - // below traps. - atomicAdd(&s_counters[1], hash_size); - } + if (old == g) { + inserted = true; + break; } + h = next_slot(h, table_size); } - __syncthreads(); - if (tid == 0) { - if (valid_counts != nullptr) { - valid_counts[static_cast(row) * valid_count_stride] = - s_counters[2]; - } - s_counters[2] = 0; + if (inserted) { + atomicMin(&s_hash_vals[h], position); + } else { + // A full table: push the union count past the cap so the check + // below traps. + atomicAdd(&s_counters[1], max_union_rows + 1); } - __syncthreads(); } - // The table and eviction slots are sized for max_union_rows (hash_size - 1); - // a larger union would never reach an empty probe slot or run past s_lru_out. - if (s_counters[1] >= hash_size) { + __syncthreads(); + // Hash capacity can exceed the union bound; eviction slots cannot. + if (s_counters[1] > max_union_rows) { if (tid == 0) { printf( "HiSparse residency: request %d resolves %d host rows, above the " "configured bound of %d\n", - request_row, s_counters[1], hash_size - 1); + request_row, s_counters[1], max_union_rows); } __trap(); } - // Fully resident rows need only request-relative page translation. Avoid - // scanning or rewriting the hot LRU when no selected row can consult it. - if (s_counters[1] == 0) { - for (int row = first_row + tid; row < end_row; row += blockDim.x) { - if (swap_counts != nullptr) swap_counts[row] = 0; + + // Phase 2: scan hot slots in slot order so device_global_indices is read + // coalesced, recording each hit's slot on its union entry and in s_done. + constexpr int kSlotBatch = 4; + for (int base = 0; base < hot_size; base += kSlotBatch * blockDim.x) { + int32_t cached_g[kSlotBatch]; +#pragma unroll + for (int b = 0; b < kSlotBatch; ++b) { + const int slot = base + b * blockDim.x + tid; + cached_g[b] = slot < hot_size ? row_dgi[slot] : -1; + } +#pragma unroll + for (int b = 0; b < kSlotBatch; ++b) { + const int slot = base + b * blockDim.x + tid; + // Every lane makes its first probe, then the warp iterates together + // until no lane is still probing, so the warp is converged at the + // ballot. CUDA 12.8-13.2.1 ptxas can fail to reconverge a per-lane + // probe loop before a ballot on sm_100. + const int32_t g = cached_g[b]; + int h = hash_slot(g, table_size); + int32_t k = s_hash_keys[h]; + bool is_hit = g >= 0 && k == g; + bool probing = g >= 0 && !is_hit && k != kHashEmpty; + while (__any_sync(0xFFFFFFFF, probing)) { + if (probing) h = next_slot(h, table_size); + k = probing ? s_hash_keys[h] : kHashEmpty; + const bool found = probing && k == g; + is_hit |= found; + probing = probing && !found && k != kHashEmpty; + } + if (is_hit) s_hash_vals[h] = kEntryHit | slot; + const unsigned int hit_bits = __ballot_sync(0xFFFFFFFF, is_hit); + if (lane_id == 0 && slot < hot_size) { + s_done[slot / kWarpSize] = hit_bits; + } } - return; } + __syncthreads(); - // Phase 2: walk hot slots in LRU order, classify hit / evictable, and - // compact them (hits forward, evictables backward) into s_lru_out. A hit - // records its slot on the union entry. + // Walk hot slots in LRU order, classify hit / evictable, and compact them + // (hits forward, evictables backward) into s_lru_out. const int iters_buffer = (num_buffer_chunks + NUM_WARPS - 1) / NUM_WARPS; int total_hit_count = 0; int total_evict_count = 0; + int16_t next_lru = tid < hot_size ? row_lru[tid] : int16_t(-1); for (int iter = 0; iter < iters_buffer; iter++) { const int chunk_idx = warp_id + iter * NUM_WARPS; const bool has_valid_chunk = chunk_idx < num_buffer_chunks; const int pos = chunk_idx * kWarpSize + lane_id; const bool has_valid_pos = has_valid_chunk && (pos < hot_size); - const int16_t slot = has_valid_pos ? row_lru[pos] : int16_t(-1); - // Corruption tripwire: lru/dgi are long-lived device state; if some - // external writer (e.g. stray RDMA into reused VRAM) corrupts a slot - // out of [0, region_stride), treat it as no-hit so it degrades to a - // re-miss instead of an unbounded dgi read (phase 3 bounds the write - // side). - const int32_t cached_g = - (slot >= 0 && slot < region_stride) ? row_dgi[slot] : -1; - - bool is_hit = false; - if (cached_g >= 0) { - int h = hash_slot(cached_g, table_size); - while (true) { - const int32_t k = s_hash_keys[h]; - if (k == cached_g) { - s_hash_vals[h] = kEntryHit | slot; - is_hit = true; - break; - } - if (k == kHashEmpty) break; - h = next_slot(h, table_size); - } - } + const int16_t slot = has_valid_pos ? next_lru : int16_t(-1); + const int next_pos = pos + NUM_WARPS * kWarpSize; + next_lru = next_pos < hot_size ? row_lru[next_pos] : int16_t(-1); + // Corruption tripwire: lru is long-lived device state; an out-of-range + // slot is treated as no-hit so it degrades to a re-miss (phase 3 bounds + // the write side). + const bool is_hit = slot >= 0 && slot < hot_size && + (s_done[slot / kWarpSize] >> (slot % kWarpSize)) & 1u; const bool is_evictable = has_valid_pos && !is_hit; int local_hit_off = 0; @@ -580,7 +686,6 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( // device_global_indices. Swaps fill the rows' swap segments in order; the // union never exceeds their combined capacity, and position chunks never // exceed buffer chunks because hot_size covers every row's top-k. - const int num_positions = (end_row - first_row) * top_k; const int num_position_chunks = (num_positions + kWarpSize - 1) / kWarpSize; const int iters_position = (num_position_chunks + NUM_WARPS - 1) / NUM_WARPS; int miss_running_total = 0; @@ -596,17 +701,10 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( int32_t entry = 0; bool is_miss = false; if (has_valid_chunk && position < num_positions) { - const int row = first_row + position / top_k; - int32_t resident_row; - g = translate_topk_entry( - global_indices[static_cast(row) * input_row_stride + - position % top_k], - request_row, state_row, source_block_table, resident_block_table, - host_rows, source_bt_stride, source_num_reqs, source_num_blocks, - source_block_size, resident_bt_stride, resident_num_reqs, - resident_num_blocks, resident_block_size, resident_null_block, - resident_row); - if (resident_row < 0 && g >= 0) { + const int32_t stashed = + hot_indices[static_cast(first_row) * top_k + position]; + if (stashed < -1) { + g = encode_host_row(stashed); h = hash_slot(g, table_size); while (s_hash_keys[h] != g) { h = next_slot(h, table_size); @@ -703,16 +801,10 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel( continue; } const int row = first_row + position / top_k; - int32_t resident_row; - const int32_t g = translate_topk_entry( - global_indices[static_cast(row) * input_row_stride + - position % top_k], - request_row, state_row, source_block_table, resident_block_table, - host_rows, source_bt_stride, source_num_reqs, source_num_blocks, - source_block_size, resident_bt_stride, resident_num_reqs, - resident_num_blocks, resident_block_size, resident_null_block, - resident_row); - if (resident_row >= 0 || g < 0) continue; + const int32_t stashed = + hot_indices[static_cast(first_row) * top_k + position]; + if (stashed > -2) continue; + const int32_t g = encode_host_row(stashed); int h = hash_slot(g, table_size); while (s_hash_keys[h] != g) { h = next_slot(h, table_size); @@ -848,30 +940,42 @@ __global__ void hisparse_gather_compact_kernel( } } +// One block per (written token, index group): drop the group's hot copy of +// the token's rewritten host slot. __global__ void hisparse_invalidate_written_slots_kernel( - int32_t* __restrict__ device_global_indices, + int32_t* const* __restrict__ device_global_indices, // [num_groups] const int32_t* __restrict__ request_state_indices, const int32_t* __restrict__ req_id_per_token, - const int64_t* __restrict__ written_slots, const int64_t num_tokens, - const int64_t num_request_ids, const int64_t num_state_rows, - const int64_t region_stride) { + const int64_t* __restrict__ written_slots, const int64_t num_request_ids, + const int64_t num_state_rows, const int64_t region_stride) { const int64_t token_idx = blockIdx.x; - if (token_idx >= num_tokens) { - return; - } const int32_t req_idx = req_id_per_token[token_idx]; if (req_idx < 0 || req_idx >= num_request_ids) { return; } const int32_t state_idx = request_state_indices[req_idx]; const int64_t written_slot = written_slots[token_idx]; - if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0) { + if (state_idx < 0 || state_idx >= num_state_rows || written_slot < 0 || + written_slot > INT32_MAX) { + return; + } + const int32_t slot = static_cast(written_slot); + int32_t* row = device_global_indices[blockIdx.y] + state_idx * region_stride; + if (region_stride % 4 == 0 && (reinterpret_cast(row) & 15) == 0) { + const int4* row4 = reinterpret_cast(row); + for (int64_t offset = threadIdx.x; offset < region_stride / 4; + offset += blockDim.x) { + const int4 v = row4[offset]; + if (v.x == slot) row[4 * offset] = -1; + if (v.y == slot) row[4 * offset + 1] = -1; + if (v.z == slot) row[4 * offset + 2] = -1; + if (v.w == slot) row[4 * offset + 3] = -1; + } return; } - int32_t* row = device_global_indices + state_idx * region_stride; for (int64_t offset = threadIdx.x; offset < region_stride; offset += blockDim.x) { - if (row[offset] == written_slot) { + if (row[offset] == slot) { row[offset] = -1; } } @@ -889,20 +993,22 @@ int64_t check_2d_rows(const torch::stable::Tensor& t, const char* name, } // namespace void hisparse_invalidate_written_slots( - torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& device_global_indices_ptrs, torch::stable::Tensor const& request_state_indices, torch::stable::Tensor const& req_id_per_token, - torch::stable::Tensor const& written_slots) { - STD_TORCH_CHECK(device_global_indices.is_cuda() && + torch::stable::Tensor const& written_slots, int64_t num_state_rows, + int64_t region_stride) { + STD_TORCH_CHECK(device_global_indices_ptrs.is_cuda() && request_state_indices.is_cuda() && req_id_per_token.is_cuda() && written_slots.is_cuda(), "HiSparse invalidation tensors must be on CUDA"); STD_TORCH_CHECK( - device_global_indices.scalar_type() == - torch::headeronly::ScalarType::Int && - device_global_indices.dim() == 2 && - device_global_indices.is_contiguous(), - "device_global_indices must be a contiguous 2D int32 CUDA tensor"); + device_global_indices_ptrs.scalar_type() == + torch::headeronly::ScalarType::Long && + device_global_indices_ptrs.dim() == 1 && + device_global_indices_ptrs.is_contiguous(), + "device_global_indices_ptrs must be a contiguous 1D int64 CUDA tensor " + "of [max_num_reqs, region_stride] int32 table addresses"); STD_TORCH_CHECK( request_state_indices.scalar_type() == torch::headeronly::ScalarType::Int && @@ -919,7 +1025,7 @@ void hisparse_invalidate_written_slots( "written_slots must be a contiguous 1D int64 CUDA tensor"); STD_TORCH_CHECK(req_id_per_token.numel() == written_slots.numel(), "req_id_per_token and written_slots must have equal length"); - const int device_index = device_global_indices.get_device_index(); + const int device_index = device_global_indices_ptrs.get_device_index(); STD_TORCH_CHECK( request_state_indices.get_device_index() == device_index && req_id_per_token.get_device_index() == device_index && @@ -927,24 +1033,25 @@ void hisparse_invalidate_written_slots( "HiSparse invalidation tensors must be on the same CUDA device"); const int64_t num_tokens = written_slots.numel(); - const int64_t num_state_rows = device_global_indices.size(0); - const int64_t region_stride = device_global_indices.size(1); - if (num_tokens == 0 || num_state_rows == 0 || region_stride == 0) { + const int64_t num_groups = device_global_indices_ptrs.numel(); + if (num_tokens == 0 || num_groups == 0 || num_state_rows == 0 || + region_stride == 0) { return; } - STD_TORCH_CHECK(num_tokens <= INT32_MAX, + STD_TORCH_CHECK(num_tokens <= INT32_MAX && num_groups <= 65535, "HiSparse invalidation grid exceeds CUDA limits"); constexpr int kBlockSize = 256; const torch::stable::accelerator::DeviceGuard device_guard(device_index); const cudaStream_t stream = get_current_cuda_stream(); - hisparse_invalidate_written_slots_kernel<<(num_tokens), - kBlockSize, 0, stream>>>( - device_global_indices.mutable_data_ptr(), + const dim3 grid(static_cast(num_tokens), + static_cast(num_groups)); + hisparse_invalidate_written_slots_kernel<<>>( + static_cast(device_global_indices_ptrs.const_data_ptr()), request_state_indices.const_data_ptr(), req_id_per_token.const_data_ptr(), - written_slots.const_data_ptr(), num_tokens, - request_state_indices.numel(), num_state_rows, region_stride); + written_slots.const_data_ptr(), request_state_indices.numel(), + num_state_rows, region_stride); const cudaError_t launch_error = cudaGetLastError(); STD_TORCH_CHECK(launch_error == cudaSuccess, "HiSparse invalidation kernel launch failed: ", @@ -973,7 +1080,8 @@ void hisparse_resolve_residency( std::optional const& swap_device_physical_rows, std::optional const& swap_counts, std::optional const& resident_block_table, - int64_t resident_block_size, int64_t resident_null_block) { + int64_t resident_block_size, int64_t resident_null_block, + std::optional const& num_valid_rows) { STD_TORCH_CHECK( host_cache.device().is_cpu() && is_pinned_cpu_tensor(host_cache), "host_cache must be pinned CPU memory"); @@ -1067,6 +1175,18 @@ void hisparse_resolve_residency( source_num_reqs = static_cast(table.size(0)); source_num_blocks = static_cast(table.size(1)); } + const int32_t* num_valid_rows_ptr = nullptr; + if (num_valid_rows.has_value()) { + auto const& valid_rows = num_valid_rows.value(); + STD_TORCH_CHECK(request_ids.has_value(), + "num_valid_rows requires request_ids"); + STD_TORCH_CHECK( + valid_rows.is_cuda() && + valid_rows.scalar_type() == torch::headeronly::ScalarType::Int && + valid_rows.numel() == 1, + "num_valid_rows must be a one-element int32 CUDA tensor"); + num_valid_rows_ptr = valid_rows.const_data_ptr(); + } const int32_t required_hot_rows = request_ids_ptr != nullptr ? source_num_reqs : launch_rows; STD_TORCH_CHECK(hot_block_table.size(0) >= required_hot_rows, @@ -1213,13 +1333,12 @@ void hisparse_resolve_residency( STD_TORCH_CHECK(max_union_rows > 0 && max_union_rows <= hot_size, "max_union_rows must be in [1, hot_size], got ", max_union_rows); - // One spare entry above the union bound keeps every probe sequence finite. - const int hash_size = static_cast(max_union_rows) + 1; const int num_buffer_chunks = (hot_size + kWarpSize - 1) / kWarpSize; - const size_t smem_bytes = - sizeof(int32_t) * (2 * hash_size + 2 * (num_buffer_chunks + 1) + - kResidencyCounters + num_buffer_chunks) + + const size_t fixed_smem_bytes = + sizeof(int32_t) * (3 * num_buffer_chunks + 2 + kResidencyCounters) + sizeof(int16_t) * hot_size; + const size_t min_smem_bytes = + fixed_smem_bytes + 2 * sizeof(int32_t) * (max_union_rows + 1); const torch::stable::accelerator::DeviceGuard device_guard( hot_cache.get_device_index()); @@ -1228,26 +1347,52 @@ void hisparse_resolve_residency( &max_smem_bytes, cudaDevAttrMaxSharedMemoryPerBlockOptin, hot_cache.get_device_index()) == cudaSuccess, "failed to query shared memory per block"); - STD_TORCH_CHECK(smem_bytes <= static_cast(max_smem_bytes), + STD_TORCH_CHECK(min_smem_bytes <= static_cast(max_smem_bytes), "HiSparse device_buffer_size=", hot_size, " needs ", - smem_bytes, + min_smem_bytes, " bytes of shared memory to resolve residency, " "but the device allows ", max_smem_bytes, "; lower device_buffer_size or the number of speculative " "tokens."); + // Keep the table sparse when possible without rejecting configurations + // that fit the minimum table with one spare entry above the union bound. + const int hash_size = + std::min(kHashLoadInverse * static_cast(max_union_rows) + 1, + static_cast((max_smem_bytes - fixed_smem_bytes) / + (2 * sizeof(int32_t)))); + const size_t smem_bytes = fixed_smem_bytes + 2 * sizeof(int32_t) * hash_size; const cudaStream_t stream = get_current_cuda_stream(); const int64_t attention_row_stride = attention_indices.has_value() ? attention_indices.value().stride(0) : 0; const int64_t valid_count_stride = valid_counts.has_value() ? valid_counts.value().stride(0) : 0; - auto kernel = hisparse_resolve_residency_kernel; + // Block arithmetic dominates the all-resident path; specialize the + // ubiquitous 64-row block. + constexpr int kFixedBlockSize = 64; + const bool fixed_block_size = hot_block_size == kFixedBlockSize && + (!source_block_table.has_value() || + source_block_size == kFixedBlockSize) && + (!resident_block_table.has_value() || + resident_block_size == kFixedBlockSize); + auto kernel = fixed_block_size + ? hisparse_resolve_residency_kernel + : hisparse_resolve_residency_kernel<0>; if (smem_bytes > 48 * 1024) { - const cudaError_t attribute_error = cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); - STD_TORCH_CHECK(attribute_error == cudaSuccess, - "failed to configure HiSparse swap-in shared memory: ", - cudaGetErrorString(attribute_error)); + static std::vector configured_smem_bytes[2]; + std::vector& configured = configured_smem_bytes[fixed_block_size]; + const size_t device_index = hot_cache.get_device_index(); + if (configured.size() <= device_index) { + configured.resize(device_index + 1, 0); + } + if (configured[device_index] < smem_bytes) { + const cudaError_t attribute_error = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); + STD_TORCH_CHECK(attribute_error == cudaSuccess, + "failed to configure HiSparse swap-in shared memory: ", + cudaGetErrorString(attribute_error)); + configured[device_index] = smem_bytes; + } } kernel<<>>( hot_block_table.const_data_ptr(), @@ -1260,13 +1405,13 @@ void hisparse_resolve_residency( device_global_indices.mutable_data_ptr(), lru_slots.mutable_data_ptr(), stats_ptr, request_state_ptr, request_state_count, host_rows, hot_block_table.stride(0), hot_block_size, - top_k, hot_size, hash_size, region_stride, attention_block_stride, - source_bt_stride, source_num_reqs, source_num_blocks, - static_cast(source_block_size), resident_bt_stride, - resident_num_reqs, resident_num_blocks, + top_k, hot_size, hash_size, static_cast(max_union_rows), + region_stride, attention_block_stride, source_bt_stride, source_num_reqs, + source_num_blocks, static_cast(source_block_size), + resident_bt_stride, resident_num_reqs, resident_num_blocks, static_cast(resident_block_size), static_cast(resident_null_block), global_indices.stride(0), - attention_row_stride, valid_count_stride); + attention_row_stride, valid_count_stride, num_valid_rows_ptr); const cudaError_t launch_error = cudaGetLastError(); STD_TORCH_CHECK(launch_error == cudaSuccess, "HiSparse residency kernel launch failed: ", diff --git a/csrc/moe/grouped_topk_kernels.cu b/csrc/moe/grouped_topk_kernels.cu index ea89aa38a1f0..443cdd14dc69 100644 --- a/csrc/moe/grouped_topk_kernels.cu +++ b/csrc/moe/grouped_topk_kernels.cu @@ -479,9 +479,9 @@ __device__ void topk_with_k2(T* output, T const* input, BiasT const* bias, template __global__ void grouped_topk_fused_kernel( T* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, - int64_t const num_tokens, int64_t const num_experts, int64_t const n_group, - int64_t const topk_group, int64_t const topk, bool renormalize, - double routed_scaling_factor) { + bool const* is_padding, int64_t const num_tokens, int64_t const num_experts, + int64_t const n_group, int64_t const topk_group, int64_t const topk, + bool renormalize, double routed_scaling_factor) { int32_t const token_id = static_cast(blockIdx.x); if (token_id >= num_tokens) { return; @@ -564,8 +564,11 @@ __global__ void grouped_topk_fused_kernel( if (!proceed) { for (int i = lane_id; i < topk_i32; i += WARP_SIZE) { - topk_indices[i] = static_cast(i); - topk_values[i] = 1.0f / static_cast(topk_i32); + topk_indices[i] = + is_padding && is_padding[token_id] ? -1 : static_cast(i); + topk_values[i] = is_padding && is_padding[token_id] + ? 0.0f + : 1.0f / static_cast(topk_i32); } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); @@ -624,8 +627,9 @@ __global__ void grouped_topk_fused_kernel( } if (lane_id < topk_i32) { - topk_indices[lane_id] = lane_idx; - topk_values[lane_id] = lane_unbiased * scale; + topk_indices[lane_id] = is_padding && is_padding[token_id] ? -1 : lane_idx; + topk_values[lane_id] = + is_padding && is_padding[token_id] ? 0.0f : lane_unbiased * scale; } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -638,8 +642,8 @@ template __global__ void grouped_topk_fused_small_expert_count_kernel( T* scores, float* topkValues, IdxT* topkIndices, BiasT const* routingBias, - int64_t const numTokens, int64_t const numGroup, int64_t const topkGroup, - int64_t const topk, int64_t const numExperts, + bool const* is_padding, int64_t const numTokens, int64_t const numGroup, + int64_t const topkGroup, int64_t const topk, int64_t const numExperts, int64_t const numExpertsPerGroup, bool const renormalize, double const routedScalingFactor) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -838,8 +842,10 @@ __global__ void grouped_topk_fused_small_expert_count_kernel( } // store the topk scores and experts to output if (laneIdx < topk) { - topkValues[laneIdx] = finalScore; - topkIndices[laneIdx] = expertIdx; + topkValues[laneIdx] = + is_padding && is_padding[blockIdx.x] ? 0.0f : finalScore; + topkIndices[laneIdx] = + is_padding && is_padding[blockIdx.x] ? -1 : expertIdx; } } @@ -890,7 +896,8 @@ __device__ __forceinline__ void write_outputs( cg::thread_block_tile const& warp, float lane_selection_score, float lane_unbiased, int32_t lane_expert, int32_t lane, int32_t token, int32_t topk, float* topk_values, IdxT* topk_indices, bool renormalize, - float routed_scaling_factor) { + float routed_scaling_factor, bool const* is_padding) { + bool const padded = is_padding && is_padding[token]; bool const finite_selection = lane < topk && lane_selection_score != InvalidScore; lane_unbiased = finite_selection ? lane_unbiased : 0.0F; @@ -912,8 +919,8 @@ __device__ __forceinline__ void write_outputs( } int64_t const output_index = int64_t{token} * topk + lane; - topk_values[output_index] = output; - topk_indices[output_index] = static_cast(lane_expert); + topk_values[output_index] = padded ? 0.0F : output; + topk_indices[output_index] = static_cast(padded ? -1 : lane_expert); } } @@ -922,8 +929,8 @@ template = 900) @@ -1095,9 +1102,9 @@ template ::BlockDim) single_group_topk_warp_kernel(T const* scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, - int64_t num_tokens, int64_t num_experts, - int64_t topk, bool renormalize, - float routed_scaling_factor, + bool const* is_padding, int64_t num_tokens, + int64_t num_experts, int64_t topk, + bool renormalize, float routed_scaling_factor, bool enable_pdl) { static constexpr int NumChunks = (MaxNumExperts + WARP_SIZE - 1) / WARP_SIZE; static constexpr int WarpBlockDim = @@ -1166,7 +1173,7 @@ __global__ void __launch_bounds__(WarpTopKLaunchConfig::BlockDim) } write_outputs(warp, lane_score, lane_unbiased, lane_expert, lane, token, topk_i32, topk_values, topk_indices, renormalize, - routed_scaling_factor); + routed_scaling_factor, is_padding); } #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) @@ -1198,9 +1205,10 @@ using PrecomputedSoftmaxBiasTiers = template void launch(T* scores, float* topk_values, IdxT* topk_indices, - BiasT const* bias, int64_t num_tokens, int64_t num_experts, - int64_t topk, bool renormalize, double routed_scaling_factor, - bool enable_pdl, cudaLaunchConfig_t& config) { + BiasT const* bias, bool const* is_padding, int64_t num_tokens, + int64_t num_experts, int64_t topk, bool renormalize, + double routed_scaling_factor, bool enable_pdl, + cudaLaunchConfig_t& config) { config.dynamicSmemBytes = 0; bool const use_block_kernel = UseTunedBlockPath || @@ -1213,8 +1221,8 @@ void launch(T* scores, float* topk_values, IdxT* topk_indices, &config, &single_group_topk_block_kernel, - scores, topk_values, topk_indices, bias, num_experts, topk, renormalize, - static_cast(routed_scaling_factor), enable_pdl); + scores, topk_values, topk_indices, bias, is_padding, num_experts, topk, + renormalize, static_cast(routed_scaling_factor), enable_pdl); } else { using WarpConfig = WarpTopKLaunchConfig; config.gridDim = WarpConfig::grid_dim(num_tokens); @@ -1223,54 +1231,58 @@ void launch(T* scores, float* topk_values, IdxT* topk_indices, &config, &single_group_topk_warp_kernel, - scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, - renormalize, static_cast(routed_scaling_factor), enable_pdl); + scores, topk_values, topk_indices, bias, is_padding, num_tokens, + num_experts, topk, renormalize, + static_cast(routed_scaling_factor), enable_pdl); } } template -bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, int64_t, int64_t, - int64_t, bool, double, bool, cudaLaunchConfig_t&) { +bool dispatch(TierList<>*, T*, float*, IdxT*, BiasT const*, bool const*, + int64_t, int64_t, int64_t, bool, double, bool, + cudaLaunchConfig_t&) { return false; } template bool dispatch(TierList*, T* scores, float* topk_values, - IdxT* topk_indices, BiasT const* bias, int64_t num_tokens, - int64_t num_experts, int64_t topk, bool renormalize, - double routed_scaling_factor, bool enable_pdl, + IdxT* topk_indices, BiasT const* bias, bool const* is_padding, + int64_t num_tokens, int64_t num_experts, int64_t topk, + bool renormalize, double routed_scaling_factor, bool enable_pdl, cudaLaunchConfig_t& config) { if (num_experts <= First::kExperts && topk <= First::kTopK) { launch( - scores, topk_values, topk_indices, bias, num_tokens, num_experts, topk, - renormalize, routed_scaling_factor, enable_pdl, config); + scores, topk_values, topk_indices, bias, is_padding, num_tokens, + num_experts, topk, renormalize, routed_scaling_factor, enable_pdl, + config); return true; } return dispatch( static_cast*>(nullptr), scores, topk_values, - topk_indices, bias, num_tokens, num_experts, topk, renormalize, - routed_scaling_factor, enable_pdl, config); + topk_indices, bias, is_padding, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); } } // namespace detail template bool invoke(T* scores, float* topk_values, IdxT* topk_indices, - BiasT const* bias, int64_t num_tokens, int64_t num_experts, - int64_t topk, bool renormalize, double routed_scaling_factor, - bool enable_pdl, cudaLaunchConfig_t& config) { + BiasT const* bias, bool const* is_padding, int64_t num_tokens, + int64_t num_experts, int64_t topk, bool renormalize, + double routed_scaling_factor, bool enable_pdl, + cudaLaunchConfig_t& config) { static_assert(SF == SCORING_NONE || SF == SCORING_SIGMOID); if constexpr (SF == SCORING_SIGMOID) { return detail::dispatch( static_cast(nullptr), scores, topk_values, - topk_indices, bias, num_tokens, num_experts, topk, renormalize, - routed_scaling_factor, enable_pdl, config); + topk_indices, bias, is_padding, num_tokens, num_experts, topk, + renormalize, routed_scaling_factor, enable_pdl, config); } else { return detail::dispatch( static_cast(nullptr), scores, - topk_values, topk_indices, bias, num_tokens, num_experts, topk, - renormalize, routed_scaling_factor, enable_pdl, config); + topk_values, topk_indices, bias, is_padding, num_tokens, num_experts, + topk, renormalize, routed_scaling_factor, enable_pdl, config); } } @@ -1278,10 +1290,11 @@ bool invoke(T* scores, float* topk_values, IdxT* topk_indices, template void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, - BiasT const* bias, int64_t const num_tokens, - int64_t const num_experts, int64_t const n_group, - int64_t const topk_group, int64_t const topk, - bool const renormalize, double const routed_scaling_factor, + BiasT const* bias, bool const* is_padding, + int64_t const num_tokens, int64_t const num_experts, + int64_t const n_group, int64_t const topk_group, + int64_t const topk, bool const renormalize, + double const routed_scaling_factor, const bool enable_pdl = false, cudaStream_t const stream = 0) { cudaLaunchConfig_t config; @@ -1293,8 +1306,9 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, config.attrs = attrs; if (n_group == 1 && topk_group == 1 && single_group_topk::invoke( - scores, topk_values, topk_indices, bias, num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, enable_pdl, config)) { + scores, topk_values, topk_indices, bias, is_padding, num_tokens, + num_experts, topk, renormalize, routed_scaling_factor, enable_pdl, + config)) { return; } @@ -1346,9 +1360,9 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, config.blockDim = num_threads; config.dynamicSmemBytes = 0; cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, - topk_indices, bias, num_tokens, n_group, topk_group, - topk, num_experts, num_experts / n_group, renormalize, - routed_scaling_factor); + topk_indices, bias, is_padding, num_tokens, n_group, + topk_group, topk, num_experts, num_experts / n_group, + renormalize, routed_scaling_factor); } else { auto* kernel_instance = &grouped_topk_fused_kernel; // One block per token; one warp per group. @@ -1366,18 +1380,20 @@ void invokeNoAuxTc(T* scores, float* topk_values, IdxT* topk_indices, size_t const extra_bytes = 16 + static_cast(n_group) * sizeof(T); config.dynamicSmemBytes = internal_bytes + extra_bytes; cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, - topk_indices, bias, num_tokens, num_experts, n_group, - topk_group, topk, renormalize, routed_scaling_factor); + topk_indices, bias, is_padding, num_tokens, num_experts, + n_group, topk_group, topk, renormalize, + routed_scaling_factor); } } #define INSTANTIATE_NOAUX_TC(T, BiasT, IdxT, SF) \ template void invokeNoAuxTc( \ T * scores, float* topk_values, IdxT* topk_indices, BiasT const* bias, \ - int64_t const num_tokens, int64_t const num_experts, \ - int64_t const n_group, int64_t const topk_group, int64_t const topk, \ - bool const renormalize, double const routed_scaling_factor, \ - const bool enable_pdl, cudaStream_t const stream); + bool const* is_padding, int64_t const num_tokens, \ + int64_t const num_experts, int64_t const n_group, \ + int64_t const topk_group, int64_t const topk, bool const renormalize, \ + double const routed_scaling_factor, const bool enable_pdl, \ + cudaStream_t const stream); INSTANTIATE_NOAUX_TC(float, float, int32_t, SCORING_SIGMOID); INSTANTIATE_NOAUX_TC(float, half, int32_t, SCORING_SIGMOID); @@ -1403,7 +1419,8 @@ INSTANTIATE_NOAUX_TC(__nv_bfloat16, __nv_bfloat16, int32_t, SCORING_NONE); std::tuple grouped_topk( torch::stable::Tensor const& scores, int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize, double routed_scaling_factor, - torch::stable::Tensor const& bias, int64_t scoring_func = 0) { + torch::stable::Tensor const& bias, int64_t scoring_func, + const std::optional& is_padding) { const auto data_type = scores.scalar_type(); const auto bias_type = bias.scalar_type(); STD_TORCH_CHECK(scores.dim() == 2, "scores must be a 2D Tensor"); @@ -1426,6 +1443,20 @@ std::tuple grouped_topk( scoring_func == vllm::moe::SCORING_SIGMOID, "scoring_func must be SCORING_NONE (0) or SCORING_SIGMOID (1)"); + const bool* is_padding_ptr = nullptr; + if (is_padding.has_value()) { + STD_TORCH_CHECK( + is_padding->scalar_type() == torch::headeronly::ScalarType::Bool, + "is_padding must be bool"); + STD_TORCH_CHECK(is_padding->dim() == 1 && is_padding->size(0) == num_tokens, + "is_padding must have one entry per token"); + STD_TORCH_CHECK(is_padding->is_contiguous(), + "is_padding must be contiguous"); + STD_TORCH_CHECK(is_padding->device() == scores.device(), + "is_padding must be on the scores device"); + is_padding_ptr = is_padding->const_data_ptr(); + } + // Always output float32 for topk_values (eliminates Python-side conversion) auto topk_values = torch::stable::new_empty( scores, {num_tokens, topk}, torch::headeronly::ScalarType::Float); @@ -1447,8 +1478,8 @@ std::tuple grouped_topk( reinterpret_cast(scores.mutable_data_ptr()), \ reinterpret_cast(topk_values.mutable_data_ptr()), \ reinterpret_cast(topk_indices.mutable_data_ptr()), \ - reinterpret_cast(bias.data_ptr()), num_tokens, \ - num_experts, n_group, topk_group, topk, renormalize, \ + reinterpret_cast(bias.data_ptr()), is_padding_ptr, \ + num_tokens, num_experts, n_group, topk_group, topk, renormalize, \ routed_scaling_factor, pdl_flag, stream); \ break; \ case vllm::moe::SCORING_SIGMOID: \ @@ -1456,8 +1487,8 @@ std::tuple grouped_topk( reinterpret_cast(scores.mutable_data_ptr()), \ reinterpret_cast(topk_values.mutable_data_ptr()), \ reinterpret_cast(topk_indices.mutable_data_ptr()), \ - reinterpret_cast(bias.data_ptr()), num_tokens, \ - num_experts, n_group, topk_group, topk, renormalize, \ + reinterpret_cast(bias.data_ptr()), is_padding_ptr, \ + num_tokens, num_experts, n_group, topk_group, topk, renormalize, \ routed_scaling_factor, pdl_flag, stream); \ break; \ default: \ diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h index 8ab24fcb33b8..484ae2a86957 100644 --- a/csrc/moe/moe_ops.h +++ b/csrc/moe/moe_ops.h @@ -20,6 +20,15 @@ void topk_sigmoid(torch::stable::Tensor& topk_weights, double routed_scaling_factor, std::optional is_padding); +#ifdef USE_CUDA +void muse_spark_routing(torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_ids, + const torch::stable::Tensor& logits, + const torch::stable::Tensor& bias, double multiplier, + double route_scale, double division_eps, + const std::optional& is_padding); +#endif + void topk_softplus_sqrt( torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, torch::stable::Tensor& token_expert_indices, @@ -82,7 +91,8 @@ torch::stable::Tensor moe_wna16_gemm( std::tuple grouped_topk( const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group, int64_t topk, bool renormalize, double routed_scaling_factor, - const torch::stable::Tensor& bias, int64_t scoring_func); + const torch::stable::Tensor& bias, int64_t scoring_func, + const std::optional& is_padding); #endif bool moe_permute_unpermute_supported(); diff --git a/csrc/moe/muse_spark_routing.cu b/csrc/moe/muse_spark_routing.cu new file mode 100644 index 000000000000..0c77237be32f --- /dev/null +++ b/csrc/moe/muse_spark_routing.cu @@ -0,0 +1,346 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Adapted from SGLang's Muse Spark CUDA routing implementation. + +#include +#include +#include +#include +#include +#include +#include + +#include "dispatch_utils.h" +#include "torch_utils.h" + +namespace vllm { +namespace moe { + +__device__ __forceinline__ float warp_max(float value) { +#if defined(__CUDA_ARCH_FAMILY_SPECIFIC__) && __CUDA_ARCH__ >= 1000 && \ + __CUDA_ARCH__ < 1100 + float result; + asm("redux.sync.max.f32 %0, %1, %2;" + : "=f"(result) + : "f"(value), "r"(0xffffffffu)); + return result; +#else + #pragma unroll + for (int offset = 16; offset; offset /= 2) + value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset)); + return value; +#endif +} + +inline constexpr float kNegInf = -__builtin_huge_valf(); +// fp32(log2(e)) == 0x3FB8AA3B, the constant Triton folds into tl.exp. +inline constexpr float kLog2e = 1.4426950408889634f; + +__device__ __forceinline__ float ex2_approx(float x) { + float r; + asm("ex2.approx.f32 %0, %1;" : "=f"(r) : "f"(x)); + return r; +} + +// Triton lowers `a / b` on fp32 to div.full.f32, not the div.rn.f32 nvcc emits. +__device__ __forceinline__ float div_full(float a, float b) { + float r; + asm("div.full.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +// triton/language/standard.py:50 -- `1 / (1 + exp(-x))`, lowered as +// sub.f32 / mul.f32 / ex2.approx.f32 / add.f32 / div.full.f32. +__device__ __forceinline__ float triton_sigmoid(float x) { + return div_full(1.0f, ex2_approx((0.0f - x) * kLog2e) + 1.0f); +} + +// kExpertsPerLane must be a multiple of 4 (128-bit row loads) and covers +// kExpertsPerLane * 32 experts. Lane `l` holds expert v*128 + l*4 + j at +// register slot v*4 + j, monotone in the slot -- both the stable sort and the +// group-ballot tie-break rely on that to reproduce tl.argmax. +template +struct RoutingTrait { + static constexpr int kTopK = 8; + static constexpr int kExpertsPerLane = 8; + static constexpr int kVecsPerLane = 2; + static constexpr int kWarpsPerBlock = kWarpsPerBlock_; + static constexpr int kMaxExperts = 256; +}; + +template +struct RoutingParams { + const float* __restrict__ gating_output; + const float* __restrict__ bias; + float* __restrict__ topk_weights; + Index* __restrict__ topk_ids; + const bool* __restrict__ is_padding; + int64_t padding_stride; + int num_tokens; + int num_experts; + int64_t stride_gm; + int64_t stride_wm; + int64_t stride_im; + float router_output_multiplier; + float route_scale; + float division_eps; +}; + +template +__global__ void routing_kernel( + const RoutingParams __grid_constant__ params) { + constexpr int kEPL = Trait::kExpertsPerLane; + constexpr int kTopK = Trait::kTopK; + constexpr int kGroupExperts = 4 * 32; + + const int lane = threadIdx.x % 32; + const int token = + blockIdx.x * Trait::kWarpsPerBlock + static_cast(threadIdx.x / 32); + if (token >= params.num_tokens) return; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if constexpr (kUsePDL) cudaGridDependencySynchronize(); +#endif + if (params.is_padding != nullptr && + params.is_padding[token * params.padding_stride]) { + if (lane < kTopK) { + params.topk_ids[token * params.stride_im + lane] = -1; + params.topk_weights[token * params.stride_wm + lane] = 0.0f; + } + return; + } + + const float* __restrict__ row = + params.gating_output + static_cast(token) * params.stride_gm; + float score[kEPL]; + int expert[kEPL]; + +#pragma unroll + for (int v = 0; v < Trait::kVecsPerLane; ++v) { + const int base = v * kGroupExperts + lane * 4; + float4 g = make_float4(kNegInf, kNegInf, kNegInf, kNegInf); + if (base < params.num_experts) + g = *reinterpret_cast(row + base); + score[v * 4 + 0] = g.x; + score[v * 4 + 1] = g.y; + score[v * 4 + 2] = g.z; + score[v * 4 + 3] = g.w; +#pragma unroll + for (int j = 0; j < 4; ++j) expert[v * 4 + j] = base + j; + } + +#pragma unroll + for (int s = 0; s < kEPL; ++s) score[s] *= params.router_output_multiplier; + +#pragma unroll + for (int s = 0; s < kEPL; ++s) score[s] = triton_sigmoid(score[s]); + + float biased[kEPL]; +#pragma unroll + for (int v = 0; v < Trait::kVecsPerLane; ++v) { + const int base = v * kGroupExperts + lane * 4; + float4 b = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + if (base < params.num_experts) + b = *reinterpret_cast(params.bias + base); + biased[v * 4 + 0] = b.x; + biased[v * 4 + 1] = b.y; + biased[v * 4 + 2] = b.z; + biased[v * 4 + 3] = b.w; + } + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if constexpr (kUsePDL) cudaTriggerProgrammaticLaunchCompletion(); +#endif + + float key[kEPL]; +#pragma unroll + for (int s = 0; s < kEPL; ++s) { + if (expert[s] < params.num_experts) { + key[s] = biased[s] + score[s]; + } else { + score[s] = 0.0f; + key[s] = kNegInf; + } + } + + // Must be a stable sort: equal keys keep their slot order and slots are + // monotone in expert id, which is tl.argmax's lowest-index-wins tie-break. +#pragma unroll + for (int r = 0; r < kEPL; ++r) { +#pragma unroll + for (int i = r % 2; i + 1 < kEPL; i += 2) { + const bool swap = key[i + 1] > key[i]; + const float k0 = swap ? key[i + 1] : key[i]; + const float k1 = swap ? key[i] : key[i + 1]; + const int e0 = swap ? expert[i + 1] : expert[i]; + const int e1 = swap ? expert[i] : expert[i + 1]; + const float s0 = swap ? score[i + 1] : score[i]; + const float s1 = swap ? score[i] : score[i + 1]; + key[i] = k0; + key[i + 1] = k1; + expert[i] = e0; + expert[i + 1] = e1; + score[i] = s0; + score[i + 1] = s1; + } + } + + float score_sum = 0.0f; + int out_id = 0; + float out_score = 0.0f; + +#pragma unroll + for (int k = 0; k < kTopK; ++k) { + const float round_max = warp_max(key[0]); + const bool is_head = key[0] == round_max; + + // Lowest expert id among the tied heads: within one kGroupExperts-wide + // group the id is monotone in the lane, so the lowest set bit of the lowest + // non-empty group's ballot owns the round. + unsigned owners = 0; +#pragma unroll + for (int g = Trait::kVecsPerLane - 1; g >= 0; --g) { + const unsigned m = + __ballot_sync(0xffffffffu, is_head && expert[0] / kGroupExperts == g); + owners = (m != 0u) ? m : owners; + } + const int owner = __ffs(owners) - 1; + const int winner = __shfl_sync(0xffffffffu, expert[0], owner); + const float selected = __shfl_sync(0xffffffffu, score[0], owner); + + // Tapered pop: after round k a lane can be popped at most kTopK-1-k more + // times, so slots past that can never surface and need not be maintained. + const bool pop = expert[0] == winner; +#pragma unroll + for (int s = 0; s + 1 < kEPL; ++s) + if (s + k <= kTopK - 2) { + key[s] = pop ? key[s + 1] : key[s]; + expert[s] = pop ? expert[s + 1] : expert[s]; + score[s] = pop ? score[s + 1] : score[s]; + } + // Only reachable when kTopK > kEPL: a lane popped kEPL times is exhausted + // and must stop winning rounds. + if (kEPL - 1 + k <= kTopK - 2) + key[kEPL - 1] = pop ? kNegInf : key[kEPL - 1]; + + score_sum += selected; + if (lane == k) { + out_id = winner; + out_score = selected; + } + } + + score_sum += params.division_eps; + const float inv_sum_scaled = div_full(params.route_scale, score_sum); + if (lane < kTopK) { + const Index id = out_id; + const float weight = out_score * inv_sum_scaled; + params.topk_ids[token * params.stride_im + lane] = id; + params.topk_weights[token * params.stride_wm + lane] = weight; + } +} + +} // namespace moe +} // namespace vllm + +namespace { +using Tensor = torch::stable::Tensor; +using ScalarType = torch::headeronly::ScalarType; + +template +void launch_routing(const vllm::moe::RoutingParams& params, + cudaStream_t stream) { + using Trait = vllm::moe::RoutingTrait; + cudaLaunchConfig_t config{}; + config.gridDim = dim3((params.num_tokens + Warps - 1) / Warps); + config.blockDim = dim3(32 * Warps); + config.stream = stream; + cudaLaunchAttribute attr{}; + if constexpr (PDL) { + attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr.val.programmaticStreamSerializationAllowed = 1; + config.attrs = &attr; + config.numAttrs = 1; + } + auto err = cudaLaunchKernelEx( + &config, vllm::moe::routing_kernel, params); + STD_TORCH_CHECK(err == cudaSuccess, "Muse Spark routing launch failed: ", + cudaGetErrorString(err)); +} +} // namespace + +void muse_spark_routing(Tensor& topk_weights, Tensor& topk_ids, + const Tensor& logits, const Tensor& bias, + double multiplier, double route_scale, + double division_eps, + const std::optional& is_padding) { + STD_TORCH_CHECK(logits.device().is_cuda() && logits.dim() == 2 && + logits.scalar_type() == ScalarType::Float && + logits.size(1) == 256 && logits.stride(1) == 1 && + logits.stride(0) >= 256 && logits.stride(0) % 4 == 0 && + reinterpret_cast(logits.data_ptr()) % 16 == 0, + "logits must be aligned FP32 CUDA rows of 256 experts"); + STD_TORCH_CHECK(logits.size(0) <= std::numeric_limits::max(), + "Too many routing tokens"); + const int tokens = logits.size(0); + STD_TORCH_CHECK(bias.device() == logits.device() && bias.dim() == 1 && + bias.size(0) == 256 && + bias.scalar_type() == ScalarType::Float && + bias.stride(0) == 1 && + reinterpret_cast(bias.data_ptr()) % 16 == 0, + "bias must be an aligned FP32 tensor of 256 experts on the " + "logits device"); + for (const auto* output : {&topk_weights, &topk_ids}) { + STD_TORCH_CHECK( + output->device() == logits.device() && output->dim() == 2 && + output->size(0) == tokens && output->size(1) == 8 && + output->stride(1) == 1 && output->stride(0) >= 8, + "routing outputs must have shape [tokens, 8] with contiguous rows"); + } + STD_TORCH_CHECK( + topk_weights.scalar_type() == ScalarType::Float && + (topk_ids.scalar_type() == ScalarType::Int || + topk_ids.scalar_type() == ScalarType::Long), + "routing outputs must be FP32 weights and int32 or int64 IDs"); + if (is_padding) { + STD_TORCH_CHECK(is_padding->device() == logits.device() && + is_padding->scalar_type() == ScalarType::Bool && + is_padding->dim() == 1 && is_padding->size(0) == tokens, + "is_padding must be a bool tensor of shape [tokens] on the " + "logits device"); + } + if (tokens == 0) return; + torch::stable::accelerator::DeviceGuard guard(logits.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(logits.get_device_index()); + VLLM_STABLE_DISPATCH_IDX_TYPES( + topk_ids.scalar_type(), "muse_spark_routing", [&] { + vllm::moe::RoutingParams params{ + .gating_output = static_cast(logits.data_ptr()), + .bias = static_cast(bias.data_ptr()), + .topk_weights = static_cast(topk_weights.data_ptr()), + .topk_ids = static_cast(topk_ids.data_ptr()), + .is_padding = is_padding + ? static_cast(is_padding->data_ptr()) + : nullptr, + .padding_stride = is_padding ? is_padding->stride(0) : 0, + .num_tokens = tokens, + .num_experts = 256, + .stride_gm = logits.stride(0), + .stride_wm = topk_weights.stride(0), + .stride_im = topk_ids.stride(0), + .router_output_multiplier = static_cast(multiplier), + .route_scale = static_cast(route_scale), + .division_eps = static_cast(division_eps), + }; + if (tokens <= 2048) { + const bool pdl = get_device_prop()->major >= 9; + if (pdl) + launch_routing<1, true>(params, stream); + else + launch_routing<1, false>(params, stream); + } else { + launch_routing<4, false>(params, stream); + } + }); +} diff --git a/csrc/moe/torch_bindings.cpp b/csrc/moe/torch_bindings.cpp index 9e8907afd714..3e2ee68a5eaa 100644 --- a/csrc/moe/torch_bindings.cpp +++ b/csrc/moe/torch_bindings.cpp @@ -23,6 +23,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding, " "Tensor? bias_vl=None, int image_sentinel_lo=0) -> ()"); +#ifdef USE_CUDA + m.def( + "muse_spark_routing(Tensor! topk_weights, Tensor! topk_ids, " + "Tensor logits, Tensor bias, float multiplier, float route_scale, " + "float division_eps, Tensor? is_padding) -> ()"); +#endif + // Calculate the result of moe by summing up the partial results // from all selected experts. topk_ids/expert_map are optional and, when // both given, enable pad-aware reduce that skips (token, expert) @@ -144,7 +151,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { m.def( "grouped_topk(Tensor scores, int n_group, int " "topk_group, int topk, bool renormalize, float " - "routed_scaling_factor, Tensor bias, int scoring_func) -> (Tensor, " + "routed_scaling_factor, Tensor bias, int scoring_func, Tensor? " + "is_padding=None) -> (Tensor, " "Tensor)"); #endif } @@ -152,6 +160,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { m.impl("topk_softmax", TORCH_BOX(&topk_softmax)); m.impl("topk_sigmoid", TORCH_BOX(&topk_sigmoid)); +#ifdef USE_CUDA + m.impl("muse_spark_routing", TORCH_BOX(&muse_spark_routing)); +#endif m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt)); m.impl("moe_sum", TORCH_BOX(&moe_sum)); m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size)); diff --git a/csrc/nvfp4_ds_mla_cache.h b/csrc/nvfp4_ds_mla_cache.h index 2b5f757846ac..b9aa9464813e 100644 --- a/csrc/nvfp4_ds_mla_cache.h +++ b/csrc/nvfp4_ds_mla_cache.h @@ -29,9 +29,9 @@ void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe, // Gather an nvfp4_ds_mla cache into a bf16 [total_tokens, 576] workspace. void launch_cp_gather_and_upconvert_nvfp4_kv_cache( const uint8_t* src_cache, void* dst, const int32_t* block_table, - const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size, - int32_t total_tokens, int64_t block_table_stride, - int64_t cache_block_stride, int64_t cache_entry_stride, - int64_t dst_entry_stride, cudaStream_t stream); + const int32_t* workspace_starts, const int32_t* seq_starts, + int32_t num_reqs, int32_t block_size, int32_t total_tokens, + int64_t block_table_stride, int64_t cache_block_stride, + int64_t cache_entry_stride, int64_t dst_entry_stride, cudaStream_t stream); } // namespace vllm diff --git a/csrc/nvfp4_ds_mla_cache_kernels.cu b/csrc/nvfp4_ds_mla_cache_kernels.cu index b1f369e5ab00..3b80f73de6c3 100644 --- a/csrc/nvfp4_ds_mla_cache_kernels.cu +++ b/csrc/nvfp4_ds_mla_cache_kernels.cu @@ -183,6 +183,7 @@ __global__ void cp_gather_and_upconvert_nvfp4_kv_cache_kernel( __nv_bfloat16* __restrict__ dst, // [total_tokens, 576] const int32_t* __restrict__ block_table, // [num_reqs, BLOCK_INDICES] const int32_t* __restrict__ workspace_starts, // [num_reqs] + const int32_t* __restrict__ seq_starts, // [num_reqs] or nullptr const int32_t num_reqs, const int32_t block_size, const int32_t total_tokens, const int64_t block_table_stride, const int64_t cache_block_stride, const int64_t cache_entry_stride, @@ -203,7 +204,10 @@ __global__ void cp_gather_and_upconvert_nvfp4_kv_cache_kernel( const int req_id = lo; const int out_token_id = flat_warp_id; - const int token_offset = out_token_id - workspace_starts[req_id]; + // seq_starts offsets the read into the request's context, so a chunk that + // continues a request reads from its start instead of the first token. + const int seq_start = seq_starts == nullptr ? 0 : seq_starts[req_id]; + const int token_offset = out_token_id - workspace_starts[req_id] + seq_start; const int cache_block_idx = token_offset / block_size; const int offset_in_block = token_offset % block_size; const int physical_block = @@ -254,18 +258,19 @@ void launch_concat_and_cache_nvfp4_ds_mla(const void* kv_c, const void* k_pe, void launch_cp_gather_and_upconvert_nvfp4_kv_cache( const uint8_t* src_cache, void* dst, const int32_t* block_table, - const int32_t* workspace_starts, int32_t num_reqs, int32_t block_size, - int32_t total_tokens, int64_t block_table_stride, - int64_t cache_block_stride, int64_t cache_entry_stride, - int64_t dst_entry_stride, cudaStream_t stream) { + const int32_t* workspace_starts, const int32_t* seq_starts, + int32_t num_reqs, int32_t block_size, int32_t total_tokens, + int64_t block_table_stride, int64_t cache_block_stride, + int64_t cache_entry_stride, int64_t dst_entry_stride, cudaStream_t stream) { constexpr int warps_per_block = 8; const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block; const int block_size_threads = warps_per_block * 32; // 256 threads cp_gather_and_upconvert_nvfp4_kv_cache_kernel<<>>( src_cache, reinterpret_cast<__nv_bfloat16*>(dst), block_table, - workspace_starts, num_reqs, block_size, total_tokens, block_table_stride, - cache_block_stride, cache_entry_stride, dst_entry_stride); + workspace_starts, seq_starts, num_reqs, block_size, total_tokens, + block_table_stride, cache_block_stride, cache_entry_stride, + dst_entry_stride); } } // namespace vllm diff --git a/csrc/ops.h b/csrc/ops.h index 7c4463ac3f1a..802de7f550e7 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -85,7 +85,8 @@ void get_cutlass_moe_mm_data( torch::stable::Tensor& output_permutation, const int64_t num_experts, const int64_t n, const int64_t k, const std::optional& blockscale_offsets, - const bool is_gated); + const bool is_gated, + const std::optional& sorted_expert_ids); void get_cutlass_moe_mm_problem_sizes_from_expert_offsets( const torch::stable::Tensor& expert_first_token_offset, @@ -412,7 +413,18 @@ void fused_minimax_m3_qknorm_rope_kv_insert( const std::string& kv_cache_dtype, bool skip_index_branch, std::optional q_fp8_out, double q_fp8_scale, std::optional kv_k_scale, - std::optional kv_v_scale); + std::optional kv_v_scale, bool write_qkv); + +void fused_minimax_m3_dspark_context_kv( + torch::stable::Tensor const& all_kv, + torch::stable::Tensor const& k_norm_weight, + torch::stable::Tensor const& cos_sin_cache, + torch::stable::Tensor const& positions, + torch::stable::Tensor const& slot_mapping, torch::stable::Tensor& kv_cache, + torch::stable::Tensor const& cache_offsets, + torch::stable::Tensor const& k_scale, + torch::stable::Tensor const& scale_offsets, int64_t rotary_dim, double eps, + const std::string& kv_cache_dtype); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE void fused_kda_decode( @@ -712,7 +724,7 @@ void concat_and_cache_mla(torch::stable::Tensor& kv_c, torch::stable::Tensor& kv_cache, torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype, - torch::stable::Tensor& scale); + torch::stable::Tensor& scale, bool skip_scale); void concat_and_cache_mla_grouped( torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe, @@ -745,13 +757,15 @@ void hisparse_resolve_residency( std::optional const& swap_device_physical_rows, std::optional const& swap_counts, std::optional const& resident_block_table, - int64_t resident_block_size, int64_t resident_null_block); + int64_t resident_block_size, int64_t resident_null_block, + std::optional const& num_valid_rows); void hisparse_invalidate_written_slots( - torch::stable::Tensor& device_global_indices, + torch::stable::Tensor const& device_global_indices_ptrs, torch::stable::Tensor const& request_state_indices, torch::stable::Tensor const& req_id_per_token, - torch::stable::Tensor const& written_slots); + torch::stable::Tensor const& written_slots, int64_t num_state_rows, + int64_t region_stride); void hisparse_gather_plan( torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache, @@ -825,7 +839,8 @@ void cp_gather_and_upconvert_nvfp4_kv_cache( torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size); + int64_t batch_size, + std::optional seq_starts = std::nullopt); // Indexer K quantization and cache function void indexer_k_quant_and_cache( @@ -869,3 +884,12 @@ void ngram_compute_n_gram_ids( torch::stable::Tensor& exclusive_req_len_sums, torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids); + +void fused_mqa_topk_prefill( + torch::stable::Tensor q, torch::stable::Tensor k, + torch::stable::Tensor scales, torch::stable::Tensor weights, + torch::stable::Tensor starts, torch::stable::Tensor ends, + torch::stable::Tensor out, torch::stable::Tensor counter, + torch::stable::Tensor values, torch::stable::Tensor indices, + torch::stable::Tensor retained, torch::stable::Tensor retained_indices, + torch::stable::Tensor counts, torch::stable::Tensor defer); diff --git a/csrc/persistent_topk.cuh b/csrc/persistent_topk.cuh index d6e41ab8a9d6..30ffed044561 100644 --- a/csrc/persistent_topk.cuh +++ b/csrc/persistent_topk.cuh @@ -12,10 +12,13 @@ #include #include "topk_histogram_4096.cuh" +#include "persistent_topk_overflow.cuh" namespace vllm { namespace persistent { +using topk_histogram_4096::convert_to_uint32_v2; + // ============================================================================ // Constants // ============================================================================ @@ -25,7 +28,7 @@ constexpr int RADIX = 256; // Medium path: all shared state in dynamic smem (no static __shared__, // which would inflate the kernel's smem footprint and kill occupancy -// for the decode/trivial paths). +// for the short and trivial paths). constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072 constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20 constexpr size_t kMediumHeaderSize = @@ -33,79 +36,19 @@ constexpr size_t kMediumHeaderSize = constexpr int MAX_BUFFERED_ITEMS = 4096; constexpr size_t kSmemMedium = kMediumHeaderSize + 2 * MAX_BUFFERED_ITEMS * sizeof(int); // 35968 +static_assert(topk_histogram_4096::kExactCandidateOffset + + MAX_BUFFERED_ITEMS * sizeof(int) <= + kSmemMedium); constexpr uint32_t RADIX_THRESHOLD = 32768; -// Decode path constants -constexpr int kDecodeBins = 2048; -constexpr uint32_t HIST2048_THRESHOLD = 8192; +// Short and medium path constants +constexpr int kCoarseBins = 2048; +constexpr uint32_t SHORT_EXACT_THRESHOLD = 8192; // Large path: fixed shared memory for histograms + scalars constexpr size_t kFixedSmemLarge = ((RADIX + RADIX + 5) * sizeof(uint32_t) + 15) & ~size_t(15); -// ============================================================================ -// Common helpers -// ============================================================================ - -__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); -} - -// ============================================================================ -// Vectorized load helpers -// ============================================================================ - -// Unconditional float4 load with cache hint (.cg = cache at global level only). -__device__ __forceinline__ void load_float4(const float* ptr, float& v0, - float& v1, float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - -// Per-element predicated scalar loads with -inf default. -__device__ __forceinline__ void load_float4_predicated(const float* ptr, - int base, int seq_len, - float& v0, float& v1, - float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - int p0 = (base < seq_len); - int p1 = (base + 1 < seq_len); - int p2 = (base + 2 < seq_len); - int p3 = (base + 3 < seq_len); - asm volatile( - "{\n" - " .reg .pred pr0, pr1, pr2, pr3;\n" - " setp.ne.u32 pr0, %4, 0;\n" - " setp.ne.u32 pr1, %5, 0;\n" - " setp.ne.u32 pr2, %6, 0;\n" - " setp.ne.u32 pr3, %7, 0;\n" - " mov.u32 %0, 0xFF800000;\n" - " mov.u32 %1, 0xFF800000;\n" - " mov.u32 %2, 0xFF800000;\n" - " mov.u32 %3, 0xFF800000;\n" - " @pr0 ld.global.cg.u32 %0, [%8];\n" - " @pr1 ld.global.cg.u32 %1, [%8+4];\n" - " @pr2 ld.global.cg.u32 %2, [%8+8];\n" - " @pr3 ld.global.cg.u32 %3, [%8+12];\n" - "}\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - // ============================================================================ // Large path: inter-CTA coordination state (one per group) // ============================================================================ @@ -122,8 +65,9 @@ struct RadixRowState { // Kernel parameters // ============================================================================ +template struct PersistentTopKParams { - const float* __restrict__ input; // [num_rows, stride] + const InT* __restrict__ input; // [num_rows, stride] int32_t* __restrict__ output; // [num_rows, top_k] const int32_t* __restrict__ lengths; // [num_rows] RadixRowState* row_states; // large path: per-group state @@ -136,284 +80,68 @@ struct PersistentTopKParams { }; // ============================================================================ -// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192) -// Uses 11-bit half-precision bins for fine granularity. -// One histogram pass typically suffices since 8192/2048 = 4 elements/bin avg. +// Coarse selection shared by the medium and filtered paths // ============================================================================ -// 11-bit bin from half-precision representation (ascending: high values -> high -// bins) -__device__ __forceinline__ uint32_t decode_bin(float x) { - __half hx = __float2half(x); - uint16_t bits = __half_as_ushort(hx); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return key >> 5; -} +template +struct CoarseSuffixSmem { + int histogram[Bins + 1]; + uint32_t warp_sum[BlockSize / 32]; +}; -template -__device__ __noinline__ void histogram_2048_topk( - const float* __restrict__ logits, int32_t* __restrict__ output_indices, - int32_t seq_len) { - extern __shared__ int decode_smem[]; +// Rewrite the coarse counts as suffix counts (histogram[b] = keys in bins >= b) +// and call on_select(bin) from the one thread whose bin holds rank k. Uses the +// shuffle scan of histogram_4096_topk: a CUB block scan spills under the +// persistent kernel's 32-register budget. +template +__device__ __forceinline__ void select_coarse_bin( + CoarseSuffixSmem* coarse, int k, OnSelect on_select) { + namespace hist4096 = topk_histogram_4096; + constexpr int kBinsPerThread = Bins / BlockSize; + static_assert(Bins % BlockSize == 0); + static_assert(BlockSize == 32 * 32, "one warp sum per lane"); const int tx = threadIdx.x; - const int lane = tx & 31; - - // ---- Layout constants ---- - constexpr int SBASE = 8192 - 8; // 8184 - constexpr int RHIST = RADIX + 128; // 384 - constexpr int BOFF = 2 * RHIST; // 768 - constexpr int DBUF = (SBASE - BOFF) / 2; // 3708 - constexpr int MAX_ITEMS_PER_THREAD = - (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; - - enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; - - // ---- Initialize scalars (prevents stale data from prior rows) ---- - if (tx < 8) { - decode_smem[SBASE + tx] = 0; - } - - // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ---- - int* histo = decode_smem; - uint16_t reg_bins[MAX_ITEMS_PER_THREAD]; - int nitems = 0; - - for (int i = tx; i < kDecodeBins; i += kThreadsPerBlock) { - histo[i] = 0; - } - __syncthreads(); - - const int n_vec = (seq_len + 3) >> 2; - const bool row_aligned = ((reinterpret_cast(logits) & 15) == 0); - - for (int i = tx; i < n_vec; i += kThreadsPerBlock) { - const int base = i << 2; - float v0, v1, v2, v3; - - if (row_aligned && base + 3 < seq_len) { - load_float4(logits + base, v0, v1, v2, v3); - } else { - load_float4_predicated(logits + base, base, seq_len, v0, v1, v2, v3); - } - - const uint16_t b0 = static_cast(decode_bin(v0)); - const uint16_t b1 = static_cast(decode_bin(v1)); - const uint16_t b2 = static_cast(decode_bin(v2)); - const uint16_t b3 = static_cast(decode_bin(v3)); - reg_bins[nitems++] = b0; - reg_bins[nitems++] = b1; - reg_bins[nitems++] = b2; - reg_bins[nitems++] = b3; - atomicAdd(&histo[b0], 1); - atomicAdd(&histo[b1], 1); - atomicAdd(&histo[b2], 1); - atomicAdd(&histo[b3], 1); - } - __syncthreads(); - - // ---- CUB suffix sum ---- - using BlockScanT = cub::BlockScan; - const int h0 = histo[2 * tx]; - const int pair_sum = h0 + histo[2 * tx + 1]; - - auto& scan_storage = *reinterpret_cast( - decode_smem + kDecodeBins); - - int pair_prefix, total; - BlockScanT(scan_storage).ExclusiveSum(pair_sum, pair_prefix, total); - - // Find threshold bin purely from registers - const int pair_suffix = total - pair_prefix; - - if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) { - decode_smem[SBASE + sTHR] = 2 * tx; - } - { - const int right_suf = pair_suffix - h0; - const int next_suf = pair_suffix - pair_sum; - if (right_suf >= TopK && next_suf < TopK) { - decode_smem[SBASE + sTHR] = 2 * tx + 1; - } + const uint32_t lane = tx % 32; + const uint32_t warp = tx / 32; + const int first_bin = kBinsPerThread * tx; + int counts[kBinsPerThread]; + uint32_t local_sum = 0; +#pragma unroll + for (int i = 0; i < kBinsPerThread; ++i) { + counts[i] = coarse->histogram[first_bin + i]; + local_sum += counts[i]; } + const uint32_t warp_inclusive = hist4096::warp_inclusive_sum(lane, local_sum); + if (lane == 31) coarse->warp_sum[warp] = warp_inclusive; __syncthreads(); - - const int threshold = decode_smem[SBASE + sTHR]; - - // ---- Phase 2: Collection with warp-aggregated atomicAdds ---- - int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF}; - const int sOUT_abs = SBASE + sOUT; - const int sBUF0_abs = SBASE + sBUF0; - - { - const uint32_t uthr = static_cast(threshold); - int item = 0; - const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; - - for (int iter = 0; iter < n_vec_iters; iter++) { - const int i = tx + iter * kThreadsPerBlock; - const bool vec_valid = (i < n_vec); - const int base_idx = i << 2; - -#pragma unroll 4 - for (int sub = 0; sub < 4; sub++) { - const int elem_idx = base_idx + sub; - uint32_t bin = 0; - if (vec_valid) bin = reg_bins[item++]; - // Lanes past seq_len are -inf padding from the row's last float4. - const bool in_range = vec_valid && (elem_idx < seq_len); - const bool is_above = in_range && (bin > uthr); - const bool is_equal = in_range && (bin == uthr); - - const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); - if (above_mask) { - const int above_count = __popc(above_mask); - const int above_rank = __popc(above_mask & ((1u << lane) - 1)); - int above_base; - if (lane == 0) { - above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); - } - above_base = __shfl_sync(0xffffffff, above_base, 0); - if (is_above) { - output_indices[above_base + above_rank] = elem_idx; - } - } - - const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); - if (equal_mask) { - const int equal_count = __popc(equal_mask); - const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); - int equal_base; - if (lane == 0) { - equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); - } - equal_base = __shfl_sync(0xffffffff, equal_base, 0); - if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) { - bufs[0][equal_base + equal_rank] = elem_idx; - } - } - } - } + const uint32_t warp_total = coarse->warp_sum[lane]; + const uint32_t total = hist4096::warp_reduce_sum_full(warp_total); + const uint32_t warps_before = + hist4096::warp_reduce_sum_full(lane < warp ? warp_total : 0); + // Each thread rewrites only its own bins, so no barrier is needed here. + int suffix = + static_cast(total - warps_before - warp_inclusive + local_sum); +#pragma unroll + for (int i = 0; i < kBinsPerThread; ++i) { + coarse->histogram[first_bin + i] = suffix; + suffix -= counts[i]; } + if (tx == 0) coarse->histogram[Bins] = 0; __syncthreads(); - int remaining_k = TopK - decode_smem[SBASE + sOUT]; - if (remaining_k <= 0) return; - - // If all buffered elements fit, output them all (common for short seqs) - const int raw_buf0 = decode_smem[SBASE + sBUF0]; - if (raw_buf0 <= remaining_k) { - const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; - const int base = decode_smem[SBASE + sOUT]; - for (int i = tx; i < nb; i += kThreadsPerBlock) { - output_indices[base + i] = bufs[0][i]; +#pragma unroll + for (int i = 0; i < kBinsPerThread; ++i) { + const int bin = first_bin + i; + if (coarse->histogram[bin] > k && coarse->histogram[bin + 1] <= k) { + on_select(bin); } - __syncthreads(); - return; - } - - // ---- Phase 3: Deferred refinement (rare path) ---- - int* refine[2] = {decode_smem, decode_smem + RHIST}; - const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; - - for (int i = tx; i < RHIST; i += kThreadsPerBlock) { - refine[0][i] = 0; } __syncthreads(); - - for (int i = tx; i < num_buf0; i += kThreadsPerBlock) { - const uint32_t fp32 = - topk_histogram_4096::convert_to_uint32_v2(logits[bufs[0][i]]); - atomicAdd(&refine[0][(fp32 >> 24) & 0xFF], 1); - } - __syncthreads(); - - auto compute_suffix_sum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (tx < RADIX) { - const int stride = 1 << i; - const int s = i & 1; - const int d = s ^ 1; - int value = refine[s][tx]; - if (tx < RADIX - stride) value += refine[s][tx + stride]; - refine[d][tx] = value; - } - __syncthreads(); - } - }; - -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { - const int src = pass & 1; - const int dst = src ^ 1; - - const int raw_buf = decode_smem[SBASE + sBUF0 + src]; - const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF; - - compute_suffix_sum(); - - if (tx < RADIX && refine[0][tx] > remaining_k && - refine[0][tx + 1] <= remaining_k) { - decode_smem[SBASE + sREF] = tx; - decode_smem[SBASE + sBUF0 + dst] = 0; - decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1]; - } - __syncthreads(); - - const int ref_thr = decode_smem[SBASE + sREF]; - remaining_k -= refine[0][ref_thr + 1]; - const int bit_offset = 24 - pass * 8; - - if (remaining_k == 0) { - for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { - const int idx = bufs[src][i]; - const uint32_t fp32 = - topk_histogram_4096::convert_to_uint32_v2(logits[idx]); - if (((fp32 >> bit_offset) & 0xFF) > static_cast(ref_thr)) { - const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; - } - } - __syncthreads(); - break; - } - - __syncthreads(); - if (tx < RADIX + 1) refine[0][tx] = 0; - __syncthreads(); - - for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { - const int idx = bufs[src][i]; - const float logit_val = logits[idx]; - const uint32_t fp32 = - topk_histogram_4096::convert_to_uint32_v2(logit_val); - const int bin = (fp32 >> bit_offset) & 0xFF; - - if (bin > ref_thr) { - const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; - } else if (bin == ref_thr) { - if (pass == 3) { - const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); - if (slot > 0) output_indices[TopK - slot] = idx; - } else { - const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); - if (__builtin_expect(bp < DBUF, 1)) { - bufs[dst][bp] = idx; - const int nbo = bit_offset - 8; - atomicAdd(&refine[0][(fp32 >> nbo) & 0xFF], 1); - } - } - } - } - __syncthreads(); - } } // ============================================================================ -// Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement -// For sequences 8K < seq_len <= 64K. +// Medium path: exact 11-bit FP32 histogram + 3-pass FP32 refinement +// For sequences 8K < seq_len <= 32K. // ============================================================================ // Adapted from: @@ -421,10 +149,12 @@ __device__ __noinline__ void histogram_2048_topk( // by: DarkSharpness // which at the same time is an optimized topk kernel copied from tilelang // kernel -template -__device__ __noinline__ void histogram_256_topk( - const float* __restrict__ logits, int* __restrict__ output_indices, - int logits_offset, int seq_len) { +// Returns where exact recovery resumes when the threshold bin overflows the +// candidate buffer; .needed is false otherwise. +template +__device__ __noinline__ overflow::MediumResume histogram_medium_topk( + const InT* __restrict__ logits, int* __restrict__ output_indices, + int seq_len) { // All shared state lives in dynamic shared memory to avoid static extern __shared__ char medium_smem[]; @@ -442,16 +172,41 @@ __device__ __noinline__ void histogram_256_topk( const int thread_id = threadIdx.x; int remaining_k = TopK; - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; + // The coarse histogram overlays the refinement state until it is consumed. + struct CoarseSmem { + CoarseSuffixSmem suffix; + int threshold; + }; + static_assert(sizeof(CoarseSmem) <= kSmemMedium); + auto* coarse = reinterpret_cast(medium_smem); + int* coarse_histogram = coarse->suffix.histogram; + + for (int bin = thread_id; bin < kCoarseBins + 1; bin += kThreadsPerBlock) { + coarse_histogram[bin] = 0; } __syncthreads(); - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const auto bin = convert_to_uint8(logits[idx + logits_offset]); - atomicAdd(&shared_histogram[0][bin], 1); + topk_histogram_4096::scan_scores( + logits, seq_len, [&](uint32_t, float score) { + atomicAdd(&coarse_histogram[convert_to_uint32_v2(score) >> 21], 1); + }); + __syncthreads(); + + select_coarse_bin(&coarse->suffix, TopK, + [&](int bin) { coarse->threshold = bin; }); + + const int threshold_bin = coarse->threshold; + const int count_above = coarse_histogram[threshold_bin + 1]; + const int threshold_bin_count = coarse_histogram[threshold_bin] - count_above; + __syncthreads(); + + if (thread_id == 0) { + shared_threshold_bin = threshold_bin; + shared_buffered_count[0] = 0; + shared_output_count = 0; } __syncthreads(); + remaining_k -= count_above; auto compute_cumulative_sum = [&]() { #pragma unroll 8 @@ -470,58 +225,50 @@ __device__ __noinline__ void histogram_256_topk( } }; - compute_cumulative_sum(); - - if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && - shared_histogram[0][thread_id + 1] <= remaining_k) { - shared_threshold_bin = thread_id; - shared_buffered_count[0] = 0; - shared_output_count = 0; - } - __syncthreads(); - - const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; - if (remaining_k == 0) { - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const int bin = convert_to_uint8(logits[idx + logits_offset]); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } - } + topk_histogram_4096::scan_scores( + logits, seq_len, [&](uint32_t idx, float score) { + const int bin = convert_to_uint32_v2(score) >> 21; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } + }); __syncthreads(); - return; + return {}; + } + + if (__builtin_expect(threshold_bin_count > MAX_BUFFERED_ITEMS, 0)) { + return {.needed = true, + .prefix = static_cast(threshold_bin) << 21, + .remaining = static_cast(remaining_k)}; } - __syncthreads(); if (thread_id < RADIX + 1) { shared_histogram[0][thread_id] = 0; } __syncthreads(); - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const float logit_value = logits[idx + logits_offset]; - const int bin = convert_to_uint8(logit_value); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); - if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { - buffered_indices[0][buffer_pos] = idx; - const uint32_t fp32_bits = - topk_histogram_4096::convert_to_uint32_v2(logit_value); - const int next_bin = (fp32_bits >> 24) & 0xFF; - atomicAdd(&shared_histogram[0][next_bin], 1); - } - } - } + topk_histogram_4096::scan_scores( + logits, seq_len, [&](uint32_t idx, float score) { + const uint32_t ordered = convert_to_uint32_v2(score); + const int bin = ordered >> 21; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[0][buffer_pos] = idx; + const int next_bin = (ordered >> 16) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + }); __syncthreads(); -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { +#pragma unroll 3 + for (int pass = 0; pass < 3; ++pass) { const int src_buffer = pass % 2; const int dst_buffer = src_buffer ^ 1; const int raw_buffered = shared_buffered_count[src_buffer]; @@ -540,13 +287,13 @@ __device__ __noinline__ void histogram_256_topk( const int threshold_bin = shared_threshold_bin; remaining_k -= shared_histogram[0][threshold_bin + 1]; - const int bit_offset = 24 - pass * 8; + const int bit_offset = 16 - pass * 8; if (remaining_k == 0) { for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { const int idx = buffered_indices[src_buffer][i]; - const uint32_t fp32_bits = topk_histogram_4096::convert_to_uint32_v2( - logits[idx + logits_offset]); + const uint32_t fp32_bits = + convert_to_uint32_v2(topk_histogram_4096::to_float(logits[idx])); const int bin = (fp32_bits >> bit_offset) & 0xFF; if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); @@ -565,15 +312,14 @@ __device__ __noinline__ void histogram_256_topk( for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { const int idx = buffered_indices[src_buffer][i]; - const float logit_value = logits[idx + logits_offset]; - const uint32_t fp32_bits = - topk_histogram_4096::convert_to_uint32_v2(logit_value); + const float logit_value = topk_histogram_4096::to_float(logits[idx]); + const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); const int bin = (fp32_bits >> bit_offset) & 0xFF; if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); output_indices[output_pos] = idx; } else if (bin == threshold_bin) { - if (pass == 3) { + if (pass == 2) { const int slot = atomicAdd(&shared_final_k, -1); if (slot > 0) { output_indices[TopK - slot] = idx; @@ -592,6 +338,7 @@ __device__ __noinline__ void histogram_256_topk( } __syncthreads(); } + return {}; } // ============================================================================ @@ -643,7 +390,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val, } // ============================================================================ -// Large path: multi-CTA radix select for sequences > 64K +// Large path: multi-CTA radix select for sequences > 32K // // Each row is processed by a group of CTAs. Each CTA loads its chunk into // shared memory as ordered uint32, then participates in 4 rounds of @@ -655,8 +402,8 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val, // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 // ============================================================================ -template -__device__ void radix_topk(const float* __restrict__ row_input, +template +__device__ void radix_topk(const InT* __restrict__ row_input, int32_t* __restrict__ row_output, uint32_t seq_len, uint32_t my_chunk_start, uint32_t chunk_size, uint32_t* local_histogram, uint32_t* suffix_sum, @@ -676,25 +423,26 @@ __device__ void radix_topk(const float* __restrict__ row_input, for (uint32_t i = tx * VEC_SIZE; i < aligned_size; i += kThreadsPerBlock * VEC_SIZE) { - const float* src = row_input + my_chunk_start + i; + const InT* src = row_input + my_chunk_start + i; if constexpr (VEC_SIZE == 4) { - float4 v = *reinterpret_cast(src); - shared_ordered[i] = topk_histogram_4096::convert_to_uint32_v2(v.x); - shared_ordered[i + 1] = topk_histogram_4096::convert_to_uint32_v2(v.y); - shared_ordered[i + 2] = topk_histogram_4096::convert_to_uint32_v2(v.z); - shared_ordered[i + 3] = topk_histogram_4096::convert_to_uint32_v2(v.w); + float4 v = topk_histogram_4096::load_score4(src); + shared_ordered[i] = convert_to_uint32_v2(v.x); + shared_ordered[i + 1] = convert_to_uint32_v2(v.y); + shared_ordered[i + 2] = convert_to_uint32_v2(v.z); + shared_ordered[i + 3] = convert_to_uint32_v2(v.w); } else if constexpr (VEC_SIZE == 2) { - float2 v = *reinterpret_cast(src); - shared_ordered[i] = topk_histogram_4096::convert_to_uint32_v2(v.x); - shared_ordered[i + 1] = topk_histogram_4096::convert_to_uint32_v2(v.y); + float2 v = topk_histogram_4096::load_score2(src); + shared_ordered[i] = convert_to_uint32_v2(v.x); + shared_ordered[i + 1] = convert_to_uint32_v2(v.y); } else { - shared_ordered[i] = topk_histogram_4096::convert_to_uint32_v2(*src); + shared_ordered[i] = + convert_to_uint32_v2(topk_histogram_4096::to_float(*src)); } } for (uint32_t i = aligned_size + tx; i < actual_chunk_size; i += kThreadsPerBlock) { - shared_ordered[i] = topk_histogram_4096::convert_to_uint32_v2( - row_input[my_chunk_start + i]); + shared_ordered[i] = convert_to_uint32_v2( + topk_histogram_4096::to_float(row_input[my_chunk_start + i])); } } __syncthreads(); @@ -755,6 +503,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, next_hist[i] = 0; } } + __syncthreads(); if (tx == 0) { red_release(&state->arrival_counter, 1); @@ -856,23 +605,23 @@ __device__ void radix_topk(const float* __restrict__ row_input, } } } + // The CTA's next row reuses this smem; finish the tie loop before it does. + __syncthreads(); } // ============================================================================ -// Persistent kernel — BS≤32, decode/medium/large paths with RadixTopK -// BS>32 uses standalone histogram_256_buffered_topk (separate kernel, -// see filtered_topk.cuh) +// Persistent kernel — short/medium/large paths with RadixTopK. // ============================================================================ -template +template __global__ void __launch_bounds__(kThreadsPerBlock, 2) - persistent_topk_kernel(PersistentTopKParams params) { + persistent_topk_kernel(PersistentTopKParams params) { const uint32_t tx = threadIdx.x; extern __shared__ uint8_t smem_raw[]; // ======================================================================== // Group mode: multi-CTA groups with static round-robin row assignment. - // Non-large rows: CTA-0 handles trivial/decode/medium. + // Non-large rows: CTA-0 handles the trivial, short, and medium paths. // Large rows: all CTAs in the group cooperate via RadixTopK. // ======================================================================== const uint32_t ctas_per_group = params.ctas_per_group; @@ -930,7 +679,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) const uint32_t seq_len = non_negative_len < row_bound ? non_negative_len : row_bound; int32_t* row_output = params.output + row_idx * params.top_k; - const float* row_input = params.input + row_idx * params.stride; + const InT* row_input = params.input + row_idx * params.stride; if (seq_len <= RADIX_THRESHOLD) { if (cta_in_group == 0) { @@ -940,10 +689,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) i += kThreadsPerBlock) { row_output[i] = (i < seq_len) ? static_cast(i) : -1; } - } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { - histogram_2048_topk(row_input, row_output, seq_len); - } else { - histogram_256_topk(row_input, row_output, 0, seq_len); + } else if (seq_len <= static_cast(SHORT_EXACT_THRESHOLD)) { + topk_histogram_4096::exact_topk_rescan_cold( + row_input, row_output, seq_len, smem_raw); + } else if (const auto resume = histogram_medium_topk( + row_input, row_output, seq_len); + __builtin_expect(resume.needed, 0)) { + overflow::recover_medium( + row_input, row_output, seq_len, smem_raw, resume.prefix, + resume.remaining); } } continue; @@ -961,19 +716,13 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) } // namespace persistent // ============================================================================ -// ============================================================================ -// Optimized FilteredTopK — single CTA per row for bs > 32. -// Kept with persistent_topk so the portable fallback owns the non-cluster path. +// FilteredTopK — one CTA per row, for the non-cluster path. +// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 // ============================================================================ namespace filtered_topk { namespace hist4096 = topk_histogram_4096; -// ============================================================================ -// FilteredTopK — single CTA per row for bs > 32 -// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 -// ============================================================================ - #define FLASHINFER_CUDA_CALL(func, ...) \ { \ cudaError_t e = (func); \ @@ -991,38 +740,37 @@ struct vec_t { FLASHINFER_INLINE T& operator[](size_t i) { return data[i]; } FLASHINFER_INLINE const T& operator[](size_t i) const { return data[i]; } - FLASHINFER_INLINE void cast_load(const T* ptr) { + template + FLASHINFER_INLINE void cast_load(const InT* ptr) { #pragma unroll for (size_t i = 0; i < N; ++i) { - data[i] = ptr[i]; + data[i] = topk_histogram_4096::to_float(ptr[i]); } } }; #undef FLASHINFER_INLINE // FilteredTopK traits for different data types -template +template struct FilteredTopKTraits; -// Specialization for float (32-bit): coarse histogram uses FP16 high 8 bits, 4 -// refinement rounds -template <> -struct FilteredTopKTraits { +// Both float paths use the exact ordered-FP32 key; only the coarse-bin width +// and refinement digit layout differ. +template +struct FilteredTopKTraits { using OrderedType = uint32_t; - static constexpr int NUM_REFINE_ROUNDS = 4; - static constexpr int FIRST_REFINE_SHIFT = 24; - - __device__ __forceinline__ static uint8_t ToCoarseKey(float x) { - // Convert to FP16 representation and extract high 8 bits - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); + static constexpr int COARSE_BITS = UseWideCoarse ? 12 : 11; + static constexpr int COARSE_RADIX = 1 << COARSE_BITS; + static constexpr int NUM_REFINE_ROUNDS = UseWideCoarse ? 2 : 3; + static constexpr int FIRST_REFINE_SHIFT = UseWideCoarse ? 10 : 16; + static constexpr int REFINE_BITS = UseWideCoarse ? 10 : 8; + + __device__ __forceinline__ static uint16_t ToCoarseKey(float x) { + return static_cast(ToOrdered(x) >> (32 - COARSE_BITS)); } __device__ __forceinline__ static OrderedType ToOrdered(float x) { - return hist4096::convert_to_uint32_v2(x); + return topk_histogram_4096::convert_to_uint32_v2(x); } }; @@ -1031,10 +779,16 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = 16 * 1024; // 16K indices per buffer constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB +constexpr uint32_t FILTERED_TOPK_WIDE_COARSE_MIN_LENGTH = 64 * 1024; +constexpr uint32_t FILTERED_TOPK_WIDE_COARSE_MAX_LENGTH = 192 * 1024; +static_assert(hist4096::kExactCandidateOffset + + FILTERED_TOPK_SMEM_INPUT_SIZE * sizeof(int) <= + FILTERED_TOPK_SMEM_DYNAMIC); -template +template struct FilteredTopKStorage { - alignas(128) int histogram[2][256 + 128]; + static constexpr int RADIX = UseWideCoarse ? 1024 : 256; + alignas(128) int histogram[2][RADIX + 128]; alignas(128) int counter; alignas(128) int threshold_bin; alignas(128) int num_input[2]; @@ -1042,16 +796,24 @@ struct FilteredTopKStorage { int last_remain; }; -// With CheckOverflow, return false before using a truncated stash so the caller -// can retry with exact full-row selection. +// With CheckOverflow, return false instead of recovering when the candidate +// stash overflows, so the caller can finish with its own exact selection. template + bool UsePredicatedShortLoads, bool CheckOverflow = false, + bool UseWideCoarse = false> __device__ __forceinline__ bool filtered_topk_row( const DType* score, IdType* dst, int length, uint32_t top_k, - FilteredTopKStorage& storage) { + FilteredTopKStorage& storage, uint32_t num_rows = 0) { + extern __shared__ __align__(16) uint8_t dynamic_smem[]; constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; - constexpr int RADIX = 256; constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; + using Traits = FilteredTopKTraits; + constexpr int RADIX = 1 << Traits::REFINE_BITS; + static_assert(RADIX == FilteredTopKStorage::RADIX); + constexpr int COARSE_RADIX = Traits::COARSE_RADIX; + using CoarseSmem = persistent::CoarseSuffixSmem; + static_assert(sizeof(CoarseSmem) <= FILTERED_TOPK_SMEM_DYNAMIC); + const int tx = threadIdx.x; // Trivial case: length <= top_k @@ -1064,13 +826,14 @@ __device__ __forceinline__ bool filtered_topk_row( // Short path if (length <= 32768) { - extern __shared__ uint8_t _smem_reg[]; if constexpr (UsePredicatedShortLoads) { - hist4096::histogram_4096_topk_predicated(score, dst, length, - _smem_reg); + hist4096::histogram_4096_topk_predicated< + MAX_K, 12, 8, hist4096::OverflowRecovery::kRegisterValues>( + score, dst, length, dynamic_smem); } else { - hist4096::histogram_4096_topk(score, dst, length, - _smem_reg); + hist4096::histogram_4096_topk< + MAX_K, 12, 8, false, hist4096::OverflowRecovery::kRegisterValues>( + score, dst, length, dynamic_smem); } return true; } @@ -1080,19 +843,22 @@ __device__ __forceinline__ bool filtered_topk_row( auto& s_threshold_bin_id = storage.threshold_bin; auto& s_num_input = storage.num_input; auto& s_indices = storage.indices; + auto& s_histogram = s_histogram_buf[0]; - // Dynamic shared memory for input double buffer - extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE]; + auto* coarse_smem = reinterpret_cast(dynamic_smem); + auto* s_coarse_histogram = coarse_smem->histogram; + auto* s_input_idx = reinterpret_cast(dynamic_smem); - using Traits = FilteredTopKTraits; int topk = top_k; - // Stage 1: 8-bit coarse histogram with vectorized loads - if (tx < RADIX + 1) s_histogram[tx] = 0; + // Stage 1: high bits of the exact ordered-FP32 key. + for (int bin = tx; bin < COARSE_RADIX + 1; bin += BLOCK_SIZE) { + s_coarse_histogram[bin] = 0; + } __syncthreads(); - vec_t score_vec; + vec_t score_vec; const int aligned_length = (length / VEC_SIZE) * VEC_SIZE; #pragma unroll 2 @@ -1102,20 +868,20 @@ __device__ __forceinline__ bool filtered_topk_row( #pragma unroll for (int j = 0; j < VEC_SIZE; ++j) { const auto bin = Traits::ToCoarseKey(score_vec[j]); - atomicAdd(&s_histogram[bin], 1); + atomicAdd(&s_coarse_histogram[bin], 1); } } // Handle tail for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - const auto bin = Traits::ToCoarseKey(score[i]); - atomicAdd(&s_histogram[bin], 1); + const auto bin = Traits::ToCoarseKey(hist4096::to_float(score[i])); + atomicAdd(&s_coarse_histogram[bin], 1); } __syncthreads(); - // Suffix sum - const auto run_cumsum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { + // Suffix sum for the exact-refinement histograms. + const auto run_refine_cumsum = [&]() { +#pragma unroll + for (int i = 0; i < Traits::REFINE_BITS; ++i) { if (tx < RADIX) { const auto j = 1 << i; const auto k = i & 1; @@ -1129,16 +895,15 @@ __device__ __forceinline__ bool filtered_topk_row( } }; - run_cumsum(); - if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { - s_threshold_bin_id = tx; + // The selecting thread also resets the refinement counters, saving a barrier. + persistent::select_coarse_bin(coarse_smem, topk, [&](int bin) { + s_threshold_bin_id = bin; s_num_input[0] = 0; s_counter = 0; - } - __syncthreads(); + }); const auto threshold_bin = s_threshold_bin_id; - topk -= s_histogram[threshold_bin + 1]; + topk -= s_coarse_histogram[threshold_bin + 1]; constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS; constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT; @@ -1160,7 +925,8 @@ __device__ __forceinline__ bool filtered_topk_row( } // Handle tail for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - const auto bin = static_cast(Traits::ToCoarseKey(score[i])); + const auto bin = + static_cast(Traits::ToCoarseKey(hist4096::to_float(score[i]))); if (bin > threshold_bin) { const auto pos = atomicAdd(&s_counter, 1); s_indices[pos] = i; @@ -1168,8 +934,21 @@ __device__ __forceinline__ bool filtered_topk_row( } __syncthreads(); } else { - __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; + // The threshold bin does not fit the candidate stash: finish exactly. + if (__builtin_expect(s_coarse_histogram[threshold_bin] - + s_coarse_histogram[threshold_bin + 1] > + SMEM_INPUT_SIZE, + 0)) { + if constexpr (CheckOverflow) return false; + overflow::recover_coarse( + score, dst, length, top_k, num_rows, s_coarse_histogram, + threshold_bin, dynamic_smem, s_histogram); + return true; + } + for (int bin = tx; bin < RADIX + 1; bin += BLOCK_SIZE) { + s_histogram[bin] = 0; + } __syncthreads(); // Filter + histogram for refinement @@ -1183,7 +962,7 @@ __device__ __forceinline__ bool filtered_topk_row( if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { s_input_idx[0][pos] = index; const auto ordered = Traits::ToOrdered(raw_input); - const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; + const auto sub_bin = (ordered >> FIRST_SHIFT) & (RADIX - 1); atomicAdd(&s_histogram[sub_bin], 1); } } @@ -1199,11 +978,11 @@ __device__ __forceinline__ bool filtered_topk_row( } // Handle tail for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - filter_and_add_to_histogram(score[i], i); + filter_and_add_to_histogram(hist4096::to_float(score[i]), i); } __syncthreads(); - // Stage 2: refine with 8bit radix passes + // Stage 2: refine the selected coarse bin. #pragma unroll for (int round = 0; round < NUM_ROUNDS; ++round) { auto& s_last_remain = storage.last_remain; @@ -1216,7 +995,7 @@ __device__ __forceinline__ bool filtered_topk_row( const auto num_input = (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; - run_cumsum(); + run_refine_cumsum(); if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { s_threshold_bin_id = tx; s_num_input[r_idx ^ 1] = 0; @@ -1227,13 +1006,15 @@ __device__ __forceinline__ bool filtered_topk_row( const auto threshold = s_threshold_bin_id; topk -= s_histogram[threshold + 1]; - const int offset = FIRST_SHIFT - round * 8; + const int offset = FIRST_SHIFT - round * Traits::REFINE_BITS; const bool is_last_round = (round == NUM_ROUNDS - 1); if (topk == 0) { for (int i = tx; i < num_input; i += BLOCK_SIZE) { const auto idx = s_input_idx[r_idx][i]; - const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF; + const auto bin = + (Traits::ToOrdered(hist4096::to_float(score[idx])) >> offset) & + (RADIX - 1); if (static_cast(bin) > threshold) { const auto pos = atomicAdd(&s_counter, 1); s_indices[pos] = idx; @@ -1243,12 +1024,15 @@ __device__ __forceinline__ bool filtered_topk_row( break; } else { __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; + for (int bin = tx; bin < RADIX + 1; bin += BLOCK_SIZE) { + s_histogram[bin] = 0; + } __syncthreads(); for (int i = tx; i < num_input; i += BLOCK_SIZE) { const auto idx = s_input_idx[r_idx][i]; - const auto raw_input = score[idx]; - const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF; + const float raw_input = hist4096::to_float(score[idx]); + const auto bin = + (Traits::ToOrdered(raw_input) >> offset) & (RADIX - 1); if (static_cast(bin) > threshold) { const auto pos = atomicAdd(&s_counter, 1); s_indices[pos] = idx; @@ -1263,7 +1047,8 @@ __device__ __forceinline__ bool filtered_topk_row( if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { s_input_idx[r_idx ^ 1][pos] = idx; const auto bin32 = Traits::ToOrdered(raw_input); - const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF; + const auto sub_bin = + (bin32 >> (offset - Traits::REFINE_BITS)) & (RADIX - 1); atomicAdd(&s_histogram[sub_bin], 1); } } @@ -1291,7 +1076,7 @@ __device__ __forceinline__ bool filtered_topk_row( * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) */ template + bool UsePredicatedShortLoads = false, bool UseWideCoarse = false> __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) FilteredTopKUnifiedKernel(const DType* __restrict__ input, IdType* __restrict__ output, @@ -1301,9 +1086,11 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const uint32_t bid = blockIdx.x; if (bid >= num_rows) return; const int length = lengths ? lengths[bid] : static_cast(max_len); - __shared__ FilteredTopKStorage storage; - filtered_topk_row( - input + bid * max_len, output + bid * top_k, length, top_k, storage); + __shared__ FilteredTopKStorage storage; + filtered_topk_row(input + bid * max_len, + output + bid * top_k, length, top_k, + storage, num_rows); } // Helper to compute GCD for VEC_SIZE selection @@ -1320,7 +1107,7 @@ constexpr uint32_t gcd(uint32_t a, uint32_t b) { // Returns 1, 2, 4, or 8 template constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { - constexpr int MAX_VEC = 16 / sizeof(DType); // 4 for float32, 8 for fp16/bf16 + constexpr int MAX_VEC = 4; // Use GCD to find largest power-of-2 divisor const uint32_t g = gcd(max_len, static_cast(MAX_VEC)); return static_cast(g); @@ -1334,24 +1121,38 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input, uint32_t max_len, cudaStream_t stream = 0) { constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; - constexpr int MAX_VEC = 16 / sizeof(DType); + constexpr int MAX_VEC = 4; dim3 grid(num_rows); dim3 block(FILTERED_TOPK_BLOCK_THREADS); void* args[] = {&input, &output_indices, &lengths, &num_rows, &top_k_val, &max_len}; - const int vec_size = ComputeFilteredTopKVecSize(max_len); - -#define DISPATCH_VEC_SIZE(VS) \ - if (vec_size == VS) { \ - auto kernel = \ - FilteredTopKUnifiedKernel; \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ - smem_size, stream)); \ - return cudaSuccess; \ + // Views like x[:, 1:] keep the stride but shift the base. + int vec_size = ComputeFilteredTopKVecSize(max_len); + while (vec_size > 1 && + reinterpret_cast(input) % (vec_size * sizeof(DType)) != 0) { + vec_size /= 2; + } + const bool use_wide_coarse = + MAX_K == 1024 && max_len >= FILTERED_TOPK_WIDE_COARSE_MIN_LENGTH && + max_len < FILTERED_TOPK_WIDE_COARSE_MAX_LENGTH; + +#define DISPATCH_VEC_SIZE(VS) \ + if (vec_size == VS) { \ + auto kernel = FilteredTopKUnifiedKernel; \ + if constexpr (MAX_K == 1024) { \ + if (use_wide_coarse) { \ + kernel = FilteredTopKUnifiedKernel; \ + } \ + } \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ + smem_size, stream)); \ + return cudaSuccess; \ } DISPATCH_VEC_SIZE(1) diff --git a/csrc/persistent_topk_overflow.cuh b/csrc/persistent_topk_overflow.cuh new file mode 100644 index 000000000000..306dcfe263a5 --- /dev/null +++ b/csrc/persistent_topk_overflow.cuh @@ -0,0 +1,265 @@ +#ifndef PERSISTENT_TOPK_OVERFLOW_CUH_ +#define PERSISTENT_TOPK_OVERFLOW_CUH_ + +#include +#include +#include +#include + +#include "topk_histogram_4096.cuh" +#include "topk_histogram_4096_overflow.cuh" + +namespace vllm::persistent::overflow { + +// Where exact recovery resumes after the medium path's coarse histogram: +// the selected 11-bit FP32 prefix and how many keys are still needed in it. +struct MediumResume { + bool needed = false; + uint32_t prefix = 0; + uint32_t remaining = 0; +}; + +template +__device__ __noinline__ void recover_medium(const InT* scores, int32_t* output, + uint32_t length, void* smem, + uint32_t prefix, + uint32_t remaining) { + topk_histogram_4096::exact_topk_rescan( + scores, output, length, smem, prefix, remaining); +} + +} // namespace vllm::persistent::overflow + +namespace vllm::filtered_topk { + +template +struct vec_t; + +template +struct FilteredTopKTraits; + +namespace overflow { + +// Sampled recovery pays off only for long rows in large batches. topk.cu +// currently routes such batches to sampled_topk_kernel, so this matters for +// direct FilteredTopKRaggedTransform callers and future dispatch changes. +constexpr uint32_t kSamplingMinRows = 96; +constexpr uint32_t kSamplingMinLength = 192 * 1024; + +// 8-bit digits over the ordered key: the sampled candidate set is small, so a +// 256-bin histogram selected by one thread is cheaper than a block scan. +struct SampledRadixDigits { + static constexpr uint32_t kBins = 256; + static constexpr uint32_t kRounds = 4; + + __device__ static constexpr uint32_t shift(uint32_t round) { + return 24 - round * 8; + } + + __device__ static constexpr uint32_t digit_mask(uint32_t) { return 0xFFu; } + + __device__ static constexpr uint32_t prefix_mask(uint32_t round) { + return round == 0 ? 0u : (~0u << (32 - round * 8)); + } +}; + +// Sample a pivot, bound the candidate set, then select it exactly by radix. +template +__device__ bool try_sampled_exact_topk(const DType* __restrict__ scores, + IdType* __restrict__ output, + uint32_t length, uint32_t top_k, + void* dynamic_smem, int* scratch) { + using Traits = + FilteredTopKTraits; + using Radix = SampledRadixDigits; + constexpr uint32_t kItemsPerThread = 2; + constexpr uint32_t kSampleSize = BlockSize * kItemsPerThread; + using SampleSort = cub::BlockRadixSort; + static_assert(sizeof(typename SampleSort::TempStorage) <= + 2 * CandidateCapacity * sizeof(int)); + + // scratch: the radix histogram followed by four counters. + int* histogram = scratch; + int* candidate_count = scratch + Radix::kBins; + int* output_counter = candidate_count + 1; + int* radix_prefix = candidate_count + 2; + int* radix_remaining = candidate_count + 3; + const uint32_t tx = threadIdx.x; + uint32_t sample_keys[kItemsPerThread]; +#pragma unroll + for (uint32_t i = 0; i < kItemsPerThread; ++i) { + const uint32_t sample_idx = tx * kItemsPerThread + i; + const uint32_t score_idx = static_cast( + (static_cast(2 * sample_idx + 1) * length) / + (2 * kSampleSize)); + sample_keys[i] = + Traits::ToOrdered(topk_histogram_4096::to_float(scores[score_idx])); + } + + auto* sort_smem = + static_cast(dynamic_smem); + SampleSort(*sort_smem).SortDescending(sample_keys); + __syncthreads(); + + const uint32_t target_candidates = + min(CandidateCapacity / 2, (top_k < 2048) ? 4 * top_k : 2 * top_k); + uint32_t sample_rank = static_cast( + (static_cast(target_candidates) * kSampleSize + length - 1) / + length); + sample_rank = max(1u, min(sample_rank, kSampleSize)); + const uint32_t rank_index = sample_rank - 1; + if (tx == rank_index / kItemsPerThread) { + *radix_prefix = static_cast(sample_keys[rank_index % kItemsPerThread]); + } + if (tx == 0) *candidate_count = 0; + __syncthreads(); + + auto* candidates = static_cast(dynamic_smem); + const uint32_t proposed_pivot = static_cast(*radix_prefix); + // Every key tied with the pivot is a candidate. When the sample already + // predicts more candidates than fit, skip the doomed collection scan. + uint32_t samples_at_or_above = 0; +#pragma unroll + for (uint32_t i = 0; i < kItemsPerThread; ++i) { + samples_at_or_above += + __syncthreads_count(sample_keys[i] >= proposed_pivot); + } + if (static_cast(samples_at_or_above) * length > + static_cast(CandidateCapacity) * kSampleSize) { + return false; + } + const auto collect_candidate = [&](uint32_t idx, float score) { + if (Traits::ToOrdered(score) >= proposed_pivot) { + const int pos = atomicAdd(candidate_count, 1); + if (pos < static_cast(CandidateCapacity)) candidates[pos] = idx; + } + }; + + vec_t score_vec; + const uint32_t aligned_length = length / VecSize * VecSize; + for (uint32_t base = tx * VecSize; base < aligned_length; + base += BlockSize * VecSize) { + score_vec.cast_load(scores + base); +#pragma unroll + for (uint32_t i = 0; i < VecSize; ++i) { + collect_candidate(base + i, score_vec[i]); + } + } + for (uint32_t idx = aligned_length + tx; idx < length; idx += BlockSize) { + collect_candidate(idx, topk_histogram_4096::to_float(scores[idx])); + } + __syncthreads(); + + const int num_candidates = *candidate_count; + if (num_candidates < static_cast(top_k) || + num_candidates > static_cast(CandidateCapacity)) { + return false; + } + + auto* candidate_keys = candidates + CandidateCapacity; + for (int i = tx; i < num_candidates; i += BlockSize) { + candidate_keys[i] = static_cast(Traits::ToOrdered( + topk_histogram_4096::to_float(scores[candidates[i]]))); + } + + if (tx == 0) { + *radix_prefix = 0; + *radix_remaining = static_cast(top_k); + } + __syncthreads(); + +#pragma unroll + for (uint32_t round = 0; round < Radix::kRounds; ++round) { + if (tx < Radix::kBins) histogram[tx] = 0; + __syncthreads(); + + const uint32_t prefix = static_cast(*radix_prefix); + const uint32_t prefix_mask = Radix::prefix_mask(round); + const uint32_t shift = Radix::shift(round); + for (int i = tx; i < num_candidates; i += BlockSize) { + const uint32_t ordered = static_cast(candidate_keys[i]); + if ((ordered & prefix_mask) == prefix) { + atomicAdd(&histogram[(ordered >> shift) & Radix::digit_mask(round)], 1); + } + } + __syncthreads(); + + if (tx == 0) { + int count_above = 0; + const int remaining = *radix_remaining; + for (int bin = static_cast(Radix::kBins) - 1; bin >= 0; --bin) { + const int count = histogram[bin]; + if (count_above + count >= remaining) { + *radix_remaining = remaining - count_above; + *radix_prefix = + static_cast(prefix | (static_cast(bin) << shift)); + break; + } + count_above += count; + } + } + __syncthreads(); + } + + if (tx == 0) { + *output_counter = 0; + histogram[0] = 0; + } + __syncthreads(); + + const uint32_t pivot = static_cast(*radix_prefix); + const int equal_count = *radix_remaining; + const int equal_base = static_cast(top_k) - equal_count; + for (int i = tx; i < num_candidates; i += BlockSize) { + const int idx = candidates[i]; + const uint32_t ordered = static_cast(candidate_keys[i]); + if (ordered > pivot) { + output[atomicAdd(output_counter, 1)] = static_cast(idx); + } else if (ordered == pivot) { + const int pos = atomicAdd(&histogram[0], 1); + if (pos < equal_count) + output[equal_base + pos] = static_cast(idx); + } + } + __syncthreads(); + return true; +} + +// Exact selection after the filtered path's coarse bin overflowed the +// candidate stash. The coarse histogram holds suffix counts over the top +// CoarseBits of the ordered key; recovery resumes from its top 11 bits. +// scratch must hold Radix::kBins + 4 ints for the sampled path. +template +__device__ __noinline__ void recover_coarse(const DType* scores, IdType* output, + uint32_t length, uint32_t top_k, + uint32_t num_rows, + const int* coarse_suffix, + uint32_t threshold_bin, + void* dynamic_smem, int* scratch) { + static_assert(CoarseBits >= 11); + constexpr uint32_t kShift = CoarseBits - 11; + const uint32_t prefix_bin = threshold_bin >> kShift; + const uint32_t remaining = + top_k - static_cast(coarse_suffix[(prefix_bin + 1) << kShift]); + // The coarse histogram shares dynamic smem with the recovery buffers. + __syncthreads(); + + // Sampling amortizes its fixed block-sort cost only for long, busy batches. + if (num_rows >= kSamplingMinRows && length >= kSamplingMinLength && + try_sampled_exact_topk(scores, output, length, top_k, + dynamic_smem, scratch)) { + return; + } + topk_histogram_4096::exact_topk_rescan( + scores, output, length, dynamic_smem, prefix_bin << 21, remaining); +} + +} // namespace overflow +} // namespace vllm::filtered_topk + +#endif // PERSISTENT_TOPK_OVERFLOW_CUH_ diff --git a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index 2064775111fb..269e50cfdf4a 100644 --- a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -38,7 +38,7 @@ namespace vllm { template __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, - int32_t num_packed_cols, + int32_t num_packed_cols, int32_t num_padded_cols, Type const* __restrict__ in, float const* __restrict__ SFScale, uint32_t* __restrict__ out, @@ -49,9 +49,16 @@ __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 = (numCols + 63) / 64; + int const sf_m = round_up(numRows, 128); + // Get the global scaling factor, which will be applied to the SF. // Note SFScale is the same as next GEMM's alpha, which is // (448.f / (Alpha_A / 6.f)). @@ -60,9 +67,10 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x; int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD; - // Input tensor row/col loops. - for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { - if (colIdx < num_packed_cols) { + // Iterate over all rows and cols including padded ones - + // ensures we visit every single scale factor address to initialize it. + for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) { + if (colIdx < num_padded_cols) { PackedVec in_vec; PackedVec in_vec2; int64_t inOffset = @@ -70,7 +78,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) + numCols / CVT_FP4_ELTS_PER_THREAD + colIdx; - bool valid = (rowIdx < numRows) && (elem_idx < numCols); + // If we are outside valid rows/columns, feed zeros + bool valid = (rowIdx < numRows) && (colIdx < num_packed_cols); if constexpr (CVT_FP4_PACK16) { ld256_cg_or_zero(reinterpret_cast(in_vec), &reinterpret_cast(in)[inOffset * 8], @@ -95,9 +104,17 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) CVT_FP4_NUM_THREADS_PER_SF>( rowIdx, colIdx, numKTiles, SFout); - auto out_val = - cvt_warp_fp16_to_fp4( - out_silu_mul, SFScaleVal, sf_out); + fp4_packed_t out_val{}; + if (valid) { + out_val = + cvt_warp_fp16_to_fp4( + out_silu_mul, SFScaleVal, sf_out); + } else if (sf_out != nullptr) { + // The swizzled scale tensor may be allocated with torch.empty. + // Explicitly initialize padded rows and columns in this kernel so + // downstream NVFP4 GEMMs never observe stale scale factors. + *sf_out = 0; + } if (valid) { if constexpr (CVT_FP4_PACK16) { @@ -108,7 +125,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) } else { int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; - out[outOffset] = out_val; + reinterpret_cast(out)[outOffset] = out_val; } } } @@ -146,18 +163,41 @@ void silu_and_mul_nvfp4_quant_sm1xxa( int num_packed_cols = int(n / CVT_FP4_ELTS_PER_THREAD); - int grid_y = vllm::div_round_up(num_packed_cols, static_cast(block.x)); - int grid_x = std::min( - int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); + // The swizzled scale tensor is padded to 128 rows x 4 scale columns; the + // grid covers the padding so the kernel can initialize it. + int sf_n_int = int(vllm::round_up(n / CVT_FP4_SF_VEC_SIZE, 4) / 4); + int num_padded_cols = + sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; + + int grid_y = vllm::div_round_up(num_padded_cols, static_cast(block.x)); + int grid_x = + std::min(vllm::computeEffectiveRows(m), + std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); dim3 grid(grid_x, grid_y); + // PDL launch pays off at small M; at large M the early-launched blocks + // slow the chain down, so keep the serialized launch there. + bool const use_pdl = + int64_t(m) * num_packed_cols <= int64_t(multiProcessorCount) * 4096; + VLLM_STABLE_DISPATCH_HALF_TYPES( input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] { using cuda_type = vllm::CUDATypeConverter::Type; auto input_ptr = static_cast(input.data_ptr()); - vllm::silu_mul_cvt_fp16_to_fp4<<>>( - m, n, num_packed_cols, input_ptr, input_sf_ptr, - reinterpret_cast(output_ptr), - reinterpret_cast(sf_out)); + 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 = use_pdl ? 1 : 0; + config.attrs = attrs; + cudaLaunchKernelEx(&config, vllm::silu_mul_cvt_fp16_to_fp4, + m, n, num_packed_cols, num_padded_cols, input_ptr, + input_sf_ptr, + reinterpret_cast(output_ptr), + reinterpret_cast(sf_out)); }); } diff --git a/csrc/quantization/fp4/mxfp4_experts_quant.cu b/csrc/quantization/fp4/mxfp4_experts_quant.cu index 4c6875f0623a..1f9fd7c8a162 100644 --- a/csrc/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/quantization/fp4/mxfp4_experts_quant.cu @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -66,7 +67,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) fp4_packed_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts, uint32_t* output_scale_offset_by_experts, - int n_experts, bool low_latency) { + int n_experts, int32_t const* sorted_expert_ids, + bool low_latency) { using PackedVec = PackedVec; static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); @@ -89,7 +91,10 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int rowIdx_in_expert = 0; int expert_idx = 0; - if constexpr (SMALL_NUM_EXPERTS) { + if (sorted_expert_ids != nullptr) { + expert_idx = __ldca(&sorted_expert_ids[rowIdx]); + rowIdx_in_expert = rowIdx - __ldca(&input_offset_by_experts[expert_idx]); + } else if constexpr (SMALL_NUM_EXPERTS) { for (int i = 0; i < n_experts; i++) { uint32_t current_offset = __ldca(&input_offset_by_experts[i]); uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]); @@ -168,7 +173,7 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) fp4_packed_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts, uint32_t* output_scale_offset_by_experts, - int n_experts) { + int n_experts, int32_t const* sorted_expert_ids) { using PackedVec = PackedVec; static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, "Vec size is not matched."); @@ -208,21 +213,26 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) int rowIdx_in_expert = 0; int expert_idx = 0; - // Binary search through experts using shared memory - int left = 0, right = n_experts - 1; - while (left <= right) { - int mid = (left + right) / 2; - uint32_t mid_offset = shared_input_offsets[mid]; - uint32_t next_offset = shared_input_offsets[mid + 1]; - - if (rowIdx >= mid_offset && rowIdx < next_offset) { - rowIdx_in_expert = rowIdx - mid_offset; - expert_idx = mid; - break; - } else if (rowIdx < mid_offset) { - right = mid - 1; - } else { - left = mid + 1; + if (sorted_expert_ids != nullptr) { + expert_idx = __ldca(&sorted_expert_ids[rowIdx]); + rowIdx_in_expert = rowIdx - shared_input_offsets[expert_idx]; + } else { + // Binary search through experts using shared memory + int left = 0, right = n_experts - 1; + while (left <= right) { + int mid = (left + right) / 2; + uint32_t mid_offset = shared_input_offsets[mid]; + uint32_t next_offset = shared_input_offsets[mid + 1]; + + if (rowIdx >= mid_offset && rowIdx < next_offset) { + rowIdx_in_expert = rowIdx - mid_offset; + expert_idx = mid; + break; + } else if (rowIdx < mid_offset) { + right = mid - 1; + } else { + left = mid + 1; + } } } @@ -259,7 +269,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) template void mxfp4_quant_impl(void* output, void* output_scale, void* input, void* input_offset_by_experts, - void* output_scale_offset_by_experts, int m_topk, int k, + void* output_scale_offset_by_experts, + int32_t const* sorted_expert_ids, int m_topk, int k, int n_experts, cudaStream_t stream) { int multiProcessorCount = get_device_attribute(cudaDevAttrMultiProcessorCount, -1); @@ -288,7 +299,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), - n_experts); + n_experts, sorted_expert_ids); } else { mxfp4_cvt_fp16_to_fp4 <<>>( @@ -297,7 +308,7 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), - n_experts); + n_experts, sorted_expert_ids); } } else { if (n_experts >= 16) { @@ -308,7 +319,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), - n_experts, /* bool low_latency */ true); + n_experts, sorted_expert_ids, + /* bool low_latency */ true); } else { mxfp4_cvt_fp16_to_fp4<<>>( m_topk, k, reinterpret_cast(input), @@ -316,7 +328,8 @@ void mxfp4_quant_impl(void* output, void* output_scale, void* input, reinterpret_cast(output_scale), reinterpret_cast(input_offset_by_experts), reinterpret_cast(output_scale_offset_by_experts), - n_experts, /* bool low_latency */ true); + n_experts, sorted_expert_ids, + /* bool low_latency */ true); } } } @@ -345,12 +358,21 @@ static void validate_mxfp4_experts_quant_inputs( torch::stable::Tensor const& input, torch::stable::Tensor const& input_offset_by_experts, torch::stable::Tensor const& output_scale_offset_by_experts, + std::optional const& sorted_expert_ids, int64_t n_experts, int64_t m_topk, int64_t k) { CHECK_INPUT(output, "output"); CHECK_INPUT(output_scale, "output_scale"); CHECK_INPUT(input, "input"); CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts"); CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts"); + if (sorted_expert_ids.has_value()) { + CHECK_INPUT(sorted_expert_ids.value(), "sorted_expert_ids"); + STD_TORCH_CHECK(sorted_expert_ids->dim() == 1); + STD_TORCH_CHECK(sorted_expert_ids->scalar_type() == INT); + STD_TORCH_CHECK(sorted_expert_ids->size(0) == m_topk); + STD_TORCH_CHECK(sorted_expert_ids->get_device_index() == + input.get_device_index()); + } STD_TORCH_CHECK(output.dim() == 2); STD_TORCH_CHECK(output_scale.dim() == 2); @@ -393,7 +415,8 @@ void mxfp4_experts_quant( torch::stable::Tensor const& input, torch::stable::Tensor const& input_offset_by_experts, torch::stable::Tensor const& output_scale_offset_by_experts, - int64_t n_experts) { + int64_t n_experts, + std::optional const& sorted_expert_ids) { #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED int32_t sm = get_sm_version_num(); STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm), @@ -405,7 +428,7 @@ void mxfp4_experts_quant( validate_mxfp4_experts_quant_inputs( output, output_scale, input, input_offset_by_experts, - output_scale_offset_by_experts, n_experts, m_topk, k); + output_scale_offset_by_experts, sorted_expert_ids, n_experts, m_topk, k); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -417,8 +440,11 @@ void mxfp4_experts_quant( vllm::mxfp4_quant_impl( output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), input_offset_by_experts.data_ptr(), - output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, - stream); + output_scale_offset_by_experts.data_ptr(), + sorted_expert_ids.has_value() + ? static_cast(sorted_expert_ids->data_ptr()) + : nullptr, + m_topk, k, n_experts, stream); }); #else STD_TORCH_CHECK_NOT_IMPLEMENTED(false, @@ -431,7 +457,8 @@ void silu_and_mul_mxfp4_experts_quant( torch::stable::Tensor const& input, torch::stable::Tensor const& input_offset_by_experts, torch::stable::Tensor const& output_scale_offset_by_experts, - int64_t n_experts) { + int64_t n_experts, + std::optional const& sorted_expert_ids) { #if VLLM_MXFP4_EXPERTS_QUANT_SUPPORTED int32_t sm = get_sm_version_num(); STD_TORCH_CHECK(mxfp4_experts_quant_sm_supported(sm), @@ -445,7 +472,7 @@ void silu_and_mul_mxfp4_experts_quant( validate_mxfp4_experts_quant_inputs( output, output_scale, input, input_offset_by_experts, - output_scale_offset_by_experts, n_experts, m_topk, k); + output_scale_offset_by_experts, sorted_expert_ids, n_experts, m_topk, k); const torch::stable::accelerator::DeviceGuard device_guard( input.get_device_index()); @@ -457,8 +484,11 @@ void silu_and_mul_mxfp4_experts_quant( vllm::mxfp4_quant_impl( output.data_ptr(), output_scale.data_ptr(), input.data_ptr(), input_offset_by_experts.data_ptr(), - output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts, - stream); + output_scale_offset_by_experts.data_ptr(), + sorted_expert_ids.has_value() + ? static_cast(sorted_expert_ids->data_ptr()) + : nullptr, + m_topk, k, n_experts, stream); }); #else STD_TORCH_CHECK_NOT_IMPLEMENTED( diff --git a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu index b6043e803306..268650ae1813 100644 --- a/csrc/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/quantization/w8a8/cutlass/moe/moe_data.cu @@ -88,6 +88,7 @@ __global__ void compute_expert_blockscale_offsets( __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids, int32_t* input_permutation, int32_t* output_permutation, + int32_t* sorted_expert_ids, int32_t* atomic_buffer, const int topk_length, const int topk) { int const blk_expert_id = blockIdx.x; @@ -108,6 +109,9 @@ __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids, int start = atomicAdd(&atomic_buffer[expert_id], 1); input_permutation[start] = i / topk; output_permutation[i] = start; + if (sorted_expert_ids != nullptr) { + sorted_expert_ids[start] = expert_id; + } } } } @@ -242,7 +246,8 @@ void get_cutlass_moe_mm_data_caller( torch::stable::Tensor& output_permutation, const int64_t num_experts, const int64_t n, const int64_t k, const std::optional& blockscale_offsets, - const bool is_gated) { + const bool is_gated, + const std::optional& sorted_expert_ids) { auto device = topk_ids.device(); const torch::stable::accelerator::DeviceGuard device_guard(device.index()); auto stream = get_current_cuda_stream(device.index()); @@ -278,6 +283,9 @@ void get_cutlass_moe_mm_data_caller( static_cast(topk_ids.data_ptr()), static_cast(input_permutation.data_ptr()), static_cast(output_permutation.data_ptr()), + sorted_expert_ids.has_value() + ? static_cast(sorted_expert_ids->data_ptr()) + : nullptr, static_cast(atomic_buffer.data_ptr()), topk_ids.numel(), topk_ids.size(1)); } diff --git a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu index 7e4f942f5445..0fd72ebd0824 100644 --- a/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -97,7 +97,8 @@ void get_cutlass_moe_mm_data_caller( torch::stable::Tensor& output_permutation, const int64_t num_experts, const int64_t n, const int64_t k, const std::optional& blockscale_offsets, - const bool is_gated); + const bool is_gated, + const std::optional& sorted_expert_ids); void get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller( const torch::stable::Tensor& expert_first_token_offset, @@ -337,17 +338,18 @@ void get_cutlass_moe_mm_data( torch::stable::Tensor& output_permutation, const int64_t num_experts, const int64_t n, const int64_t k, const std::optional& blockscale_offsets, - const bool is_gated) { + const bool is_gated, + const std::optional& sorted_expert_ids) { // This function currently gets compiled only if we have a valid cutlass moe // mm to run it for. int32_t version_num = get_sm_version_num(); #if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \ (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \ (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120) - get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1, - problem_sizes2, input_permutation, - output_permutation, num_experts, n, k, - blockscale_offsets, is_gated); + get_cutlass_moe_mm_data_caller( + topk_ids, expert_offsets, problem_sizes1, problem_sizes2, + input_permutation, output_permutation, num_experts, n, k, + blockscale_offsets, is_gated, sorted_expert_ids); return; #endif STD_TORCH_CHECK_NOT_IMPLEMENTED( diff --git a/csrc/quantization/w8a8/fp8/common.cu b/csrc/quantization/w8a8/fp8/common.cu index 471499786977..837948fdd2a7 100644 --- a/csrc/quantization/w8a8/fp8/common.cu +++ b/csrc/quantization/w8a8/fp8/common.cu @@ -95,7 +95,7 @@ __global__ void scaled_fp8_quant_kernel_strided_group_shape( int size, float inv_scale) { vectorize_with_alignment( in, out, size, tid, blockDim.x, - [=] __device__(fp8_type & dst, const scalar_t& src) { + [=] __device__(fp8_type& dst, const scalar_t& src) { dst = scaled_fp8_conversion(static_cast(src), inv_scale); }); @@ -194,7 +194,7 @@ __global__ void scaled_fp8_quant_kernel_strided_dynamic( const float reciprocal_scale = 1.0f / (*scale); vectorize_with_alignment<16>( token_in, token_out, hidden_size, tid, blockDim.x, - [=] __device__(fp8_type & dst, const scalar_t& src) { + [=] __device__(fp8_type& dst, const scalar_t& src) { dst = scaled_fp8_conversion(static_cast(src), reciprocal_scale); }); @@ -238,7 +238,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided( // 2) quantize vectorize_with_alignment<16>( token_in, token_out, hidden_size, tid, blockDim.x, - [=] __device__(fp8_type & dst, const scalar_t& src) { + [=] __device__(fp8_type& dst, const scalar_t& src) { dst = scaled_fp8_conversion(static_cast(src), token_scale); }); diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index 40d3f92efb39..162a2d879f67 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -49,7 +49,7 @@ __device__ __forceinline__ float ComputeGroupScale( constexpr int vec_size = 16 / sizeof(T); // copy global -> shared & compute absmax - auto scalar_op_cache = [&] __device__(T & dst, const T& src) { + auto scalar_op_cache = [&] __device__(T& dst, const T& src) { float abs_v = fabsf(static_cast(src)); local_absmax = fmaxf(local_absmax, abs_v); dst = src; @@ -85,7 +85,7 @@ __device__ __forceinline__ void QuantizeGroup( const float inv_y = 1.0f / y_s; // quantize shared -> global 8-bit - auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) { + auto scalar_op_quant = [&] __device__(DST_DTYPE& dst, const T& src) { float q = fminf(fmaxf(static_cast(src) * inv_y, min_8bit), max_8bit); dst = DST_DTYPE(q); }; diff --git a/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu b/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu new file mode 100644 index 000000000000..8cc03ea5397b --- /dev/null +++ b/csrc/rocm/mxfp4_w4a8_gemv_rdna3.cu @@ -0,0 +1,254 @@ +// MXFP4 x int8 (W4A8) GEMV for RDNA3/RDNA3.5 (gfx11) small-batch decode. +// +// Activations are quantized to int8 per 32-element group (symmetric, one fp32 +// scale per group), E2M1 weight codes map exactly to int8 (2 * e2m1), and each +// 32-element group is reduced with eight v_dot4_i32_iu8 instructions. One fp32 +// multiply per group applies the E8M0 weight scale and the activation scale. +// The scheme follows ggml's MMVQ / q8_1 path (llama.cpp, MIT License). +// +// Layouts (Quark OCP MX checkpoints, no repack): +// a [M, K] fp16 / bf16, 1 <= M <= 8 +// b_q [N, K/2] uint8, two E2M1 codes per byte, low nibble = even k +// b_scale [N, K/32] uint8 E8M0 +// out [M, N] same dtype as a + +#include +#include +#include + +#include +#include +#include + +#include + +#include "core/cuda_compat.h" + +#if defined(__GFX11__) + #define __HIP__MXFP4_W4A8_GFX11__ +#endif + +namespace { + +constexpr int kMaxM = 8; +constexpr int kGroup = 32; +constexpr int kWarpsPerBlock = 8; + +template +__device__ __forceinline__ float to_float(T v); +template <> +__device__ __forceinline__ float to_float(half v) { + return __half2float(v); +} +template <> +__device__ __forceinline__ float to_float(__hip_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ half from_float(float v) { + return __float2half(v); +} +template <> +__device__ __forceinline__ __hip_bfloat16 from_float(float v) { + return __float2bfloat16(v); +} + +#if defined(__HIP__MXFP4_W4A8_GFX11__) || !defined(__HIP_DEVICE_COMPILE__) + +// 2 * E2M1 value for each 4-bit code (sign in bit 3). Exact in int8. +__device__ __constant__ int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, + 0, -1, -2, -3, -4, -6, -8, -12}; + +template +__global__ void mxfp4_w4a8_quant_act_kernel(const T* __restrict__ a, + int8_t* __restrict__ a_q, + float* __restrict__ a_s, int M, + int K) { + const int groups = K / kGroup; + const int m = blockIdx.x; + const int g = blockIdx.y * blockDim.x + threadIdx.x; + if (m >= M || g >= groups) return; + + const T* row = a + (int64_t)m * K + g * kGroup; + float v[kGroup]; + float amax = 0.f; + #pragma unroll + for (int i = 0; i < kGroup; ++i) { + v[i] = to_float(row[i]); + amax = fmaxf(amax, fabsf(v[i])); + } + const float scale = fmaxf(amax / 127.0f, 1e-12f); + a_s[(int64_t)m * groups + g] = scale; + + int8_t* out = a_q + (int64_t)m * K + g * kGroup; + #pragma unroll + for (int i = 0; i < kGroup; ++i) { + out[i] = (int8_t)fmaxf(-127.f, fminf(127.f, rintf(v[i] / scale))); + } +} + +// One wave (32 lanes) per output column n. Lanes stride over the K groups; +// each lane decodes its weight group once and reuses it for all M rows. +template +__global__ void mxfp4_w4a8_gemv_rdna3_kernel( + const int* __restrict__ a_q, const float* __restrict__ a_s, + const uint8_t* __restrict__ b_q, const uint8_t* __restrict__ b_scale, + T* __restrict__ out, int N, int K) { + const int groups = K / kGroup; + const int lane = threadIdx.x; + const int n = blockIdx.x * blockDim.y + threadIdx.y; + if (n >= N) return; + + const uint8_t* w_row = b_q + (int64_t)n * (K / 2); + const uint8_t* s_row = b_scale + (int64_t)n * groups; + + float acc[M]; + #pragma unroll + for (int m = 0; m < M; ++m) acc[m] = 0.f; + + for (int g = lane; g < groups; g += 32) { + const uint8_t* gp = w_row + g * (kGroup / 2); + int w4[8]; + #pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t two = gp[2 * j] | (gp[2 * j + 1] << 8); + const int w0 = kE2M1x2[two & 0xF]; + const int w1 = kE2M1x2[(two >> 4) & 0xF]; + const int w2 = kE2M1x2[(two >> 8) & 0xF]; + const int w3 = kE2M1x2[(two >> 12) & 0xF]; + w4[j] = (w0 & 0xFF) | ((w1 & 0xFF) << 8) | ((w2 & 0xFF) << 16) | + ((w3 & 0xFF) << 24); + } + // 2^(e - 127) / 2: undoes the factor 2 folded into kE2M1x2. + const float w_scale = __int_as_float((int)s_row[g] << 23) * 0.5f; + #pragma unroll + for (int m = 0; m < M; ++m) { + const int* ap = a_q + (int64_t)m * (K / 4) + g * (kGroup / 4); + int isum = 0; + #pragma unroll + for (int j = 0; j < 8; ++j) { + isum = __builtin_amdgcn_sudot4(true, w4[j], true, ap[j], isum, false); + } + acc[m] += (float)isum * (w_scale * a_s[(int64_t)m * groups + g]); + } + } + + #pragma unroll + for (int m = 0; m < M; ++m) { + float v = acc[m]; + #pragma unroll + for (int o = 16; o > 0; o >>= 1) v += VLLM_SHFL_XOR_SYNC_WIDTH(v, o, 32); + if (lane == 0) out[(int64_t)m * N + n] = from_float(v); + } +} + +#else // non-gfx11 device pass: empty kernels for symbol parity. + +template +__global__ void mxfp4_w4a8_quant_act_kernel(const T*, int8_t*, float*, int, + int) {} + +template +__global__ void mxfp4_w4a8_gemv_rdna3_kernel(const int*, const float*, + const uint8_t*, const uint8_t*, T*, + int, int) {} + +#endif // __HIP__MXFP4_W4A8_GFX11__ || !__HIP_DEVICE_COMPILE__ + +bool mxfp4_w4a8_on_gfx11() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string arch = dprops->gcnArchName; + return arch.rfind("gfx11", 0) == 0; + }(); + return result; +} + +template +void launch(const at::Tensor& a, const at::Tensor& b_q, + const at::Tensor& b_scale, at::Tensor& out, int M, int N, int K) { + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int groups = K / kGroup; + + auto a_q = torch::empty({M, K}, a.options().dtype(torch::kInt8)); + auto a_s = torch::empty({M, groups}, a.options().dtype(torch::kFloat32)); + + const dim3 q_block(256); + const dim3 q_grid(M, (groups + 255) / 256); + mxfp4_w4a8_quant_act_kernel<<>>( + reinterpret_cast(a.data_ptr()), a_q.data_ptr(), + a_s.data_ptr(), M, K); + + const dim3 block(32, kWarpsPerBlock); + const dim3 grid((N + kWarpsPerBlock - 1) / kWarpsPerBlock); + const int* a_q32 = reinterpret_cast(a_q.data_ptr()); + const float* a_s_ptr = a_s.data_ptr(); + const uint8_t* b_q_ptr = b_q.data_ptr(); + const uint8_t* b_s_ptr = b_scale.data_ptr(); + T* out_ptr = reinterpret_cast(out.data_ptr()); + +#define MXFP4_W4A8_LAUNCH(MM) \ + case MM: \ + mxfp4_w4a8_gemv_rdna3_kernel<<>>( \ + a_q32, a_s_ptr, b_q_ptr, b_s_ptr, out_ptr, N, K); \ + break; + + switch (M) { + MXFP4_W4A8_LAUNCH(1) + MXFP4_W4A8_LAUNCH(2) + MXFP4_W4A8_LAUNCH(3) + MXFP4_W4A8_LAUNCH(4) + MXFP4_W4A8_LAUNCH(5) + MXFP4_W4A8_LAUNCH(6) + MXFP4_W4A8_LAUNCH(7) + MXFP4_W4A8_LAUNCH(8) + } +#undef MXFP4_W4A8_LAUNCH +} + +} // namespace + +torch::Tensor mxfp4_w4a8_gemv_rdna3(const at::Tensor& a, const at::Tensor& b_q, + const at::Tensor& b_scale) { + TORCH_CHECK(a.is_cuda() && b_q.is_cuda() && b_scale.is_cuda(), + "mxfp4_w4a8_gemv_rdna3: all inputs must be on the GPU"); + TORCH_CHECK(mxfp4_w4a8_on_gfx11(), + "mxfp4_w4a8_gemv_rdna3 requires an RDNA3/RDNA3.5 (gfx11) GPU"); + TORCH_CHECK(a.scalar_type() == at::kHalf || a.scalar_type() == at::kBFloat16, + "mxfp4_w4a8_gemv_rdna3: a must be float16 or bfloat16"); + TORCH_CHECK( + b_q.scalar_type() == at::kByte && b_scale.scalar_type() == at::kByte, + "mxfp4_w4a8_gemv_rdna3: b_q and b_scale must be uint8"); + TORCH_CHECK(a.dim() == 2 && b_q.dim() == 2 && b_scale.dim() == 2, + "mxfp4_w4a8_gemv_rdna3: inputs must be 2D"); + TORCH_CHECK( + a.is_contiguous() && b_q.is_contiguous() && b_scale.is_contiguous(), + "mxfp4_w4a8_gemv_rdna3: inputs must be contiguous"); + + const int M = a.size(0); + const int K = a.size(1); + const int N = b_q.size(0); + TORCH_CHECK(M >= 1 && M <= kMaxM, + "mxfp4_w4a8_gemv_rdna3 supports 1 <= M <= ", kMaxM, + ", got M=", M); + TORCH_CHECK(K % kGroup == 0, + "mxfp4_w4a8_gemv_rdna3: K must be a multiple of 32"); + TORCH_CHECK(b_q.size(1) * 2 == K, + "mxfp4_w4a8_gemv_rdna3: b_q must be [N, K/2]"); + TORCH_CHECK(b_scale.size(0) == N && b_scale.size(1) == K / kGroup, + "mxfp4_w4a8_gemv_rdna3: b_scale must be [N, K/32]"); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); + auto out = torch::empty({M, N}, a.options()); + + if (a.scalar_type() == at::kHalf) { + launch(a, b_q, b_scale, out, M, N, K); + } else { + launch<__hip_bfloat16>(a, b_q, b_scale, out, M, N, K); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/rocm/ops.h b/csrc/rocm/ops.h index de963913ea5d..f36cd3d90184 100644 --- a/csrc/rocm/ops.h +++ b/csrc/rocm/ops.h @@ -24,6 +24,9 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b, const at::Tensor& scale_a, const at::Tensor& scale_b, const int64_t CuCount); +torch::Tensor mxfp4_w4a8_gemv_rdna3(const at::Tensor& a, const at::Tensor& b_q, + const at::Tensor& b_scale); + torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight, torch::Tensor b_qzeros, torch::Tensor b_scales, bool use_v2_format); diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index b80a71e36d3e..d1b96abbf6b4 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -30,6 +30,9 @@ static constexpr int get_lds_size_int4() { return 64 * 1024; } +// K must be a multiple of this; checked host-side in wvSplitK_int4_g. +static constexpr int kInt4KAlign = 16; + // Host-side twin of __HIP__GFX115X__, for the tuple choices the dispatch makes. static bool on_gfx115x_int4() { static const bool result = [] { @@ -226,7 +229,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) float f[A_CHUNK / 8]; }; - __shared__ scalar_t s[max_lds_len]; + // k_ is a multiple of A_CHUNK and K of kInt4KAlign, so s[k_ + K * n] is + // aligned to the smaller of the two. Deep-K N = 1 tuples use A_CHUNK = 32. + static constexpr size_t kActLdsAlign = + sizeof(bigTypeA) < kInt4KAlign * sizeof(scalar_t) + ? sizeof(bigTypeA) + : kInt4KAlign * sizeof(scalar_t); + static_assert(kActLdsAlign >= 16, "LDS activation loads must stay b128"); + __shared__ alignas(sizeof(bigTypeA)) scalar_t s[max_lds_len]; // CHUNKED fills s[] per chunk instead, and needs the waves above _WvPrGrp // gone before the first barrier: s_barrier is untagged on gfx11, so a wave @@ -339,7 +349,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (uint32_t p = 0; p < A_CHUNK / ACT_LDS_PIECE; p++) dst[p] = src[p * THRDS]; } else { - bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n]))); + bigA[n][k2] = *((const bigTypeA*)__builtin_assume_aligned( + &(s[k_ + K * n]), kActLdsAlign)); } } } @@ -605,7 +616,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) float f[A_CHUNK / 8]; }; - __shared__ scalar_t s[max_lds_len]; + // k_ is a multiple of A_CHUNK and K of kInt4KAlign, so s[k_ + K * n] is + // aligned to the smaller of the two. Deep-K N = 1 tuples use A_CHUNK = 32. + static constexpr size_t kActLdsAlign = + sizeof(bigTypeA) < kInt4KAlign * sizeof(scalar_t) + ? sizeof(bigTypeA) + : kInt4KAlign * sizeof(scalar_t); + static_assert(kActLdsAlign >= 16, "LDS activation loads must stay b128"); + __shared__ alignas(sizeof(bigTypeA)) scalar_t s[max_lds_len]; uint32_t commitColumn[YTILE]; for (uint32_t i = 0; i < YTILE; i++) { @@ -671,7 +689,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS) for (int n = 0; n < N; n++) { if (k_ + K * n < max_lds_len) - bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n]))); + bigA[n][k2] = *((const bigTypeA*)__builtin_assume_aligned( + &(s[k_ + K * n]), kActLdsAlign)); else bigA[n][k2] = *((const bigTypeA*)(&(A[k_ + K * n]))); } @@ -953,7 +972,7 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_a, const at::Tensor& in_b, num_groups, "] but got [", in_zero_points->size(0), ", ", in_zero_points->size(1), "]"); } - TORCH_CHECK(K_in % 16 == 0, "K must be divisible by 16"); + TORCH_CHECK(K_in % kInt4KAlign == 0, "K must be divisible by ", kInt4KAlign); // Every kernel indexes A linearly, which a non-contiguous activation would // break. diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 077c3f35cef1..5c34d6676b96 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -54,6 +54,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { #endif // VLLM_SKIP_SKINNY_GEMMS #ifdef VLLM_ROCM_GFX1100 + // MXFP4 weights x int8 activations GEMV for RDNA3/RDNA3.5 (gfx11) decode, + // 1 <= M <= 8. Weights [N, K/2] uint8 E2M1, scales [N, K/32] uint8 E8M0. + rocm_ops.def( + "mxfp4_w4a8_gemv_rdna3(Tensor a, Tensor b_q, Tensor b_scale) -> Tensor"); + rocm_ops.impl("mxfp4_w4a8_gemv_rdna3", torch::kCUDA, &mxfp4_w4a8_gemv_rdna3); + // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). rocm_ops.def( "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, " diff --git a/csrc/sampled_topk.cuh b/csrc/sampled_topk.cuh index c25f8ed8e731..9923dcbf417e 100644 --- a/csrc/sampled_topk.cuh +++ b/csrc/sampled_topk.cuh @@ -69,9 +69,8 @@ __device__ void select_bin(Storage& s) { // Find a coarse bin, then refine its full FP32 keys one byte at a time. // Read shared candidates normally, or the full row for exact fallback. -template -__device__ uint32_t threshold(Storage& s, const float* row, int count, - int rank) { +template +__device__ uint32_t threshold(Storage& s, const InT* row, int count, int rank) { const int tid = threadIdx.x; if (tid == 0) { s.remaining = rank; @@ -81,7 +80,8 @@ __device__ uint32_t threshold(Storage& s, const float* row, int count, __syncthreads(); for (int i = tid; i < count; i += kThreads) { const uint32_t key = - Buffered ? s.candidates[i].y : hist4096::convert_to_uint32_v2(row[i]); + Buffered ? s.candidates[i].y + : hist4096::convert_to_uint32_v2(hist4096::to_float(row[i])); atomicAdd(&s.histogram[coarse(key)], 1); } __syncthreads(); @@ -94,7 +94,8 @@ __device__ uint32_t threshold(Storage& s, const float* row, int count, const uint32_t prefix = s.prefix; for (int i = tid; i < count; i += kThreads) { const uint32_t key = - Buffered ? s.candidates[i].y : hist4096::convert_to_uint32_v2(row[i]); + Buffered ? s.candidates[i].y + : hist4096::convert_to_uint32_v2(hist4096::to_float(row[i])); if (coarse(key) == coarse_bin && (key & mask) == prefix) { atomicAdd(&s.histogram[(key >> shift) & 255], 1); } @@ -118,8 +119,8 @@ __device__ __forceinline__ int reserve(bool hit, int* counter) { return base + __popc(mask & ((1u << lane) - 1)); } -template -__device__ void emit(Storage& s, const float* row, int32_t* dst, int count, +template +__device__ void emit(Storage& s, const InT* row, int32_t* dst, int count, uint32_t cutoff) { const int tid = threadIdx.x; if (tid == 0) s.count = 0; @@ -127,10 +128,10 @@ __device__ void emit(Storage& s, const float* row, int32_t* dst, int count, for (int base = 0; base < count; base += kThreads) { const int i = base + tid; const uint2 pair = - i < count - ? (Buffered ? s.candidates[i] - : make_uint2(i, hist4096::convert_to_uint32_v2(row[i]))) - : make_uint2(0, 0); + i < count ? (Buffered ? s.candidates[i] + : make_uint2(i, hist4096::convert_to_uint32_v2( + hist4096::to_float(row[i])))) + : make_uint2(0, 0); bool keep = i < count && pair.y > cutoff; if (i < count && pair.y == cutoff) { keep = atomicSub(&s.remaining, 1) > 0; @@ -140,9 +141,9 @@ __device__ void emit(Storage& s, const float* row, int32_t* dst, int count, } } -template +template __global__ void __launch_bounds__(kThreads) - sampled_topk_kernel(const float* __restrict__ input, + sampled_topk_kernel(const InT* __restrict__ input, const int32_t* __restrict__ lengths, int32_t* __restrict__ output, int64_t stride, int max_length) { @@ -150,7 +151,7 @@ __global__ void __launch_bounds__(kThreads) auto& s = *reinterpret_cast(smem); const int tid = threadIdx.x; const int length = max(0, min(lengths[blockIdx.x], max_length)); - const float* row = input + blockIdx.x * stride; + const InT* row = input + blockIdx.x * stride; int32_t* dst = output + blockIdx.x * K; // Graph allocations can hold rows much shorter than the host length bound. if (length < kMinSampledLength) { @@ -158,11 +159,11 @@ __global__ void __launch_bounds__(kThreads) bool selected; if ((stride & 3) == 0 && (reinterpret_cast(row) & 15) == 0) { selected = - filtered_topk::filtered_topk_row( + filtered_topk::filtered_topk_row( row, dst, length, K, fallback_storage); } else { selected = - filtered_topk::filtered_topk_row( + filtered_topk::filtered_topk_row( row, dst, length, K, fallback_storage); } if (!selected) { @@ -177,7 +178,8 @@ __global__ void __launch_bounds__(kThreads) // Sample contiguous warps spread across the row to keep loads coalesced. for (int i = tid; i < kSample; i += kThreads) { const int index = (i / 32) * (length / (kSample / 32)) + i % 32; - atomicAdd(&s.histogram[coarse(hist4096::convert_to_uint32_v2(row[index]))], + atomicAdd(&s.histogram[coarse(hist4096::convert_to_uint32_v2( + hist4096::to_float(row[index])))], 1); } __syncthreads(); @@ -200,7 +202,7 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int j = 0; j < kItems / 4; ++j) { const float4 values = - reinterpret_cast(row + base)[tid + j * kThreads]; + hist4096::load_score4(row + base + (tid + j * kThreads) * 4); keys[4 * j] = hist4096::convert_to_uint32_v2(values.x); keys[4 * j + 1] = hist4096::convert_to_uint32_v2(values.y); keys[4 * j + 2] = hist4096::convert_to_uint32_v2(values.z); @@ -210,7 +212,10 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int j = 0; j < kItems; ++j) { const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; - keys[j] = i < length ? hist4096::convert_to_uint32_v2(row[i]) : 0; + keys[j] = + i < length + ? hist4096::convert_to_uint32_v2(hist4096::to_float(row[i])) + : 0; } } #pragma unroll diff --git a/csrc/topk.cu b/csrc/topk.cu index fb3d946b96b0..1322ca3b0dd4 100644 --- a/csrc/topk.cu +++ b/csrc/topk.cu @@ -15,7 +15,7 @@ namespace { #ifndef USE_ROCM -template +template void launch_persistent_topk(const torch::stable::Tensor& logits, const torch::stable::Tensor& lengths, torch::stable::Tensor& output, @@ -28,6 +28,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, const int64_t num_rows = logits.size(0); const int64_t stride = logits.stride(0); const cudaStream_t stream = get_current_cuda_stream(); + const auto* input = reinterpret_cast(logits.const_data_ptr()); static int num_sms = 0; static int max_smem_per_block = 0; @@ -37,11 +38,15 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, max_smem_per_block = device_prop->sharedMemPerBlockOptin; } + STD_TORCH_CHECK( + sizeof(InT) == sizeof(float) || max_smem_per_block >= 128 * 1024, + "bfloat16 persistent_topk requires 128 KiB shared memory"); + // Allow static fallback storage in addition to the 128 KiB dynamic buffer. if (num_rows > 64 && max_seq_len >= vllm::sampled_topk::kMinSampledLength && max_smem_per_block >= 144 * 1024) { - auto kernel = vllm::sampled_topk::sampled_topk_kernel; + auto kernel = vllm::sampled_topk::sampled_topk_kernel; constexpr size_t smem_size = vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC; cudaError_t status = cudaFuncSetAttribute( @@ -49,15 +54,14 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, STD_TORCH_CHECK(status == cudaSuccess, "sampled_topk smem failed: ", cudaGetErrorString(status)); kernel<<>>( - logits.const_data_ptr(), lengths.const_data_ptr(), + input, lengths.const_data_ptr(), output.mutable_data_ptr(), stride, static_cast(std::min(max_seq_len, logits.size(1)))); } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { - cudaError_t status = - vllm::FilteredTopKRaggedTransform( - logits.const_data_ptr(), output.mutable_data_ptr(), - lengths.const_data_ptr(), static_cast(num_rows), - static_cast(TopK), static_cast(stride), stream); + cudaError_t status = vllm::FilteredTopKRaggedTransform( + input, output.mutable_data_ptr(), + lengths.const_data_ptr(), static_cast(num_rows), + static_cast(TopK), static_cast(stride), stream); STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK failed: ", cudaGetErrorString(status)); } else { @@ -82,10 +86,12 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, uint32_t max_chunk_elements = static_cast(available_for_ordered / sizeof(uint32_t)); + // Views like x[:, 1:] keep the stride but shift the base. + const auto base = reinterpret_cast(input); uint32_t vec_size = 1; - if (stride % 4 == 0) + if (stride % 4 == 0 && base % 16 == 0) vec_size = 4; - else if (stride % 2 == 0) + else if (stride % 2 == 0 && base % 8 == 0) vec_size = 2; max_chunk_elements = (max_chunk_elements / vec_size) * vec_size; @@ -109,16 +115,16 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, cudaError_t occ_err = cudaSuccess; if (vec_size == 4) { occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, - smem_size); + &occupancy, P::persistent_topk_kernel, + P::kThreadsPerBlock, smem_size); } else if (vec_size == 2) { occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, - smem_size); + &occupancy, P::persistent_topk_kernel, + P::kThreadsPerBlock, smem_size); } else { occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, - smem_size); + &occupancy, P::persistent_topk_kernel, + P::kThreadsPerBlock, smem_size); } STD_TORCH_CHECK(occ_err == cudaSuccess, "persistent_topk occupancy query failed: ", @@ -159,9 +165,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, return; } cudaError_t status = - vllm::FilteredTopKRaggedTransform( - logits.const_data_ptr(), - output.mutable_data_ptr(), + vllm::FilteredTopKRaggedTransform( + input, output.mutable_data_ptr(), lengths.const_data_ptr(), static_cast(num_rows), static_cast(TopK), static_cast(stride), stream); @@ -206,8 +211,8 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, "row_states memset failed: ", cudaGetErrorString(mz_err)); } - P::PersistentTopKParams params; - params.input = logits.const_data_ptr(); + P::PersistentTopKParams params; + params.input = input; params.output = output.mutable_data_ptr(); params.lengths = lengths.const_data_ptr(); params.num_rows = static_cast(num_rows); @@ -221,7 +226,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \ do { \ - auto kernel = &P::persistent_topk_kernel; \ + auto kernel = &P::persistent_topk_kernel; \ cudaError_t err = cudaFuncSetAttribute( \ kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \ STD_TORCH_CHECK(err == cudaSuccess, \ @@ -243,6 +248,21 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, STD_TORCH_CHECK(err == cudaSuccess, "persistent_topk failed: ", cudaGetErrorString(err)); } + +template +void launch_persistent_topk_dtype(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, + int64_t max_seq_len) { + if (logits.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + launch_persistent_topk(logits, lengths, output, + workspace, max_seq_len); + } else { + launch_persistent_topk(logits, lengths, output, workspace, + max_seq_len); + } +} #endif } // anonymous namespace @@ -256,8 +276,14 @@ void persistent_topk(const torch::stable::Tensor& logits, STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); - STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float, - "Only float32 supported"); + STD_TORCH_CHECK( + logits.scalar_type() == torch::headeronly::ScalarType::Float || + logits.scalar_type() == torch::headeronly::ScalarType::BFloat16, + "logits must be float32 or bfloat16"); + STD_TORCH_CHECK( + logits.scalar_type() != torch::headeronly::ScalarType::BFloat16 || + logits.stride(1) == 1, + "bfloat16 logits columns must be contiguous"); STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int, "lengths must be int32"); STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int, @@ -281,14 +307,14 @@ void persistent_topk(const torch::stable::Tensor& logits, logits.get_device_index()); if (k == 512) { - launch_persistent_topk<512>(logits, lengths, output, workspace, - max_seq_len); + launch_persistent_topk_dtype<512>(logits, lengths, output, workspace, + max_seq_len); } else if (k == 1024) { - launch_persistent_topk<1024>(logits, lengths, output, workspace, - max_seq_len); + launch_persistent_topk_dtype<1024>(logits, lengths, output, workspace, + max_seq_len); } else { - launch_persistent_topk<2048>(logits, lengths, output, workspace, - max_seq_len); + launch_persistent_topk_dtype<2048>(logits, lengths, output, workspace, + max_seq_len); } #else STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); diff --git a/csrc/topk_histogram_4096.cuh b/csrc/topk_histogram_4096.cuh index 539eaddffbe8..86e7b680461e 100644 --- a/csrc/topk_histogram_4096.cuh +++ b/csrc/topk_histogram_4096.cuh @@ -5,6 +5,7 @@ #ifndef TOPK_HISTOGRAM_4096_CUH_ #define TOPK_HISTOGRAM_4096_CUH_ +#include #include #include #include @@ -66,6 +67,39 @@ __device__ __forceinline__ void load_float4_predicated(const float* ptr, v3 = __uint_as_float(r3); } +__device__ __forceinline__ float to_float(float x) { return x; } +__device__ __forceinline__ float to_float(__nv_bfloat16 x) { + return __bfloat162float(x); +} + +__device__ __forceinline__ float4 load_score4(const float* ptr) { + return *reinterpret_cast(ptr); +} +__device__ __forceinline__ float4 load_score4(const __nv_bfloat16* ptr) { + const uint2 raw = *reinterpret_cast(ptr); + return make_float4( + __uint_as_float(raw.x << 16), __uint_as_float(raw.x & 0xFFFF0000u), + __uint_as_float(raw.y << 16), __uint_as_float(raw.y & 0xFFFF0000u)); +} +__device__ __forceinline__ float2 load_score2(const float* ptr) { + return *reinterpret_cast(ptr); +} +__device__ __forceinline__ float2 load_score2(const __nv_bfloat16* ptr) { + const uint32_t raw = *reinterpret_cast(ptr); + return make_float2(__uint_as_float(raw << 16), + __uint_as_float(raw & 0xFFFF0000u)); +} + +__device__ __forceinline__ void load_float4_predicated(const __nv_bfloat16* ptr, + int base, int seq_len, + float& v0, float& v1, + float& v2, float& v3) { + v0 = base < seq_len ? to_float(ptr[0]) : __uint_as_float(0xFF800000u); + v1 = base + 1 < seq_len ? to_float(ptr[1]) : __uint_as_float(0xFF800000u); + v2 = base + 2 < seq_len ? to_float(ptr[2]) : __uint_as_float(0xFF800000u); + v3 = base + 3 < seq_len ? to_float(ptr[3]) : __uint_as_float(0xFF800000u); +} + // converts the float32 score to a 32-bit ordered unsigned integer — the full // precision key for radix sorting __device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { @@ -333,18 +367,33 @@ struct Histogram4096Smem { // Defined in topk_histogram_4096_overflow.cuh. Only kFp32Rescan instantiations // call it, and their callers include that header. -template +template __device__ __noinline__ void exact_topk_rescan_cold( - const float* __restrict__ scores, int32_t* __restrict__ output, + const InT* __restrict__ scores, int32_t* __restrict__ output, uint32_t length, void* smem); -enum class OverflowRecovery { kDisabled, kFp32Rescan }; +// kRegisterValues refines from the values already held in registers and +// falls back to the FP32 rescan when the coarse bin cannot bound them. +enum class OverflowRecovery { kFp32Rescan, kRegisterValues }; + +// Defined in topk_histogram_4096_overflow.cuh, like exact_topk_rescan_cold. +template +__device__ __forceinline__ bool exact_topk_refine_registers( + const float4 (&vecs)[VecsPerThread], int32_t* __restrict__ output, + uint32_t length, uint32_t coarse_bin, uint32_t num_above, void* _smem); + +template +__device__ void exact_topk_rescan(const InT* __restrict__ scores, + int32_t* __restrict__ output, uint32_t length, + void* _smem, uint32_t initial_prefix, + uint32_t initial_remaining); template -__device__ void histogram_4096_topk(const float* __restrict__ scores, + bool UsePredicatedLoads = false, OverflowRecovery Recovery, + typename InT> +__device__ void histogram_4096_topk(const InT* __restrict__ scores, int32_t* __restrict__ output, uint32_t length, void* _smem) { constexpr uint32_t HIST_BINS = 1 << HIST_BITS; @@ -374,14 +423,14 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, smem->counter_gt = 0; smem->counter_eq = 0; } - if constexpr (UsePredicatedLoads) { + if constexpr (UsePredicatedLoads || sizeof(InT) == 2) { const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; #pragma unroll for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { const uint32_t base = (tx + v * kBlockSize) * 4; if (base < length) { if (row_aligned && base + 3 < length) { - vecs[v] = *reinterpret_cast(scores + base); + vecs[v] = load_score4(scores + base); } else { load_float4_predicated(scores + base, static_cast(base), static_cast(length), vecs[v].x, vecs[v].y, @@ -394,7 +443,7 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { const uint32_t base = (tx + v * kBlockSize) * 4; if (base < length) { - vecs[v] = *reinterpret_cast(scores + base); + vecs[v] = load_score4(scores + base); } } } @@ -468,6 +517,16 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, return; } } + if constexpr (Recovery == OverflowRecovery::kRegisterValues) { + if (need_tie && num_equal > TopK) { + if (!exact_topk_refine_registers( + vecs, output, length, thr_bin, num_above, _smem)) { + exact_topk_rescan( + scores, output, length, _smem, 0, TopK); + } + return; + } + } done = false; #pragma unroll @@ -571,12 +630,13 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, } template + uint32_t VECS_PER_THREAD = kHist4096VecsPerThread, + OverflowRecovery Recovery, typename InT> __device__ __noinline__ void histogram_4096_topk_predicated( - const float* __restrict__ scores, int32_t* __restrict__ output, + const InT* __restrict__ scores, int32_t* __restrict__ output, uint32_t length, void* _smem) { - histogram_4096_topk(scores, output, - length, _smem); + histogram_4096_topk( + scores, output, length, _smem); } } // namespace topk_histogram_4096 diff --git a/csrc/topk_histogram_4096_overflow.cuh b/csrc/topk_histogram_4096_overflow.cuh index 617d312c79d3..f788cadac2c9 100644 --- a/csrc/topk_histogram_4096_overflow.cuh +++ b/csrc/topk_histogram_4096_overflow.cuh @@ -9,15 +9,15 @@ namespace vllm { namespace topk_histogram_4096 { // Visit a row with vectorized loads and a scalar tail. -template -__device__ __forceinline__ void scan_scores(const float* __restrict__ scores, +template +__device__ __forceinline__ void scan_scores(const InT* __restrict__ scores, uint32_t length, Visit visit) { static_assert(VecSize == 1 || VecSize == 2 || VecSize == 4); const uint32_t tx = threadIdx.x; if constexpr (VecSize == 1) { for (uint32_t idx = tx; idx < length; idx += BlockSize) { - visit(idx, scores[idx]); + visit(idx, to_float(scores[idx])); } return; } @@ -27,20 +27,20 @@ __device__ __forceinline__ void scan_scores(const float* __restrict__ scores, for (uint32_t base = tx * VecSize; base < aligned_length; base += BlockSize * VecSize) { if constexpr (VecSize == 4) { - const float4 values = *reinterpret_cast(scores + base); + const float4 values = load_score4(scores + base); visit(base, values.x); visit(base + 1, values.y); visit(base + 2, values.z); visit(base + 3, values.w); } else { - const float2 values = *reinterpret_cast(scores + base); + const float2 values = load_score2(scores + base); visit(base, values.x); visit(base + 1, values.y); } } for (uint32_t idx = aligned_length + tx; idx < length; idx += BlockSize) { - visit(idx, scores[idx]); + visit(idx, to_float(scores[idx])); } } @@ -62,36 +62,52 @@ struct ExactRadixTraits { } }; +// Final-round candidates are staged past the rescan's own shared state. +constexpr uint32_t kExactCandidateOffset = 16 * 1024; + // Exact bounded-memory radix selection. Each round rescans the row, so memory -// use and correctness are independent of coarse-bin occupancy. -template -__device__ void exact_topk_rescan(const float* __restrict__ scores, +// use and correctness are independent of coarse-bin occupancy. Recovery can +// resume at StartRound from a known prefix, and with FinalCandidateCapacity +// the last round stages its candidates so the output pass skips a row scan. +template +__device__ void exact_topk_rescan(const InT* __restrict__ scores, int32_t* __restrict__ output, uint32_t length, - void* _smem) { + void* _smem, uint32_t initial_prefix = 0, + uint32_t initial_remaining = TopK) { static_assert(VecSize == 1 || VecSize == 2 || VecSize == 4); using Radix = ExactRadixTraits; constexpr uint32_t kHistogramBins = Radix::kBins; static_assert(kHistogramBins % BlockSize == 0); - using BlockScan = cub::BlockScan; + static_assert(StartRound < Radix::kRounds); + static_assert(BlockSize == kWarpSize * kWarpSize, "one warp sum per lane"); struct ExactSmem { uint32_t histogram[kHistogramBins]; - typename BlockScan::TempStorage scan; + uint32_t warp_sum[BlockSize / kWarpSize]; uint32_t prefix; uint32_t remaining; uint32_t output_counter; + uint32_t candidate_count; + uint32_t pivot_count; }; + static_assert(FinalCandidateCapacity == 0 || + sizeof(ExactSmem) <= kExactCandidateOffset); + static_assert(FinalCandidateCapacity == 0 || TopK <= FinalCandidateCapacity); auto* smem = static_cast(_smem); + auto* final_candidates = reinterpret_cast( + static_cast(_smem) + kExactCandidateOffset); const uint32_t tx = threadIdx.x; if (tx == 0) { - smem->prefix = 0; - smem->remaining = TopK; + smem->prefix = initial_prefix; + smem->remaining = initial_remaining; } __syncthreads(); #pragma unroll - for (uint32_t round = 0; round < Radix::kRounds; ++round) { + for (uint32_t round = StartRound; round < Radix::kRounds; ++round) { for (uint32_t bin = tx; bin < kHistogramBins; bin += BlockSize) { smem->histogram[bin] = 0; } @@ -101,13 +117,39 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores, const uint32_t digit_mask = Radix::digit_mask(round); const uint32_t prefix_mask = Radix::prefix_mask(round); const uint32_t prefix = smem->prefix; + const bool final_round = round == Radix::kRounds - 1; - scan_scores(scores, length, [&](uint32_t, float score) { - const uint32_t ordered = convert_to_uint32_v2(score); - if ((ordered & prefix_mask) == prefix) { - atomicAdd(&smem->histogram[(ordered >> shift) & digit_mask], 1); + if constexpr (FinalCandidateCapacity > 0) { + if (final_round) { + if (tx == 0) { + smem->output_counter = 0; + smem->candidate_count = 0; + } + __syncthreads(); } - }); + } + + scan_scores( + scores, length, [&](uint32_t idx, float score) { + const uint32_t ordered = convert_to_uint32_v2(score); + const uint32_t ordered_prefix = ordered & prefix_mask; + if (ordered_prefix == prefix) { + atomicAdd(&smem->histogram[(ordered >> shift) & digit_mask], 1); + if constexpr (FinalCandidateCapacity > 0) { + if (final_round) { + const uint32_t pos = atomicAdd(&smem->candidate_count, 1); + if (pos < FinalCandidateCapacity) { + final_candidates[pos] = static_cast(idx); + } + } + } + } else if constexpr (FinalCandidateCapacity > 0) { + if (final_round && ordered_prefix > prefix) { + const uint32_t pos = atomicAdd(&smem->output_counter, 1); + output[pos] = static_cast(idx); + } + } + }); __syncthreads(); constexpr uint32_t kBinsPerThread = kHistogramBins / BlockSize; @@ -119,12 +161,19 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores, local_sum += counts[i]; } - uint32_t lower_prefix; - uint32_t total; - // Snapshot before the block scan; its barrier precedes any update. + // Shuffle scan: lighter on registers than a CUB block scan. Snapshot + // remaining first; the barrier below precedes any update. const uint32_t remaining = smem->remaining; - BlockScan(smem->scan).ExclusiveSum(local_sum, lower_prefix, total); - uint32_t count_above = total - lower_prefix - local_sum; + const uint32_t lane = tx % kWarpSize; + const uint32_t warp = tx / kWarpSize; + const uint32_t warp_inclusive = warp_inclusive_sum(lane, local_sum); + if (lane == kWarpSize - 1) smem->warp_sum[warp] = warp_inclusive; + __syncthreads(); + const uint32_t warp_total = smem->warp_sum[lane]; + uint32_t count_above = + warp_reduce_sum_full(lane > warp ? warp_total : 0) + + __shfl_sync(0xFFFFFFFFu, warp_inclusive, kWarpSize - 1) - + warp_inclusive; #pragma unroll for (int i = static_cast(kBinsPerThread) - 1; i >= 0; --i) { const uint32_t count = counts[i]; @@ -132,6 +181,7 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores, const uint32_t bin = tx * kBinsPerThread + i; smem->remaining = remaining - count_above; smem->prefix = prefix | (bin << shift); + if constexpr (FinalCandidateCapacity > 0) smem->pivot_count = count; } count_above += count; } @@ -140,32 +190,60 @@ __device__ void exact_topk_rescan(const float* __restrict__ scores, const uint32_t pivot = smem->prefix; if (tx == 0) { - smem->output_counter = 0; + if constexpr (FinalCandidateCapacity == 0) smem->output_counter = 0; smem->histogram[0] = 0; } __syncthreads(); const uint32_t equal_count = smem->remaining; const uint32_t equal_base = TopK - equal_count; - scan_scores( - scores, length, [&](uint32_t idx, float score) { - const uint32_t ordered = convert_to_uint32_v2(score); - if (ordered > pivot) { - const uint32_t pos = atomicAdd(&smem->output_counter, 1); - output[pos] = static_cast(idx); - } else if (ordered == pivot) { - const uint32_t pos = atomicAdd(&smem->histogram[0], 1); - if (pos < equal_count) { - output[equal_base + pos] = static_cast(idx); - } - } - }); + const auto collect = [&](uint32_t idx, float score) { + const uint32_t ordered = convert_to_uint32_v2(score); + if (ordered > pivot) { + const uint32_t pos = atomicAdd(&smem->output_counter, 1); + output[pos] = static_cast(idx); + } else if (ordered == pivot) { + const uint32_t pos = atomicAdd(&smem->histogram[0], 1); + if (pos < equal_count) { + output[equal_base + pos] = static_cast(idx); + } + } + }; + + if constexpr (FinalCandidateCapacity > 0) { + // Keys above the final round's prefix are already in the output. + const uint32_t candidate_count = smem->candidate_count; + if (candidate_count <= FinalCandidateCapacity) { + for (uint32_t i = tx; i < candidate_count; i += BlockSize) { + const uint32_t idx = final_candidates[i]; + collect(idx, to_float(scores[idx])); + } + } else if (smem->pivot_count == candidate_count) { + // Every key under the final prefix equals the pivot (heavy ties), so + // any equal_count staged candidates complete the output. + for (uint32_t i = tx; i < equal_count; i += BlockSize) { + output[equal_base + i] = final_candidates[i]; + } + } else { + constexpr uint32_t kFinalPrefixMask = + Radix::prefix_mask(Radix::kRounds - 1); + scan_scores( + scores, length, [&](uint32_t idx, float score) { + if ((convert_to_uint32_v2(score) & kFinalPrefixMask) == + (pivot & kFinalPrefixMask)) { + collect(idx, score); + } + }); + } + } else { + scan_scores(scores, length, collect); + } __syncthreads(); } -template +template __device__ __noinline__ void exact_topk_rescan_cold( - const float* __restrict__ scores, int32_t* __restrict__ output, + const InT* __restrict__ scores, int32_t* __restrict__ output, uint32_t length, void* smem) { exact_topk_rescan(scores, output, length, smem); } @@ -191,6 +269,189 @@ coarse_refine_range(uint32_t coarse_bin) { .finite = (lower_half_bits & 0x7C00u) != 0x7C00u}; } +constexpr uint32_t kCoarseRefineBins = 2048; +constexpr uint32_t kCoarseRefineDigitBits = 11; +using CoarseRefineScan = cub::BlockScan; + +struct CoarseRefineSmem { + uint32_t histogram[kCoarseRefineBins]; + CoarseRefineScan::TempStorage scan; + uint32_t prefix; + uint32_t remaining; + uint32_t output_counter; + uint32_t equal_counter; + uint32_t valid; + uint32_t key_min; + uint32_t key_max; +}; + +__device__ __forceinline__ uint32_t warp_reduce_min_full(uint32_t v) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_min_sync(0xFFFFFFFFu, v); +#else + #pragma unroll + for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) { + v = min(v, __shfl_xor_sync(0xFFFFFFFFu, v, mask)); + } + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_reduce_max_full(uint32_t v) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_max_sync(0xFFFFFFFFu, v); +#else + #pragma unroll + for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) { + v = max(v, __shfl_xor_sync(0xFFFFFFFFu, v, mask)); + } + return v; +#endif +} + +__device__ __forceinline__ void select_coarse_refine_digit( + CoarseRefineSmem* smem, uint32_t prefix, uint32_t shift) { + const uint32_t bin0 = 2 * threadIdx.x; + const uint32_t count0 = smem->histogram[bin0]; + const uint32_t count1 = smem->histogram[bin0 + 1]; + const uint32_t local_sum = count0 + count1; + const uint32_t remaining = smem->remaining; + uint32_t lower_prefix; + uint32_t total; + CoarseRefineScan(smem->scan).ExclusiveSum(local_sum, lower_prefix, total); + uint32_t count_above = total - lower_prefix - local_sum; + if (count_above < remaining && count_above + count1 >= remaining) { + smem->remaining = remaining - count_above; + smem->prefix = prefix | ((bin0 + 1) << shift); + } + count_above += count1; + if (count_above < remaining && count_above + count0 >= remaining) { + smem->remaining = remaining - count_above; + smem->prefix = prefix | (bin0 << shift); + } +} + +// Refine an overflowing FP16 histogram bin from values already resident in +// registers. A 12-bit FP16 bin spans at most 22 ordered-FP32 bits for finite +// values, so two radix-2048 passes select the exact pivot. A bin holding a +// single key (e.g. all-equal scores) needs no refinement: that key is the +// pivot, so the first pass also tracks the bin's min and max key. +template +__device__ __forceinline__ bool exact_topk_refine_registers( + const float4 (&vecs)[VecsPerThread], int32_t* __restrict__ output, + uint32_t length, uint32_t coarse_bin, uint32_t num_above, void* _smem) { + constexpr uint32_t kBins = kCoarseRefineBins; + constexpr uint32_t kDigitBits = kCoarseRefineDigitBits; + constexpr uint32_t kDigitMask = kBins - 1; + constexpr uint32_t kDeltaLimit = 1u << (2 * kDigitBits); + static_assert(HistBits == 12); + static_assert(kBins == 2 * kBlockSize); + auto* smem = static_cast(_smem); + const uint32_t tx = threadIdx.x; + const auto range = coarse_refine_range(coarse_bin); + const uint32_t base_key = range.base_key; + + if (tx == 0) { + smem->prefix = 0; + smem->remaining = TopK - num_above; + smem->valid = range.finite; + smem->key_min = 0xFFFFFFFFu; + smem->key_max = 0; + } + __syncthreads(); + +#pragma unroll + for (uint32_t round = 0; round < 2; ++round) { + for (uint32_t bin = tx; bin < kBins; bin += kBlockSize) { + smem->histogram[bin] = 0; + } + __syncthreads(); + + const uint32_t prefix = smem->prefix; + const uint32_t shift = round == 0 ? kDigitBits : 0; + uint32_t key_min = 0xFFFFFFFFu; + uint32_t key_max = 0; + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VecsPerThread && !done; ++v) { + const float* values = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; ++e) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else if (extract_coarse_bin_N(values[e]) == coarse_bin) { + const uint32_t key = convert_to_uint32_v2(values[e]); + const uint32_t delta = key - base_key; + if (round == 0) { + key_min = min(key_min, key); + key_max = max(key_max, key); + } + if (key < base_key || delta >= kDeltaLimit) { + atomicExch(&smem->valid, 0); + } else if (round == 0 || (delta & ~kDigitMask) == prefix) { + atomicAdd(&smem->histogram[(delta >> shift) & kDigitMask], 1); + } + } + } + } + if (round == 0) { + key_min = warp_reduce_min_full(key_min); + key_max = warp_reduce_max_full(key_max); + if (tx % kWarpSize == 0) { + atomicMin(&smem->key_min, key_min); + atomicMax(&smem->key_max, key_max); + } + } + __syncthreads(); + + if (smem->valid == 0) return false; + + if (round == 0 && smem->key_min == smem->key_max) { + if (tx == 0) smem->prefix = smem->key_min - base_key; + __syncthreads(); + break; + } + + select_coarse_refine_digit(smem, prefix, shift); + __syncthreads(); + } + + if (tx == 0) { + smem->output_counter = 0; + smem->equal_counter = 0; + } + __syncthreads(); + + const uint32_t pivot = base_key + smem->prefix; + const uint32_t equal_base = TopK - smem->remaining; + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VecsPerThread && !done; ++v) { + const float* values = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; ++e) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t key = convert_to_uint32_v2(values[e]); + if (key > pivot) { + const uint32_t pos = atomicAdd(&smem->output_counter, 1); + output[pos] = static_cast(idx); + } else if (key == pivot) { + const uint32_t pos = atomicAdd(&smem->equal_counter, 1); + if (equal_base + pos < TopK) { + output[equal_base + pos] = static_cast(idx); + } + } + } + } + } + __syncthreads(); + return true; +} + } // namespace topk_histogram_4096 } // namespace vllm diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a64d844e4d77..9969f6309b8f 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -8,6 +8,17 @@ // Note: We register under namespace "_C" so ops are accessible as // torch.ops._C. for compatibility with existing code. STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { +#ifdef VLLM_ENABLE_FUSED_MQA_TOPK + // Single-launch fused FP8 MQA-logits + top-k (512/1024/2048) prefill + // selection (SM100 family). Workspace tensors are caller-owned; counter must + // be zero. + ops.def( + "fused_mqa_topk_prefill(Tensor q, Tensor k, Tensor scales, " + "Tensor weights, Tensor starts, Tensor ends, Tensor! out, " + "Tensor! counter, Tensor! values, Tensor! indices, Tensor! retained, " + "Tensor! retained_indices, Tensor! counts, Tensor! defer) -> ()"); +#endif + // Compute per-token-group FP8 quantized tensor and scaling factor. // The dummy arguments are here so we can correctly fuse with RMSNorm. ops.def( @@ -151,7 +162,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { " Tensor! input_permutation, " " Tensor! output_permutation, int num_experts, " " int n, int k, Tensor? blockscale_offsets, " - " bool is_gated) -> ()"); + " bool is_gated, " + " Tensor!? sorted_expert_ids=None) -> ()"); // compute per-expert problem sizes from expert_first_token_offset // produced by vLLM's moe_permute kernel @@ -235,14 +247,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { ops.def( "mxfp4_experts_quant(Tensor! output, Tensor! output_scale," "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + "Tensor output_scale_offset_by_experts, int n_experts, " + "Tensor? sorted_expert_ids=None) -> ()"); // Fused SiLU+Mul+MXFP4 experts quantization. ops.def( "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! " "output_scale," "Tensor input, Tensor input_offset_by_experts," - "Tensor output_scale_offset_by_experts, int n_experts) -> ()"); + "Tensor output_scale_offset_by_experts, int n_experts, " + "Tensor? sorted_expert_ids=None) -> ()"); // Fused SiLU+Mul+NVFP4 quantization. ops.def( @@ -535,7 +549,17 @@ 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, " "Tensor!? q_fp8_out=None, float q_fp8_scale=1.0, " - "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None) -> ()"); + "Tensor? kv_k_scale=None, Tensor? kv_v_scale=None, " + "bool write_qkv=True) -> ()"); + + // MiniMax-M3 DSpark draft: Gemma k-norm + partial NeoX RoPE + K/V cache + // insert for every draft layer's context K/V in one launch. + ops.def( + "fused_minimax_m3_dspark_context_kv(" + "Tensor all_kv, Tensor k_norm_weight, Tensor cos_sin_cache, " + "Tensor positions, Tensor slot_mapping, Tensor! kv_cache, " + "Tensor cache_offsets, Tensor k_scale, Tensor scale_offsets, " + "int rotary_dim, float eps, str kv_cache_dtype) -> ()"); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE ops.def( @@ -779,6 +803,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { +#ifdef VLLM_ENABLE_FUSED_MQA_TOPK + ops.impl("fused_mqa_topk_prefill", TORCH_BOX(&fused_mqa_topk_prefill)); +#endif + // LongCat n-gram embedding index kernel. ops.impl("ngram_compute_n_gram_ids", TORCH_BOX(&ngram_compute_n_gram_ids)); ops.impl("vocab_parallel_embedding", TORCH_BOX(&vocab_parallel_embedding)); @@ -876,6 +904,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { TORCH_BOX(&fused_muse_spark_qknorm_rope_kv_insert)); ops.impl("fused_minimax_m3_qknorm_rope_kv_insert", TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert)); + ops.impl("fused_minimax_m3_dspark_context_kv", + TORCH_BOX(&fused_minimax_m3_dspark_context_kv)); #ifdef VLLM_ENABLE_FUSED_KDA_DECODE ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); #endif @@ -1021,7 +1051,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor! kv_cache," " Tensor slot_mapping," " str kv_cache_dtype," - " Tensor scale) -> ()"); + " Tensor scale," + " bool skip_scale=False) -> ()"); // Grouped concat_and_cache_mla across all layers. Each layer's cache base // pointer and optional plain-FP8 scale are read from device tensors. @@ -1060,13 +1091,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { " Tensor(h!)? swap_counts=None," " Tensor? resident_block_table=None," " int resident_block_size=0," - " int resident_null_block=0) -> ()"); + " int resident_null_block=0," + " Tensor? num_valid_rows=None) -> ()"); ops.def( - "hisparse_invalidate_written_slots(Tensor! device_global_indices," + "hisparse_invalidate_written_slots(Tensor device_global_indices_ptrs," " Tensor request_state_indices," " Tensor req_id_per_token," - " Tensor written_slots) -> ()"); + " Tensor written_slots," + " int num_state_rows," + " int region_stride) -> ()"); ops.def( "hisparse_gather_plan(Tensor host_cache," @@ -1128,7 +1162,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { ops.def( "cp_gather_and_upconvert_nvfp4_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor workspace_starts, int batch_size) -> ()"); + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " + "seq_starts=None) -> ()"); ops.def( "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " diff --git a/docker/Dockerfile b/docker/Dockerfile index c359a93c378a..ff5b273c56fd 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -31,7 +31,8 @@ ARG TRITON_INSTALL_FROM_SOURCE_REVISION= # DeepEPv2 requires NCCL >= 2.30.4 (GIN backend). # This version is used for CUDA 12+ builds to keep the DeepEP build and # runtime NCCL versions compatible. -ARG NCCL_VERSION=2.30.7 +# DeepEPv2 GDA over EFA requires NCCL 2.31.x or later. +ARG NCCL_VERSION=2.31.2 # 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 @@ -92,11 +93,6 @@ ARG UV_KEYRING_PROVIDER=${PIP_KEYRING_PROVIDER} # Flag enables built-in KV-connector dependency libs into docker images ARG INSTALL_KV_CONNECTORS=false -# Flag builds NIXL, including its NIXL EP extension for the installed PyTorch, -# from source and installs it in place of the nixl packages from the KV-connector -# requirements -ARG BUILD_NIXL=false - #################### BASE BUILD IMAGE #################### # prepare basic build environment FROM ${BUILD_BASE_IMAGE} AS base @@ -252,8 +248,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \ && uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ + # Keep the CUDA-specific PyTorch index first. Using it as an extra + # index with unsafe-best-match can select a higher, incompatible + # torch build from the base image (for example cu130 in a cu129 image). uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi # Track PyTorch lib versions used during build and match in downstream instances. @@ -371,7 +372,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ echo "Installing build requirements..." \ && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi WORKDIR /workspace @@ -560,24 +563,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \ (cd "${TRITON_WORKSPACE}/dist" \ && sha256sum build-state > wheel.sha256); \ fi - -# Build NIXL and its NIXL EP extension against the PyTorch installed here when -# requested. The wheels are installed by the final image, which has no compiler. -ARG BUILD_NIXL -COPY tools/build_nixl_with_torch215_for_rubin.py /tmp/build_nixl_for_rubin.py -RUN --mount=type=bind,source=requirements,target=/tmp/vllm-requirements,ro \ - set -eux; \ - NIXL_WORKSPACE=/tmp/nixl_workspace; \ - mkdir -p "${NIXL_WORKSPACE}/dist"; \ - if [ "${BUILD_NIXL}" = "true" ]; then \ - dnf install -y --setopt=install_weak_deps=False autoconf automake libtool; \ - dnf clean all; \ - rm -rf /var/cache/dnf; \ - NIXL_REQUIREMENTS_FILE=/tmp/vllm-requirements/kv_connectors.txt \ - WHEELS_CACHE_HOME="${NIXL_WORKSPACE}/dist" \ - MAX_JOBS="${max_jobs:-}" \ - python3 /tmp/build_nixl_for_rubin.py --skip-system-deps --build-only; \ - fi #################### EXTENSIONS BUILD IMAGE #################### #################### WHEEL BUILD IMAGE #################### @@ -731,7 +716,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ echo "Installing dev requirements..." \ && uv pip install --python /opt/venv/bin/python3 -r requirements/dev.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi #################### DEV IMAGE #################### @@ -940,7 +927,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ uv pip install --system \ -r /tmp/requirements-cuda.txt \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi && \ rm /tmp/requirements-cuda.txt /tmp/common.txt @@ -975,7 +964,7 @@ RUN set -eux; \ # From versions.json: .openai_server_extras.timm, .openai_server_extras.runai_model_streamer, # .openai_server_extras.gguf (NVFP4 GGUF checkpoints, e.g. RecursiveSI/flint-nvfp4-gguf) ARG TIMM_VERSION=">=1.0.17" -ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7" +ARG RUNAI_MODEL_STREAMER_VERSION=">=0.17.1" ARG GGUF_VERSION=">=0.19.0" RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system accelerate 'modelscope<1.38' \ @@ -1016,7 +1005,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist else \ echo "Installing vLLM..." \ && uv pip install --system dist/*.whl --verbose \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi # Expose the Rust frontend bundled in the wheel on the container's PATH. @@ -1139,7 +1130,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist else \ echo "Installing vLLM..." \ && uv pip install --system dist/*.whl --verbose \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match; \ fi # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH @@ -1149,7 +1142,9 @@ COPY --from=build /tmp/ep_kernels_workspace/dist/wheels.sha256 /tmp/ep-kernels-w RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \ --mount=type=cache,target=/opt/uv/cache \ uv pip install --system ep_kernels/dist/*.whl --verbose \ - --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') + --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ + --extra-index-url https://pypi.org/simple \ + --index-strategy unsafe-first-match # See the matching comment in the vllm-base stage. ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:${LD_LIBRARY_PATH} @@ -1292,25 +1287,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \ /tmp/install_triton_from_source.sh \ /tmp/triton-source-wheel.sha256 -# Install source-built NIXL last, replacing the nixl packages installed above, so -# the image has a NIXL EP build for its PyTorch version. There is no GPU driver -# at build time, so check for the extension instead of importing it. -ARG BUILD_NIXL -RUN --mount=type=cache,target=/opt/uv/cache \ - --mount=type=bind,from=extensions-build,source=/tmp/nixl_workspace/dist,target=/tmp/nixl-dist,ro \ - set -eu; \ - if [ "${BUILD_NIXL}" = "true" ]; then \ - ls /tmp/nixl-dist/*.whl; \ - uv pip uninstall --system nixl nixl-cu12 nixl-cu13; \ - uv pip install --system --no-deps /tmp/nixl-dist/*.whl; \ - python3 -c "import glob, importlib.util as u, os, torch; \ - spec = u.find_spec('nixl_ep_cu' + os.environ['CUDA_VERSION'].split('.')[0]); \ - mm = ''.join(torch.__version__.split('.')[:2]); \ - ext = glob.glob(os.path.join(spec.submodule_search_locations[0], 'nixl_ep_cpp_torch' + mm + '.*')) if spec else []; \ - assert ext, 'NIXL EP has no build for torch ' + torch.__version__; \ - print('NIXL EP extension:', ext[0])"; \ - fi - ENV VLLM_USAGE_SOURCE production-docker-image ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \ VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index db146475c830..b8fb7e07ccba 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -446,6 +446,10 @@ ADD ./tools/ ./tools/ # ship either file, and this stage does not inherit the vllm-src tree. ADD ./setup.py ./setup.py ADD ./requirements/cpu.txt ./requirements/cpu.txt +# tests/v1/core/test_prompt_block_hash_vectors.py checks the engine against the +# prompt block hash vectors that the Rust frontend's unit tests also use. +ADD ./rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json \ + ./rust/src/engine-core-client/src/tests/prompt_block_hash_vectors.json # install development dependencies (for testing) RUN --mount=type=cache,target=/root/.cache/uv \ diff --git a/docker/Dockerfile.efa b/docker/Dockerfile.efa new file mode 100644 index 000000000000..00c22248c132 --- /dev/null +++ b/docker/Dockerfile.efa @@ -0,0 +1,156 @@ +# AWS EFA variant of a vllm-openai image. Adds the EFA userspace stack +# (libfabric, aws-ofi-nccl, hwloc), the EFA Mooncake wheel (x86_64 only), +# and DeepEP rebuilt from the amazon-contributing fork for NCCL GIN over EFA. +# +# The base must be built with INSTALL_KV_CONNECTORS=true (NIXL and Mooncake): +# docker build -f docker/Dockerfile.efa \ +# --build-arg BASE_IMAGE=vllm/vllm-openai:latest -t vllm:openai-efa docker/ +# +# The host still provides the EFA kernel module and memlock limits. +ARG BASE_IMAGE=vllm/vllm-openai:latest +FROM ${BASE_IMAGE} + +ARG TARGETARCH +ARG EFA_INSTALLER_VERSION=1.51.0 +ARG EFA_INSTALLER_SHA256=83993044fa4df3dc2f8b5c2f7051d61b9689441be023a2c92ea51d3f53f2465c + +# Without the KV connectors there is no NIXL plugin or Mooncake wheel to use EFA. +RUN set -eu; \ + if ! python3 -c "import importlib.metadata as m; m.version('nixl')" \ + 2>/dev/null; then \ + echo "Dockerfile.efa needs a base image built with \ +--build-arg INSTALL_KV_CONNECTORS=true" >&2; \ + exit 1; \ + fi + +# Install the AWS EFA userspace stack. +RUN set -eux; \ + EFA_INSTALLER_ARCHIVE="/tmp/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz"; \ + curl -fsSL --retry 3 --retry-all-errors \ + "https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz" \ + -o "${EFA_INSTALLER_ARCHIVE}"; \ + echo "${EFA_INSTALLER_SHA256} ${EFA_INSTALLER_ARCHIVE}" | sha256sum -c -; \ + tar -xzf "${EFA_INSTALLER_ARCHIVE}" -C /tmp; \ + # The installer apt-installs dependencies; the base image purged the lists. + apt-get update -y; \ + # Minimal install: libfabric, aws-ofi-nccl and hwloc only. --build-ngc + # implies --skip-kmod --skip-limit-conf --no-verify. + (cd /tmp/aws-efa-installer && \ + ./efa_installer.sh --skip-mpi --skip-rdma-core --build-ngc); \ + test -x /opt/amazon/efa/bin/fi_info; \ + /opt/amazon/efa/bin/fi_info --version; \ + # `-l` lists providers without EFA hardware, so this works in a build. + /opt/amazon/efa/bin/fi_info -l | grep -q '^efa:'; \ + EFA_NIXL_PLUGIN=$(find /usr/local/lib -name libplugin_LIBFABRIC.so \ + -print -quit); \ + test -n "${EFA_NIXL_PLUGIN}"; \ + ldd "${EFA_NIXL_PLUGIN}" | grep -q 'libfabric\.so\.1 => /opt/amazon/efa/lib/'; \ + # Plugins that NCCL_NET_PLUGIN=ofi and NCCL_TUNER_PLUGIN=ofi resolve to. + ldconfig -p | grep -q '/opt/amazon/ofi-nccl/lib/libnccl-net-ofi\.so$'; \ + ldconfig -p | grep -q '/opt/amazon/ofi-nccl/lib/libnccl-tuner-ofi\.so$'; \ + apt-get clean; \ + rm -rf "${EFA_INSTALLER_ARCHIVE}" /tmp/aws-efa-installer \ + /var/lib/apt/lists/* + +# Route NCCL over EFA. Override with `-e NCCL_NET_PLUGIN= -e NCCL_TUNER_PLUGIN=` +# to run this image on a non-EFA host. +ENV NCCL_NET_PLUGIN=ofi \ + NCCL_TUNER_PLUGIN=ofi + +# Swap in the EFA Mooncake wheel at the base image's Mooncake version. Both +# wheels own the `mooncake` module; the EFA wheels are published for x86_64 only. +RUN --mount=type=cache,target=/opt/uv/cache \ + set -eu; \ + pkg_version() { python3 -c "import importlib.metadata as m, sys; \ +print(m.version(sys.argv[1]))" "$1" 2>/dev/null; }; \ + if MOONCAKE_VERSION=$(pkg_version mooncake-transfer-engine-cuda13); then \ + MOONCAKE_PKG=mooncake-transfer-engine-cuda13; \ + MOONCAKE_EFA_PKG=mooncake-transfer-engine-efa-cuda13; \ + elif MOONCAKE_VERSION=$(pkg_version mooncake-transfer-engine); then \ + MOONCAKE_PKG=mooncake-transfer-engine; \ + MOONCAKE_EFA_PKG=mooncake-transfer-engine-efa; \ + else \ + echo "No Mooncake wheel in the base image; build it with \ +--build-arg INSTALL_KV_CONNECTORS=true" >&2; \ + exit 1; \ + fi; \ + if [ "${TARGETARCH}" != "amd64" ]; then \ + echo "TARGETARCH=${TARGETARCH}: ${MOONCAKE_EFA_PKG} is x86_64-only, \ +keeping ${MOONCAKE_PKG} (Mooncake EFA protocol unavailable)"; \ + else \ + echo "Swapping ${MOONCAKE_PKG}==${MOONCAKE_VERSION} for \ +${MOONCAKE_EFA_PKG}==${MOONCAKE_VERSION}"; \ + uv pip uninstall --system "${MOONCAKE_PKG}"; \ + uv pip install --system --no-deps \ + "${MOONCAKE_EFA_PKG}==${MOONCAKE_VERSION}"; \ + test "$(pkg_version "${MOONCAKE_EFA_PKG}")" = "${MOONCAKE_VERSION}"; \ + # `! cmd` is exempt from `set -e`, so check explicitly. + if pkg_version "${MOONCAKE_PKG}" >/dev/null; then \ + echo "${MOONCAKE_PKG} is still installed" >&2; \ + exit 1; \ + fi; \ + fi + +# Replace the base image's upstream DeepEP with the amazon-contributing fork, +# whose deepep_v2 backend carries EP traffic over EFA through NCCL GIN. +ARG DEEPEP_REPO="https://github.com/amazon-contributing/DeepEP.git" +ARG DEEPEP_REF="874779c9ccd2294b56304bd6cc5f138f1f71d097" +ARG DEEPEP_CUDA_ARCH_LIST="9.0a;10.0a;10.3a" + +# Headers and git for DeepEP's nvcc build. `reinstall` also installs packages +# that are absent. +RUN set -eux; \ + CUDA_PKG_VERSION=$(nvcc --version \ + | grep -oP 'release \K[0-9]+\.[0-9]+' | tr '.' '-'); \ + DEBIAN_FRONTEND=noninteractive apt-get update -y; \ + DEBIAN_FRONTEND=noninteractive apt-get reinstall -y --no-install-recommends \ + git \ + libcusparse-dev-${CUDA_PKG_VERSION} \ + libcusolver-dev-${CUDA_PKG_VERSION} \ + cuda-nvrtc-dev-${CUDA_PKG_VERSION}; \ + rm -rf /var/lib/apt/lists/* + +RUN --mount=type=cache,target=/opt/uv/cache \ + set -eux; \ + git clone "${DEEPEP_REPO}" /tmp/deepep; \ + cd /tmp/deepep; \ + git checkout "${DEEPEP_REF}"; \ + git submodule update --init --recursive; \ + uv pip uninstall --system deep_ep 2>/dev/null || true; \ + NCCL_ROOT="$(python3 -c "import importlib.util, os; \ +s = importlib.util.find_spec('nvidia.nccl'); \ +print(os.path.dirname(s.origin or s.submodule_search_locations[0]))")/nccl"; \ + export EP_NCCL_ROOT_DIR="${NCCL_ROOT}"; \ + export LD_LIBRARY_PATH="${NCCL_ROOT}/lib:${LD_LIBRARY_PATH:-}"; \ + # The CUDA driver stubs let DeepEP link libcuda in a builder with no driver. + export LIBRARY_PATH="${NCCL_ROOT}/lib:/usr/local/cuda/lib64/stubs:\ +/usr/local/cuda/targets/x86_64-linux/lib/stubs:${LIBRARY_PATH:-}"; \ + export DISABLE_AGGRESSIVE_PTX_INSTRS=1; \ + export TORCH_CUDA_ARCH_LIST="${DEEPEP_CUDA_ARCH_LIST}"; \ + uv pip install --system --no-deps -vv --no-build-isolation .; \ + cd /; \ + rm -rf /tmp/deepep; \ + # vLLM only logs when the deepep_v2 backend is unavailable, so fail here. + python3 -c "import deep_ep; assert hasattr(deep_ep, 'ElasticBuffer'), \ +'DeepEP built without the ElasticBuffer (v2) API'" + +# Provenance. Build args from docker-build-metadata-args.sh override the values +# inherited from the base image; without them the base's values are kept. +ARG BASE_IMAGE +ARG VLLM_BUILD_COMMIT +ARG VLLM_BUILD_PIPELINE +ARG VLLM_BUILD_URL +ARG VLLM_IMAGE_TAG +ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT} \ + VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE} \ + VLLM_BUILD_URL=${VLLM_BUILD_URL} \ + VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG} +LABEL org.opencontainers.image.revision="${VLLM_BUILD_COMMIT}" \ + org.opencontainers.image.version="${VLLM_IMAGE_TAG}" \ + org.opencontainers.image.url="${VLLM_BUILD_URL}" \ + ai.vllm.build.commit="${VLLM_BUILD_COMMIT}" \ + ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}" \ + ai.vllm.build.url="${VLLM_BUILD_URL}" \ + ai.vllm.image.tag="${VLLM_IMAGE_TAG}" \ + ai.vllm.efa.base-image="${BASE_IMAGE}" \ + ai.vllm.efa.installer-version="${EFA_INSTALLER_VERSION}" diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0d3686f58c11..30de51609dd0 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,390 +1,40 @@ -# default base image -ARG REMOTE_VLLM="0" -ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" -ARG VLLM_BRANCH="main" -ARG COMMON_WORKDIR=/app -ARG BASE_IMAGE=rocm/vllm-dev:base -ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base -ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557 -ARG INSTALL_LMCACHE=false -# NIC backend for MoRI RDMA support. -# By default (all), drivers and userspace libraries for all supported NIC types -# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. -# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: -# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only -# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only -# --build-arg NIC_BACKEND=none # Install nothing. -ARG NIC_BACKEND=all -# AMD AINIC apt repo settings -# Users can specify a custom version compatible with their host drivers. -# The default version has been tested with ionic-dkms=26.03.3.001 -# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. -ARG AINIC_VERSION=1.117.5-a-77 -ARG UBUNTU_CODENAME=jammy - -# Sccache configuration. Release builds use this today; CI can opt in when a -# shared S3-compatible cache backend is available. -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 - -FROM ${BASE_IMAGE} AS base - -ARG ARG_PYTORCH_ROCM_ARCH -ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} - -# Install build dependencies and utilities -RUN apt-get update -q -y && apt-get install -q -y \ +# syntax=docker/dockerfile:1.27@sha256:4edf897a3ffa55b89f906fc8cc78afdb3f1834cc9c7083565e611a8a7d5fe99e +# vLLM on ROCm. Build it with docker/docker-bake-rocm.hcl, which provides the +# `stack` named context (the `final` stage of docker/Dockerfile.rocm_base) +# and optionally prebuilt `runtime` and `test-deps` images. +# Everything above the per-commit marker is hashed into the runtime key by .buildkite/scripts/rocm/build.sh. + +FROM stack AS dev +# Build tools plus the runtime system packages (RDMA userspace, FFmpeg for torchcodec). +# mold is installed but not the default linker: some packages JIT-compile at runtime +# with flags mold does not support. Build stages opt in via LDFLAGS="-fuse-ld=mold". +RUN apt-get update -q -y -o Acquire::Retries=3 && apt-get install -q -y -o Acquire::Retries=3 \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ - apt-transport-https ca-certificates wget curl \ - libnuma-dev ccache mold + apt-transport-https ca-certificates wget curl libnuma-dev ccache mold \ + librdmacm1 libibverbs1 ibverbs-providers ibverbs-utils util-linux unzip libsodium23 \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ + && rm -rf /var/lib/apt/lists/* RUN --mount=type=cache,target=/root/.cache/pip \ python3 -m pip install --upgrade pip -# Note: mold is installed but not set as the system default linker because -# some packages use JIT compilation at runtime with flags mold does not support. -# Build stages opt in via LDFLAGS="-fuse-ld=mold". -# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) -ARG USE_SCCACHE -RUN if [ "$USE_SCCACHE" != "1" ]; then \ - apt-get purge -y sccache || true; \ - python3 -m pip uninstall -y sccache || true; \ - rm -f "$(which sccache)" || true; \ - fi - -# Install UV — download first, then run, so a curl failure is not masked by the pipe +# Download first, then run, so a curl failure is not masked by the pipe RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ && rm -f /tmp/uv-install.sh \ && uv --version - -# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out -# Reference: https://github.com/astral-sh/uv/pull/1694 +# Some dependencies time out with uv's default; see https://github.com/astral-sh/uv/pull/1694 ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" # Use copy mode to avoid hardlink failures with Docker cache mounts ENV UV_LINK_MODE=copy -# ccache directory - persisted across layer rebuilds via cache mounts. ENV CCACHE_DIR=/root/.cache/ccache ENV CCACHE_COMPILERCHECK=content -# Empty by default so build steps fall back to $(nproc); CI can override. -ARG max_jobs -ENV MAX_JOBS=${max_jobs} - -# Install sccache if USE_SCCACHE is enabled (for release builds) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME -ARG SCCACHE_REGION_NAME -ARG SCCACHE_S3_NO_CREDENTIALS -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - if command -v sccache >/dev/null 2>&1; then \ - echo "sccache already installed, skipping installation"; \ - sccache --version; \ - else \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi; \ - fi - -# Set sccache environment variables only when USE_SCCACHE=1 -# This prevents S3 config from leaking into images when sccache is not used -ARG USE_SCCACHE -ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} -ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} -ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} -ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} - -ARG COMMON_WORKDIR -WORKDIR ${COMMON_WORKDIR} - - -# ----------------------- -# vLLM fetch stages -FROM base AS fetch_vllm_0 -ONBUILD COPY ./ vllm/ -FROM base AS fetch_vllm_1 -ARG VLLM_REPO -ARG VLLM_BRANCH -ONBUILD RUN git clone ${VLLM_REPO} \ - && cd vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git checkout FETCH_HEAD \ - && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ - git remote add upstream "https://github.com/vllm-project/vllm.git" \ - && git fetch upstream ; fi -FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm - -# Resolve the device-independent source version once for every downstream -# artifact. setuptools-scm reads Git metadata for local/remote sources and the -# archive metadata generated for canonical CI contexts. -FROM fetch_vllm AS vllm-version -ARG COMMON_WORKDIR RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - cd ${COMMON_WORKDIR}/vllm \ - && uv pip install --system -r requirements/build/rust.txt \ - && python3 -c \ - "from setuptools_scm import get_version; print(get_version(git_describe_command=['git', 'describe', '--dirty', '--tags', '--long', '--abbrev=40', '--match', 'v[0-9]*']))" \ - > /vllm-version.txt - -# ----------------------- -# Rust toolchain inputs -# -# Keep these separate from Rust source inputs so Rust source-only changes do not -# invalidate the long-lived ci_base image that carries the toolchain. -FROM scratch AS rust_toolchain_input_0 -COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml - -FROM base AS rust_toolchain_input_1 -ARG VLLM_REPO -ARG VLLM_BRANCH -RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-input/vllm \ - && cd /rust-toolchain-input/vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git sparse-checkout init --no-cone \ - && git sparse-checkout set \ - rust-toolchain.toml \ - && git checkout FETCH_HEAD - -FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input - -# ----------------------- -# Rust build inputs -# -# REMOTE_VLLM=0: use only the local files that can affect Rust artifacts. -# REMOTE_VLLM=1: fetch only those same files from the requested remote ref. -FROM scratch AS rust_input_0 -COPY requirements/build/rust.txt /rust-input/vllm/requirements/build/rust.txt -COPY rust/Cargo.lock rust/Cargo.toml /rust-input/vllm/rust/ -COPY rust/proto /rust-input/vllm/rust/proto -COPY rust/src /rust-input/vllm/rust/src -COPY rust-toolchain.toml /rust-input/vllm/rust-toolchain.toml -COPY tools/build_rust.py /rust-input/vllm/tools/build_rust.py -COPY tools/build_rust.sh /rust-input/vllm/tools/build_rust.sh - -FROM base AS rust_input_1 -ARG VLLM_REPO -ARG VLLM_BRANCH -RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \ - && cd /rust-input/vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git sparse-checkout init --no-cone \ - && git sparse-checkout set \ - requirements/build/rust.txt \ - rust/Cargo.lock \ - rust/Cargo.toml \ - rust/proto/** \ - rust/src/** \ - rust-toolchain.toml \ - tools/build_rust.py \ - tools/build_rust.sh \ - && git checkout FETCH_HEAD - -FROM rust_input_${REMOTE_VLLM} AS rust-input - -# ----------------------- -# Rust toolchain -FROM base AS rust-toolchain - -ENV CARGO_HOME=/root/.cargo -ENV RUSTUP_HOME=/root/.rustup -ENV PATH=${CARGO_HOME}/bin:${PATH} - -RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - ca-certificates curl \ - && rm -rf /var/lib/apt/lists/* - -COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml -RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \ - && if ! command -v rustup >/dev/null 2>&1; then \ - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ - | sh -s -- -y --default-toolchain none; \ - fi \ - && rustup toolchain install "${TOOLCHAIN}" \ - && rustup default "${TOOLCHAIN}" \ - && rustc --version \ - && cargo --version - -# ----------------------- -# Rust build stage -# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain. -FROM rust-toolchain AS rust-build -ARG COMMON_WORKDIR -ARG USE_SCCACHE - -# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit -# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). -ENV CARGO_BUILD_JOBS=4 -ENV CARGO_NET_RETRY=10 -ENV RUSTUP_MAX_RETRIES=10 - -WORKDIR ${COMMON_WORKDIR}/vllm - -COPY --from=rust-input /rust-input/vllm/requirements/build/rust.txt requirements/build/rust.txt - -# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on -# a separate local sccache daemon while sharing the same remote cache backend. -ENV SCCACHE_SERVER_PORT=4227 - -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -r requirements/build/rust.txt - -COPY --from=rust-input /rust-input/vllm/rust rust -COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml -COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py -COPY --from=rust-input /rust-input/vllm/tools/build_rust.sh tools/build_rust.sh -COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt - -# Build the release binary. Cargo's registry/git caches can be written by -# concurrent BuildKit jobs on shared workers, so lock those cache mounts while -# keeping the cache benefit. Do not cache target/, because stale target metadata -# can outlive source updates across BuildKit cache reuse. -RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ - --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ - if [ "$USE_SCCACHE" = "1" ]; then \ - export RUSTC_WRAPPER=sccache \ - && sccache --show-stats; \ - fi \ - && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ - && bash tools/build_rust.sh \ - && test -x vllm/vllm-rs \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - sccache --show-stats; \ - fi - -# ----------------------- -# Shared vLLM build dependencies -# -# Install requirements once, before bringing source-specific inputs into either -# consumer. Native compilation and wheel packaging then share this expensive -# layer while retaining their independent source cache boundaries. -FROM base AS build_vllm_dependencies -ARG COMMON_WORKDIR -ENV VLLM_TARGET_DEVICE=rocm -WORKDIR ${COMMON_WORKDIR}/vllm - -COPY --from=fetch_vllm \ - ${COMMON_WORKDIR}/vllm/requirements/common.txt \ - ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \ - requirements/ - -# Use constraints file to protect triton, pybind11, and torch dependencies. -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -c /etc/rocm-constraints.txt -r requirements/rocm.txt - -# Fetch only the pinned ROCm Triton kernels package. Keeping this before the -# vLLM source COPYs makes the download reusable when native sources change. -FROM base AS rocm-triton-kernels -ARG ROCM_TRITON_KERNELS_COMMIT -ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \ - /tmp/rocm-triton.tar.gz -RUN mkdir -p /opt/rocm-triton-kernels \ - && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \ - --directory /opt/rocm-triton-kernels --strip-components=4 \ - --no-same-owner \ - "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \ - && test -f /opt/rocm-triton-kernels/__init__.py \ - && rm /tmp/rocm-triton.tar.gz - -# ----------------------- -# vLLM native build stages -# -# csrc-build intentionally copies only files that affect ROCm native extension -# compilation. That keeps unrelated CI/test/docs edits from invalidating the -# expensive HIP/C++ build layer. -FROM build_vllm_dependencies AS csrc-build -ARG ROCM_TRITON_KERNELS_COMMIT -COPY --from=rocm-triton-kernels \ - /opt/rocm-triton-kernels /opt/rocm-triton-kernels -ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels - -# pyproject.toml is bind-mounted so it is not persisted in the csrc-build image, -# but it still participates in the BuildKit cache key for this RUN step. -COPY setup.py CMakeLists.txt ./ -COPY tools/build_rust.py tools/build_rust.py -COPY cmake cmake/ -COPY csrc csrc/ -COPY vllm/envs.py vllm/envs.py -COPY vllm/__init__.py vllm/__init__.py - -ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" - -RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ - cmake/external_projects/triton_kernels.cmake \ - || echo "WARNING: Dockerfile.rocm ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ - && export CCACHE_BASEDIR="$PWD" \ - && echo "=== ccache stats before ROCm native build ===" \ - && (ccache --show-stats || true) \ - && (ccache --zero-stats || true) \ - && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ - && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ - && PYTHONDONTWRITEBYTECODE=1 LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ - && test -d dist \ - && ls dist/*.whl >/dev/null \ - && echo "=== ccache stats after ROCm native build ===" \ - && (ccache --show-stats || true) - -# Build the full vLLM ROCm wheel by reusing the native extension wheel from -# csrc-build. This stage still rebuilds for Python/package changes, but skips -# the expensive dependency installation and HIP/C++ compilation when their -# respective inputs are unchanged. -FROM build_vllm_dependencies AS build_vllm -COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm -COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels -COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt - -# Drop the pre-built Rust artifacts into the source tree. setup.py detects -# them and ships them as-is, skipping the local Rust build. -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ - -RUN cd "${COMMON_WORKDIR}/vllm" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ - && export VLLM_USE_PRECOMPILED=1 \ - && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ - && export VLLM_DOCKER_BUILD_CONTEXT=1 \ - && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ - && PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=dist \ - && test -d dist \ - && ls dist/*.whl >/dev/null -FROM scratch AS export_vllm -ARG COMMON_WORKDIR -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 + uv pip install --system -c /etc/rocm-constraints.txt cmake ninja "setuptools<80" setuptools-scm \ + setuptools-rust wheel jinja2 packaging regex # NIXL/UCX build stages -FROM base AS build_nixl +FROM dev AS build_nixl ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96" ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git" ARG UCX_BRANCH="2ff0b47ce7626584cb9b56c647d08f712c6c3edf" @@ -409,11 +59,11 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ ibverbs-providers \ && rm -rf /var/lib/apt/lists/* -RUN --mount=type=cache,target=/root/.cache/uv \ +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ uv pip install --system meson meson-python pyyaml types-PyYAML \ auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0" -RUN --mount=type=cache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd /usr/local/src && \ git clone ${UCX_REPO} && \ cd ucx && \ @@ -438,7 +88,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ ENV PATH=/usr/local/ucx/bin:$PATH ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} -RUN --mount=type=cache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ git clone ${NIXL_REPO} /opt/nixl && \ cd /opt/nixl && \ git checkout ${NIXL_BRANCH} && \ @@ -487,12 +137,12 @@ RUN cd /opt/nixl && \ cp /tmp/nixl_wheels/repaired/*.whl /app/install # Fetch and verify pinned LMCache source. -FROM base AS lmcache_source +FROM dev AS lmcache_source ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git" ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b" ARG LMCACHE_VERSION="0.5.3" -# Remove NVIDIA-only dependencies and accept the release image's NumPy version. +# Remove NVIDIA-only dependencies and accept the runtime's NumPy version. RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \ && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \ "${LMCACHE_REPO}" /app/lmcache \ @@ -507,7 +157,7 @@ RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \ /app/lmcache/requirements/common.txt \ && rm -rf /app/lmcache/.git -FROM base AS build_lmcache +FROM dev AS build_lmcache COPY --from=lmcache_source \ /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ @@ -515,7 +165,6 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ ARG LMCACHE_VERSION="0.5.3" ARG LMCACHE_ROCM_ARCH="gfx942;gfx950" -ARG USE_SCCACHE COPY --from=lmcache_source /app/lmcache /app/lmcache WORKDIR /app/lmcache @@ -526,9 +175,6 @@ RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \ offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \ | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \ - if [ "${USE_SCCACHE:-0}" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers; \ - fi; \ mkdir -p /app/install; \ SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \ PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \ @@ -551,10 +197,8 @@ RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ exit 1; \ fi -# DeepEP build stage. ROCShmem is provided by TheRock SDK -# ROCSHMEM_DIR set to the SDK install path so DeepEP finds the headers, etc. -FROM base AS build_deepep -ENV ROCSHMEM_DIR=${ROCM_PATH} +# DeepEP build stage. The stack sets ROCSHMEM_DIR via the TheRock SDK. +FROM dev AS build_deepep ARG DEEPEP_BRANCH="0f63d3e9" ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" ARG DEEPEP_NIC="cx7" @@ -564,20 +208,53 @@ ARG DEEPEP_NIC="cx7" # We pass --rocm-gfx942-fp8-fnuz-max 224 to align the DeepEP dispatch FP8 # quantization bound with vLLM's own clamp # (vllm/model_executor/layers/quantization/utils/fp8_utils.py): 224 for E4M3FNUZ -RUN --mount=type=cache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ && git clone ${DEEPEP_REPO} \ && cd DeepEP \ && git checkout ${DEEPEP_BRANCH} \ && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install -# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do -# not force users to rebuild the long-lived Dockerfile.rocm_base image. -FROM base AS mori_base -ARG NIC_BACKEND -ARG AINIC_VERSION -ARG UBUNTU_CODENAME -RUN /bin/bash -lc 'set -euo pipefail; \ +# mamba-ssm and causal-conv1d, the fast path transformers uses for the hybrid-model references in the +# language model tests. CI only. Their setup.py targets HIP_ARCHITECTURES (default: the local GPU). +FROM dev AS build_mamba +ARG MAMBA_COMMIT="269cbe654d10883374736eb936e61d3ee5625695" +ARG CAUSAL_CONV1D_COMMIT="da6dbaa9fd5a919967f14d3fd031da1288ad5025" +ARG HIP_ARCHITECTURES="gfx942,gfx950" +RUN MAMBA_FORCE_BUILD=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE \ + python3 -m pip wheel --no-build-isolation --no-deps -w /app/mamba_install \ + "git+https://github.com/AndreasKaratzas/mamba@${MAMBA_COMMIT}" \ + "git+https://github.com/Dao-AILab/causal-conv1d@${CAUSAL_CONV1D_COMMIT}" \ + && for w in /app/mamba_install/*.whl; do python3 -m zipfile -e "$w" /tmp/whl; done \ + && for so in /tmp/whl/*.so; do for a in $(echo "$HIP_ARCHITECTURES" | tr , ' '); do \ + llvm-objdump --offloading "$so" | grep -qw "$a" || { echo "$so lacks $a" >&2; exit 1; }; done; done \ + && rm -rf /tmp/whl + +# rdma-core recent enough for DeepEP's rocshmem. CI only: it replaces the distro libibverbs, which +# the vendor NIC providers in the runtime (AINIC, Broadcom bnxt_re) are built against. +FROM dev AS build_rdma_core +RUN git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core \ + && cd /tmp/rdma-core && mkdir build && cd build \ + && cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. \ + && ninja && DESTDIR=/out ninja install + +# The runtime: stack + every compiled component + all pip dependencies + all env. Published as base-nightly. +FROM dev AS runtime +# NIC backend for MoRI RDMA support. +# By default (all), drivers and userspace libraries for all supported NIC types +# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. +# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: +# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only +# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only +# --build-arg NIC_BACKEND=none # Install nothing. +ARG NIC_BACKEND=all +# AMD AINIC apt repo settings +# Users can specify a custom version compatible with their host drivers. +# The default version has been tested with ionic-dkms=26.03.3.001 +# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. +ARG AINIC_VERSION=1.117.5-a-77 +ARG UBUNTU_CODENAME=jammy +RUN --mount=type=cache,target=/root/.cache/pip /bin/bash -lc 'set -euo pipefail; \ \ install_ainic() { \ apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ @@ -610,9 +287,7 @@ RUN /bin/bash -lc 'set -euo pipefail; \ rm -rf /var/lib/apt/lists/*; \ }; \ \ - echo "[MORI] Install MoRI proxy deps"; \ - pip install --quiet --ignore-installed blinker && \ - pip install --quiet quart msgpack aiohttp pyzmq; \ + pip install --quiet --ignore-installed blinker; \ echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ \ # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ @@ -624,395 +299,297 @@ RUN /bin/bash -lc 'set -euo pipefail; \ bnxt) install_bnxt ;; \ *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ esac' - -# ----------------------- -# vLLM wheel release build stage (for building distributable wheels) -# This stage pins dependencies to custom ROCm wheel versions and handles version detection -FROM fetch_vllm AS build_vllm_wheel_release - -ARG COMMON_WORKDIR -ARG ROCM_TRITON_KERNELS_COMMIT -COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt -COPY --from=rocm-triton-kernels \ - /opt/rocm-triton-kernels /opt/rocm-triton-kernels -ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels - -# Drop the pre-built Rust artifacts into the source tree. setup.py detects -# them and ships them as-is, skipping the local Rust build. -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ - -# Create /install directory for custom wheels -RUN mkdir -p /install - -# Copy custom ROCm wheels from docker/context if they exist -# COPY ensures Docker cache is invalidated when wheels change -# .keep file ensures directory always exists for COPY to work -COPY docker/context/base-wheels/ /tmp/base-wheels/ -# This is how we know if we are building for a wheel release or not. -# If there are not wheels found there, we are not building for a wheel release. -# So we exit with an error. To skip this stage. -RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ - echo "Found custom wheels - copying to /install"; \ - cp /tmp/base-wheels/*.whl /install/ && \ - echo "Copied custom wheels:"; \ - ls -lh /install/; \ - else \ - echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ - echo "Wheel releases require pre-built ROCm wheels."; \ - exit 1; \ - fi - -# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) -# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm -ARG GIT_REPO_CHECK=0 -RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ - echo "Running repository checks..."; \ - cd vllm && bash tools/check_repo.sh; \ - fi - -# Install the packaging helpers before modifying requirements/rocm.txt. -RUN --mount=type=cache,target=/root/.cache/uv \ - { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ - vllm/cmake/external_projects/triton_kernels.cmake \ - || echo "WARNING: Dockerfile.rocm ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ - && cd vllm \ - && uv pip install --system setuptools_scm regex - -# Fail if git-based package dependencies are found in requirements files -# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) -# Extra notes: pip install is able to handle git+ URLs, but uv doesn't. -RUN echo "Checking for git-based packages in requirements files..." \ - && echo "Checking common.txt for git-based packages:" \ - && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ - echo "ERROR: Git-based packages found in common.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in common.txt"; \ - fi \ - && echo "Checking rocm.txt for git-based packages:" \ - && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ - echo "ERROR: Git-based packages found in rocm.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in rocm.txt"; \ - fi \ - && echo "All requirements files are clean - no git-based packages found" - -# Pin vLLM dependencies to exact versions of custom ROCm wheels -# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi -COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py -RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ - && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt - -# Install dependencies using custom wheels from /install -RUN --mount=type=cache,target=/root/.cache/uv \ - cd vllm \ - && echo "Building vLLM with custom wheels from /install" \ - && uv pip install --system --find-links /install -r requirements/rocm.txt - -# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt -# (setup.py auto-detects ccache/sccache in PATH) -RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - cd vllm \ - && export CCACHE_BASEDIR="$PWD" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ - && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ - && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist - -FROM scratch AS export_vllm_wheel_release -ARG COMMON_WORKDIR -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl / -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/benchmark_helion_kernels.py /tools/benchmark_helion_kernels.py -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_base /docker/Dockerfile.rocm_base -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 - -# ----------------------- -# CI base image (Tier 1) - stable, rarely changing CI dependencies. -# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers -# in the vLLM artifacts for the current commit. -FROM mori_base AS ci_base -ARG COMMON_WORKDIR - -ENV CARGO_HOME=/root/.cargo -ENV RUSTUP_HOME=/root/.rustup -ENV PATH=${CARGO_HOME}/bin:${PATH} - -# Bake the Rust toolchain into ci_base so Rust build/test steps do not -# need to fetch it on each per-commit image build. -COPY --from=rust-toolchain /root/.cargo /root/.cargo -COPY --from=rust-toolchain /root/.rustup /root/.rustup - -# Update rdma-core to support latest rocshmem. -ARG DEEPEP_NIC -RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ - git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ - cd /tmp/rdma-core && \ - mkdir -p build && cd build && \ - cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ - ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ -fi - -# Install NIXL + DeepEP wheels. -RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ - --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ - uv pip install --system /nixl_install/*.whl /deep_install/*.whl - -# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. -RUN timeout --kill-after=30s 5m apt-get update -q -y --error-on=any \ - -o Acquire::Retries=3 \ - -o Acquire::http::Timeout=30 \ - -o Acquire::https::Timeout=30 \ - && apt-get install -q -y --no-install-recommends \ - -o Acquire::Retries=3 \ - -o Acquire::http::Timeout=30 \ - -o Acquire::https::Timeout=30 \ - librdmacm1 \ - libibverbs1 \ - ibverbs-providers \ - ibverbs-utils \ - util-linux \ - unzip \ - libsodium23 \ - pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ - libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ - && rm -rf /var/lib/apt/lists/* +RUN echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> /app/versions.txt \ + && echo "AINIC_VERSION=${AINIC_VERSION}" >> /app/versions.txt # Install torchcodec from source for ROCm/torch ABI compatibility. -COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh -RUN --mount=type=cache,target=/root/.cache/uv \ +RUN --mount=type=bind,source=tools/install_torchcodec_rocm.sh,target=/tmp/install_torchcodec.sh \ + --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ --mount=type=cache,target=/root/.cache/pip \ - --mount=type=cache,target=/root/.cache/torchcodec-wheels \ - bash /tmp/install_torchcodec.sh \ - && rm /tmp/install_torchcodec.sh \ - && apt-get clean && rm -rf /var/lib/apt/lists/* - -# Pre-install shared ROCm runtime dependencies. -COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -c /etc/rocm-constraints.txt \ - -r /tmp/ci-base-requirements/rocm.txt \ - && rm -rf /tmp/ci-base-requirements - -# Enable fast and less brittle model downloads in tests. -ENV HF_XET_HIGH_PERFORMANCE=1 -ENV HF_HUB_DOWNLOAD_TIMEOUT=60 - -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - -# Pre-install vLLM test dependencies. -COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r /tmp/rocm-test-reqs.txt - + --mount=type=tmpfs,target=/root/.cache/torchcodec-wheels \ + ldconfig && bash /tmp/install_torchcodec.sh +# All Python dependencies, at the versions CI tests (the test lock is a constraint). +# quart, msgpack, aiohttp, pyzmq: MoRI proxy dependencies (quart: not in the test lock). +RUN --mount=type=bind,source=requirements/common.txt,target=/r/common.txt \ + --mount=type=bind,source=requirements/rocm.txt,target=/r/rocm.txt \ + --mount=type=bind,source=requirements/test/rocm.txt,target=/r/test.txt \ + --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl \ + --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache \ + --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -c /etc/rocm-constraints.txt -c /r/test.txt -r /r/rocm.txt \ + /nixl/*.whl /lmcache/*.whl quart==0.22.0 msgpack aiohttp pyzmq \ + && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" # Rebuild fastsafetensors from source so its C++ extension is compiled with # USE_ROCM and can detect libamdhip64.so at runtime. -RUN --mount=type=cache,target=/root/.cache/pip \ - FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ +RUN --mount=type=bind,source=requirements/test/rocm.txt,target=/r/test.txt \ + --mount=type=cache,target=/root/.cache/pip \ + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /r/test.txt | head -1)" \ && test -n "${FASTSAFETENSORS_REQ}" \ && python3 -m pip install --force-reinstall --no-deps \ - --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ - && rm /tmp/rocm-test-reqs.txt + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > /app/libkineto.conf +ENV KINETO_CONFIG="/app/libkineto.conf" # Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. # See: https://github.com/pytorch/pytorch/issues/169857 ENV MIOPEN_DEBUG_CONV_DIRECT=0 ENV MIOPEN_DEBUG_CONV_GEMM=0 - # Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. # See: https://github.com/ROCm/rocm-libraries/issues/6266 ENV HSA_ENABLE_IPC_MODE_LEGACY=1 ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf - -# ROCm profiler limits workaround. -RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf -ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" - -# Install vllm_test_utils in ci_base for ci_base + wheel parity. -COPY tests/vllm_test_utils /tmp/vllm_test_utils -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system /tmp/vllm_test_utils \ - && rm -rf /tmp/vllm_test_utils - -RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \ - --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --constraint /tmp/lmcache-constraints.txt \ - /lmcache_install/*.whl \ - && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" - -# ----------------------- -# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. -FROM ${CI_BASE_IMAGE} AS test -ARG COMMON_WORKDIR - -# Install the vLLM wheel (--no-deps: all deps already in ci_base). -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && uv pip install --system --no-deps *.whl - -# Store the vLLM wheel in the image for python-only install tests. -COPY --from=export_vllm /*.whl /opt/vllm-wheels/ - -WORKDIR /vllm-workspace -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace - -# Copy in the v1 package (for python-only install test group). -COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 - -# Hide source under src/ so it won't shadow the installed package in tests. -RUN mkdir src && mv vllm src/vllm - -# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. -ENV VLLM_GPU_SYNC_CHECK=error - -FROM test AS test_smoke -ARG ROCM_SMOKE_ID -RUN --network=none \ - bash /vllm-workspace/.buildkite/scripts/rocm/smoke-test-image.sh --inside \ - && printf '%s\n' "${ROCM_SMOKE_ID:-local}" > /vllm-smoke-ok - -FROM scratch AS export_test_smoke -COPY --from=test_smoke /vllm-smoke-ok / - -# ----------------------- -# Final vLLM image -FROM mori_base AS final_common -ENV UV_COMPILE_BYTECODE=1 - -RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* - -# Clean up sccache from release image (not needed at runtime) -# This removes the binary and wrappers that may have been installed during build -RUN rm -f /usr/bin/sccache || true \ - && rm -rf /opt/sccache-wrappers || true - -# Unset sccache environment variables for the release image -# This prevents S3 bucket config from leaking into production images -ENV SCCACHE_BUCKET= -ENV SCCACHE_REGION= -ENV SCCACHE_ENDPOINT= -ENV SCCACHE_S3_NO_CREDENTIALS= -ENV SCCACHE_IDLE_TIMEOUT= - -# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. -# Manually remove it so that later steps of numpy upgrade can continue -RUN case "$(which python3)" in \ - *"/opt/conda/envs/py_3.9"*) \ - rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ - *) ;; esac - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --upgrade huggingface-hub - -# Install vLLM using uv (inherited from base stage) -# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version -# Keep OpenTelemetry within LMCache's cap. -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && echo "opentelemetry-api<=1.40.0" > /tmp/otel-constraints.txt \ - && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt -r requirements/rocm.txt \ - && pip uninstall -y vllm \ - && uv pip install --system -c /etc/rocm-constraints.txt -c /tmp/otel-constraints.txt *.whl \ - && rm /tmp/otel-constraints.txt - -# Install NIXL ROCm wheel -RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ - uv pip install --system /nixl_install/*.whl - -# Recompute the base's triton fingerprint to detect resolver interference. -RUN EXPECTED=$(grep '^TRITON_SHA256: ' /app/versions.txt | cut -d' ' -f2) \ - && ACTUAL=$(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ - -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \ - && echo "triton sha256 expected=${EXPECTED} actual=${ACTUAL}" \ - && test -n "${EXPECTED}" \ - && test "${EXPECTED}" = "${ACTUAL}" \ - && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" - -# Recompute the base's pybind11 fingerprint to detect resolver interference. -RUN EXPECTED=$(grep '^PYBIND11_SHA256: ' /app/versions.txt | cut -d' ' -f2) \ - && EXPECTED_ABI=$(grep '^PYBIND11_INTERNALS_VERSION: ' /app/versions.txt | cut -d' ' -f2) \ - && INC=$(python3 -c 'import pybind11; print(pybind11.get_include())') \ - && ACTUAL=$(find "${INC}" -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1) \ - && ACTUAL_ABI=$(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' "${INC}" | awk '{print $NF}') \ - && echo "pybind11 $(pip show pybind11 | awk '/^Version:/ {print $2}')" \ - "sha256 expected=${EXPECTED} actual=${ACTUAL}" \ - "internals expected=${EXPECTED_ABI} actual=${ACTUAL_ABI}" \ - && test -n "${EXPECTED}" \ - && test "${EXPECTED}" = "${ACTUAL}" \ - && test "${EXPECTED_ABI}" = "${ACTUAL_ABI}" - -ARG COMMON_WORKDIR -ARG BASE_IMAGE -ARG NIC_BACKEND -ARG AINIC_VERSION - -# Copy over the benchmark scripts as well -COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks -COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples -COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker - -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc -# See: https://github.com/ROCm/rocm-libraries/issues/6266 -ENV HSA_ENABLE_IPC_MODE_LEGACY=1 -ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf -ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf - ENV TOKENIZERS_PARALLELISM=false - # ENV that can improve safe tensor loading, and end-to-end time ENV SAFETENSORS_FAST_GPU=1 - # Performance environment variable. ENV HIP_FORCE_DEV_KERNARG=1 - # Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 +ARG STACK +LABEL ai.vllm.rocm.stack=${STACK} -# Workaround for ROCm profiler limits -RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf -ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt +FROM dev AS rust-toolchain +ENV CARGO_HOME=/root/.cargo +ENV RUSTUP_HOME=/root/.rustup +ENV PATH=${CARGO_HOME}/bin:${PATH} +RUN --mount=type=bind,source=rust-toolchain.toml,target=/tmp/rust-toolchain.toml \ + TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \ + && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain none \ + && rustup toolchain install "${TOOLCHAIN}" \ + && rustup default "${TOOLCHAIN}" \ + && rustc --version \ + && cargo --version -CMD ["/bin/bash"] +# Test-only dependencies, a content-keyed image like the runtime (build.sh: test-deps--), +# so every build of a key reuses these layers by digest instead of rebuilding them from cache. +FROM runtime AS test-deps +COPY --from=rust-toolchain /root/.cargo /root/.cargo +COPY --from=rust-toolchain /root/.rustup /root/.rustup +ENV CARGO_HOME=/root/.cargo RUSTUP_HOME=/root/.rustup PATH=/root/.cargo/bin:$PATH +# DeepEP and the rdma-core its rocshmem needs: CI only, as before (see build_rdma_core). +COPY --from=build_rdma_core /out / +RUN --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deepep \ + --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + ldconfig && uv pip install --system -c /etc/rocm-constraints.txt /deepep/*.whl +RUN --mount=type=bind,source=requirements/test/rocm.txt,target=/t.txt --mount=type=bind,source=tests/vllm_test_utils,target=/tu,rw \ + --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + sp="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \ + && pip freeze --path "$sp" | sort >/tmp/shipped && uv pip install --system -c /etc/rocm-constraints.txt -r /t.txt /tu \ + && ! pip freeze --path "$sp" | sort | comm -23 /tmp/shipped - | grep . # a test dep changed a shipped package (apt's python3-* may be shadowed) +RUN --mount=type=bind,from=build_mamba,src=/app/mamba_install,target=/mamba \ + --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system --no-deps /mamba/*.whl +ENV HF_XET_HIGH_PERFORMANCE=1 -FROM final_common AS final_lmcache_false +# ==== per-commit stages: nothing below is part of the runtime key ==== -FROM final_common AS final_lmcache_true -RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system /lmcache_install/*.whl \ - && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \ - && uv pip check --system \ - && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" +# The pin lives below the marker: a bump rebuilds csrc, not the runtime. +FROM dev AS triton-kernels +ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557 +ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \ + /tmp/rocm-triton.tar.gz +RUN mkdir -p /opt/rocm-triton-kernels \ + && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \ + --directory /opt/rocm-triton-kernels --strip-components=4 \ + --no-same-owner \ + "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \ + && test -f /opt/rocm-triton-kernels/__init__.py \ + && rm /tmp/rocm-triton.tar.gz +ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels ROCM_TRITON_KERNELS_COMMIT=${ROCM_TRITON_KERNELS_COMMIT} + +# The version every per-commit artifact carries. build.sh builds from an export of HEAD (no .git) +# and writes .git_archival.txt from the host's `git describe`; without it, e.g. locally, .git and a +# v* tag are used. +FROM dev AS version-calc +RUN --mount=type=bind,target=/src,rw cd /src \ + && if [ -f .git_archival.txt ]; then rm -rf .git; \ + else git config --global --add safe.directory '*'; fi \ + && python3 -c "from setuptools_scm import get_version; print(get_version(git_describe_command=['git', 'describe', '--dirty', '--tags', '--long', '--abbrev=40', '--match', 'v[0-9]*']))" \ + > /vllm-version.txt \ + && cat /vllm-version.txt \ + && if grep -q '^0\.1\.dev' /vllm-version.txt; then echo "no v* tag reachable from HEAD: git fetch --tags" >&2; exit 1; fi +FROM scratch AS version +COPY --from=version-calc /vllm-version.txt / + +# Compiles the `vllm-rs` frontend once per change to the rust sources, so the wheel build does not +# need the rust toolchain. The crates, their downloads and target/ stay in this stage's layers (no +# cache mounts), so rust-build always continues from the target/ of these exact sources. +FROM rust-toolchain AS rust-deps +# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit +# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). +ENV CARGO_BUILD_JOBS=4 +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 +WORKDIR /app/vllm +COPY requirements/build/rust.txt requirements/build/rust.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/build/rust.txt +COPY rust/Cargo.lock rust/Cargo.toml rust/ +COPY rust/proto rust/proto +COPY rust/src rust/src +COPY rust-toolchain.toml rust-toolchain.toml +COPY tools/build_rust.py tools/build_rust.sh tools/ +RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 bash tools/build_rust.sh + +# Per commit: vllm-rs embeds the vLLM version, so cargo recompiles only the crates that read it and +# relinks (~30 s). Those few rustc processes stay far from the open-file limit, and the job count +# also caps the thin-LTO threads, so lift the cap. +FROM rust-deps AS rust-build +COPY --from=version /vllm-version.txt /tmp/vllm-version.txt +RUN CARGO_BUILD_JOBS="$(nproc)" SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ + bash tools/build_rust.sh \ + && test -x vllm/vllm-rs + +# Native extensions only: copies just the files that affect the HIP/C++ build (and the map tool). +FROM triton-kernels AS csrc +WORKDIR /app/vllm +COPY setup.py CMakeLists.txt pyproject.toml ./ +COPY tools/build_rust.py tools/build_rust.py +COPY cmake cmake/ +COPY csrc csrc/ +COPY vllm/envs.py vllm/__init__.py vllm/ +# Empty: the stack's full PYTORCH_ROCM_ARCH (release); CI passes the arches it can test. +ARG ARCH +# setup.py reads requirements/rocm.txt for install_requires. This wheel only carries the .so files +# the wheel stage repackages with the real requirements, so empty files keep a requirements bump +# from recompiling HIP. +# CI builds (ARCH set) also write the kernel symbol map (CI test selection: which source file each +# GPU kernel came from). It reads the build tree, so it runs here, before the tree is removed, and +# an edit to its tool rebuilds this stage. It never fails the build. +# Release (SCCACHE_BUCKET_NAME, us-west-2): its agents run one job each, so the ccache mount starts +# empty; sccache shares per-file results in S3 through the agent's instance role. setup.py prefers +# sccache when it is on PATH; without S3 access it is removed and ccache compiles as in CI. +ARG SCCACHE_BUCKET_NAME +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + --mount=type=bind,source=tools/ci/rocm_kernel_symbol_map.py,target=/tmp/rocm_kernel_symbol_map.py \ + mkdir -p requirements && : > requirements/rocm.txt \ + && { grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + cmake/external_projects/triton_kernels.cmake \ + || echo "WARNING: Dockerfile.rocm ROCM_TRITON_KERNELS_COMMIT (${ROCM_TRITON_KERNELS_COMMIT}) does not match cmake TRITON_KERNELS_TAG; the Docker build uses the Dockerfile pin, cmake governs only non-Docker builds"; } \ + && { [ -z "${SCCACHE_BUCKET_NAME}" ] \ + || { export SCCACHE_BUCKET="${SCCACHE_BUCKET_NAME}" SCCACHE_REGION=us-west-2 SCCACHE_IDLE_TIMEOUT=0 \ + && curl -fsSL --retry 3 -o /tmp/sccache.tgz https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-unknown-linux-musl.tar.gz \ + && echo "e0ee621fb16b6940666cd770b091c62cadafd3e062dd12e3a49d9caaff3b795f /tmp/sccache.tgz" | sha256sum -c - \ + && tar -xzf /tmp/sccache.tgz -C /usr/local/bin --strip-components=1 --wildcards '*/sccache' \ + && rm /tmp/sccache.tgz && sccache --start-server; } \ + || { echo "WARNING: no S3 sccache; compiling with ccache only"; rm -f /usr/local/bin/sccache /tmp/sccache.tgz; }; } \ + && CCACHE_BASEDIR="$PWD" PYTORCH_ROCM_ARCH="${ARCH:-$PYTORCH_ROCM_ARCH}" \ + SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0+rocm.csrc.build PYTHONDONTWRITEBYTECODE=1 \ + LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=/dist \ + && { sccache --stop-server 2>/dev/null || true; } \ + && mkdir /kernel-map \ + && { [ -z "${ARCH}" ] \ + || timeout 30m python3 /tmp/rocm_kernel_symbol_map.py --build-root build --source-root . \ + --out /kernel-map/kernel_symbol_map.rocm.json.gz \ + || echo "WARNING: no kernel symbol map; CI test selection falls back to its static rule"; } \ + && rm -rf build + +# The map, stamped with this build's commit (CI selection only uses a map of the commit it ran). A +# csrc layer reused from another commit was built from the same sources, so its map holds here. +FROM dev AS kernel-symbol-map-stamp +ARG BUILDKITE_COMMIT +RUN --mount=type=bind,from=csrc,src=/kernel-map,target=/in python3 - "${BUILDKITE_COMMIT}" <<'EOF' +import gzip, json, os, sys +os.makedirs("/out") +if os.path.exists(src := "/in/kernel_symbol_map.rocm.json.gz"): + with gzip.open(src, "rt") as f: + m = json.load(f) + m["commit"] = sys.argv[1] + with gzip.open("/out/kernel_symbol_map.rocm.json.gz", "wt") as f: + json.dump(m, f, separators=(",", ":")) +EOF +FROM scratch AS kernel-symbol-map +COPY --from=kernel-symbol-map-stamp /out/ / + +# The full vLLM wheel, repackaged around the csrc extensions and the rust artifacts. +FROM dev AS wheel +COPY --from=csrc /dist /precompiled +COPY --from=rust-build /app/vllm/vllm/vllm-rs /app/vllm/vllm/_rust_*.so /rust/ +COPY --from=version /vllm-version.txt / +# The published wheel pins the exact stack builds (torch, triton, ...) it was compiled against. +RUN --mount=type=bind,target=/src,rw cd /src && cp /rust/* vllm/ \ + && ! grep -n 'git+' requirements/common.txt requirements/rocm.txt \ + && pip list --format=freeze \ + | grep -iE '^(torch|triton|torchvision|torchaudio|amdsmi|flash[-_]attn|amd[-_]aiter)==' >> requirements/rocm.txt \ + && SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /vllm-version.txt)" VLLM_USE_PRECOMPILED=1 \ + VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled/*.whl)" VLLM_DOCKER_BUILD_CONTEXT=1 \ + PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=/dist + +# The installed csrc extensions (~700 MB), as their own layer: it depends only on the csrc stage, +# so a commit that leaves csrc alone reuses the layer and pulls only the Python one below. +FROM dev AS native +RUN --mount=type=bind,from=csrc,src=/dist,target=/csrc \ + sp="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \ + && mkdir -p "/out$sp" && unzip -q /csrc/*.whl 'vllm/*.so' -d "/out$sp" + +# Bytecode compiled as if installed at / (not /out), so tracebacks name the real paths. +# The csrc extensions come from the native layer instead. +FROM runtime AS install +RUN --mount=type=bind,from=wheel,src=/dist,target=/dist \ + --mount=type=bind,from=csrc,src=/dist,target=/csrc \ + pip install --no-deps --no-compile --root /out /dist/*.whl \ + && sp="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" \ + && unzip -Z1 /csrc/*.whl 'vllm/*.so' | sed "s|^|/out$sp/|" | xargs rm \ + && python3 -m compileall -q -j 0 -s /out -p / /out + +FROM runtime AS vllm-openai +COPY --link --from=native /out / +COPY --link --from=install /out / +COPY --link benchmarks /app/vllm/benchmarks +COPY --link examples /app/vllm/examples +ARG CI_ARCH +LABEL ai.vllm.rocm.ci-arch=${CI_ARCH} +ARG VLLM_VERSION +LABEL org.opencontainers.image.version=${VLLM_VERSION} +ENTRYPOINT ["vllm", "serve"] -FROM final_lmcache_${INSTALL_LMCACHE} AS final +# Release check (cacheonly): vLLM imports offline and _C carries every arch of PYTORCH_ROCM_ARCH. +FROM vllm-openai AS openai-smoke +RUN --network=none so="$(python3 -c 'import os, vllm, vllm._C; print(os.path.dirname(vllm.__file__))')/_C.abi3.so" \ + && for a in $(echo "$PYTORCH_ROCM_ARCH" | tr ';' ' '); do \ + llvm-objdump --offloading "$so" | grep -qw "$a" || { echo "_C lacks $a" >&2; exit 1; }; done + +FROM test-deps AS test-base +# AITER nightly: replace the stack's aiter with the newest nightly wheel that fits this image's ROCm, +# torch and Python (never cached, see the bake file). AITER_SPEC pins a pip spec instead. +ARG AITER_NIGHTLY +ARG AITER_SPEC +RUN --mount=type=bind,source=.buildkite/scripts/rocm/select_aiter_nightly_wheel.py,target=/tmp/select.py \ + set -e; spec="$AITER_SPEC"; \ + if [ -z "$spec" ] && [ "$AITER_NIGHTLY" = 1 ]; then \ + out="$(python3 -P /tmp/select.py)" || { echo "$out" >&2; exit 1; }; \ + echo "$out"; spec="$(echo "$out" | head -n 1)"; \ + fi; \ + if [ -n "$spec" ]; then \ + pip install --no-cache-dir --pre --no-deps --force-reinstall "$spec" || exit 1; \ + # Locate aiter without importing it: import needs a GPU, which a build has not. A wheel built + # against a newer glibc or libstdc++ than this image installs fine but fails every test that + # loads its modules; fail here instead. + aiter="$(python3 -c 'import importlib.util as u; print(u.find_spec("aiter").origin)')"; \ + ! find "${aiter%/*}" -name '*.so' -exec ldd {} + 2>&1 | grep "version .* not found"; \ + fi -#Set entrypoint for vllm-openai official images -FROM final AS vllm-openai -ENTRYPOINT ["vllm", "serve"] +FROM test-base AS test +COPY --link --from=native /out / +COPY --link --from=install /out / +COPY --link --exclude=vllm --exclude=.git . /vllm-workspace/ +# Source under src/ so it does not shadow the installed package in tests. +COPY --link vllm /vllm-workspace/src/vllm +WORKDIR /vllm-workspace + +FROM test AS test-smoke +RUN --network=none python3 -c "import vllm, vllm._C, vllm._rocm_C; print(vllm.__version__)" + +FROM scratch AS wheels +COPY --link --from=stack-wheels / / +COPY --link --from=wheel /dist/ / + +# CI: the vLLM wheel alone (a build artifact for the python-only install test), and the expensive +# compile stages, whose cache PR builds export. +FROM scratch AS vllm-wheel +COPY --link --from=wheel /dist/ / +FROM scratch AS compile +COPY --link --from=csrc /dist /csrc +COPY --link --from=rust-deps /app/vllm/rust/Cargo.lock /rust/ diff --git a/docker/Dockerfile.rocm_72 b/docker/Dockerfile.rocm_72 deleted file mode 100644 index eeaa9d6c1b18..000000000000 --- a/docker/Dockerfile.rocm_72 +++ /dev/null @@ -1,1018 +0,0 @@ -# default base image -ARG REMOTE_VLLM="0" -ARG VLLM_REPO="https://github.com/vllm-project/vllm.git" -ARG VLLM_BRANCH="main" -ARG COMMON_WORKDIR=/app -ARG BASE_IMAGE=vllm/vllm-openai-rocm:base-nightly-rocm72 -ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base -ARG ROCM_TRITON_KERNELS_COMMIT=669b31acc1dd1b3fd93286afd5db67f65d9f7557 -ARG INSTALL_LMCACHE=false -# NIC backend for MoRI RDMA support. -# By default (all), drivers and userspace libraries for all supported NIC types -# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. -# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: -# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only -# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only -# --build-arg NIC_BACKEND=none # Install nothing. -ARG NIC_BACKEND=all -# AMD AINIC apt repo settings -# Users can specify a custom version compatible with their host drivers. -# The default version has been tested with ionic-dkms=26.03.3.001 -# (libionic1=54.0-187-1), required for MoRI-EP/NIXL on Pollara/AINIC. -ARG AINIC_VERSION=1.117.5-a-77 -ARG UBUNTU_CODENAME=jammy - -# Sccache configuration. Release builds use this today; CI can opt in when a -# shared S3-compatible cache backend is available. -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 - -FROM ${BASE_IMAGE} AS base - -ARG ARG_PYTORCH_ROCM_ARCH -ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} - -# Install build dependencies and utilities -RUN apt-get update -q -y && apt-get install -q -y \ - sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ - apt-transport-https ca-certificates wget curl \ - libnuma-dev ccache mold -RUN --mount=type=cache,target=/root/.cache/pip \ - python3 -m pip install --upgrade pip -# Note: mold is installed but not set as the system default linker because -# some packages use JIT compilation at runtime with flags mold does not support. -# Build stages opt in via LDFLAGS="-fuse-ld=mold". -# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_72_base) -ARG USE_SCCACHE -RUN if [ "$USE_SCCACHE" != "1" ]; then \ - apt-get purge -y sccache || true; \ - python3 -m pip uninstall -y sccache || true; \ - rm -f "$(which sccache)" || true; \ - fi - -# Install UV — download first, then run, so a curl failure is not masked by the pipe -RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ - && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ - && rm -f /tmp/uv-install.sh \ - && uv --version - -# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out -# Reference: https://github.com/astral-sh/uv/pull/1694 -ENV UV_HTTP_TIMEOUT=500 -ENV UV_INDEX_STRATEGY="unsafe-best-match" -# Use copy mode to avoid hardlink failures with Docker cache mounts -ENV UV_LINK_MODE=copy -# ccache directory - persisted across layer rebuilds via cache mounts. -ENV CCACHE_DIR=/root/.cache/ccache -ENV CCACHE_COMPILERCHECK=content -# Empty by default so build steps fall back to $(nproc); CI can override. -ARG max_jobs -ENV MAX_JOBS=${max_jobs} - -# Install sccache if USE_SCCACHE is enabled (for release builds) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME -ARG SCCACHE_REGION_NAME -ARG SCCACHE_S3_NO_CREDENTIALS -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - if command -v sccache >/dev/null 2>&1; then \ - echo "sccache already installed, skipping installation"; \ - sccache --version; \ - else \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi; \ - fi - -# Set sccache environment variables only when USE_SCCACHE=1 -# This prevents S3 config from leaking into images when sccache is not used -ARG USE_SCCACHE -ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} -ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} -ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} -ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} - -ARG COMMON_WORKDIR -WORKDIR ${COMMON_WORKDIR} - - -# ----------------------- -# vLLM fetch stages -FROM base AS fetch_vllm_0 -ONBUILD COPY ./ vllm/ -FROM base AS fetch_vllm_1 -ARG VLLM_REPO -ARG VLLM_BRANCH -ONBUILD RUN git clone ${VLLM_REPO} \ - && cd vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git checkout FETCH_HEAD \ - && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ - git remote add upstream "https://github.com/vllm-project/vllm.git" \ - && git fetch upstream ; fi -FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm - -# Resolve the device-independent source version once for every downstream -# artifact. setuptools-scm reads Git metadata for local/remote sources and the -# archive metadata generated for canonical CI contexts. -FROM fetch_vllm AS vllm-version -ARG COMMON_WORKDIR -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - cd ${COMMON_WORKDIR}/vllm \ - && uv pip install --system -r requirements/build/rust.txt \ - && python3 -c \ - "from setuptools_scm import get_version; print(get_version(git_describe_command=['git', 'describe', '--dirty', '--tags', '--long', '--abbrev=40', '--match', 'v[0-9]*']))" \ - > /vllm-version.txt - -# ----------------------- -# Rust toolchain inputs -# -# Keep these separate from Rust source inputs so Rust source-only changes do not -# invalidate the long-lived ci_base image that carries the toolchain. -FROM scratch AS rust_toolchain_input_0 -COPY rust-toolchain.toml /rust-toolchain-input/vllm/rust-toolchain.toml - -FROM base AS rust_toolchain_input_1 -ARG VLLM_REPO -ARG VLLM_BRANCH -RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-toolchain-input/vllm \ - && cd /rust-toolchain-input/vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git sparse-checkout init --no-cone \ - && git sparse-checkout set \ - rust-toolchain.toml \ - && git checkout FETCH_HEAD - -FROM rust_toolchain_input_${REMOTE_VLLM} AS rust-toolchain-input - -# ----------------------- -# Rust build inputs -# -# REMOTE_VLLM=0: use only the local files that can affect Rust artifacts. -# REMOTE_VLLM=1: fetch only those same files from the requested remote ref. -FROM scratch AS rust_input_0 -COPY requirements/build/rust.txt /rust-input/vllm/requirements/build/rust.txt -COPY rust/Cargo.lock rust/Cargo.toml /rust-input/vllm/rust/ -COPY rust/proto /rust-input/vllm/rust/proto -COPY rust/src /rust-input/vllm/rust/src -COPY rust-toolchain.toml /rust-input/vllm/rust-toolchain.toml -COPY tools/build_rust.py /rust-input/vllm/tools/build_rust.py -COPY tools/build_rust.sh /rust-input/vllm/tools/build_rust.sh - -FROM base AS rust_input_1 -ARG VLLM_REPO -ARG VLLM_BRANCH -RUN git clone --no-checkout --filter=blob:none ${VLLM_REPO} /rust-input/vllm \ - && cd /rust-input/vllm \ - && git fetch -v --prune -- origin ${VLLM_BRANCH} \ - && git sparse-checkout init --no-cone \ - && git sparse-checkout set \ - requirements/build/rust.txt \ - rust/Cargo.lock \ - rust/Cargo.toml \ - rust/proto/** \ - rust/src/** \ - rust-toolchain.toml \ - tools/build_rust.py \ - tools/build_rust.sh \ - && git checkout FETCH_HEAD - -FROM rust_input_${REMOTE_VLLM} AS rust-input - -# ----------------------- -# Rust toolchain -FROM base AS rust-toolchain - -ENV CARGO_HOME=/root/.cargo -ENV RUSTUP_HOME=/root/.rustup -ENV PATH=${CARGO_HOME}/bin:${PATH} - -RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ - ca-certificates curl \ - && rm -rf /var/lib/apt/lists/* - -COPY --from=rust-toolchain-input /rust-toolchain-input/vllm/rust-toolchain.toml /tmp/rust-toolchain.toml -RUN TOOLCHAIN="$(grep '^channel' /tmp/rust-toolchain.toml | sed 's/.*= *"\(.*\)"/\1/')" \ - && if ! command -v rustup >/dev/null 2>&1; then \ - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ - | sh -s -- -y --default-toolchain none; \ - fi \ - && rustup toolchain install "${TOOLCHAIN}" \ - && rustup default "${TOOLCHAIN}" \ - && rustc --version \ - && cargo --version - -# ----------------------- -# Rust build stage -# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages -# don't need the rust toolchain. -FROM rust-toolchain AS rust-build -ARG COMMON_WORKDIR -ARG USE_SCCACHE - -# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit -# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). -ENV CARGO_BUILD_JOBS=4 -ENV CARGO_NET_RETRY=10 -ENV RUSTUP_MAX_RETRIES=10 - -WORKDIR ${COMMON_WORKDIR}/vllm - -COPY --from=rust-input /rust-input/vllm/requirements/build/rust.txt requirements/build/rust.txt - -# BuildKit can run this stage in parallel with ROCm native builds. Keep Rust on -# a separate local sccache daemon while sharing the same remote cache backend. -ENV SCCACHE_SERVER_PORT=4227 - -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -r requirements/build/rust.txt - -COPY --from=rust-input /rust-input/vllm/rust rust -COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml -COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py -COPY --from=rust-input /rust-input/vllm/tools/build_rust.sh tools/build_rust.sh -COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt - -# Build the release binary. Cargo's registry/git caches can be written by -# concurrent BuildKit jobs on shared workers, so lock those cache mounts while -# keeping the cache benefit. Do not cache target/, because stale target metadata -# can outlive source updates across BuildKit cache reuse. -RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ - --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ - if [ "$USE_SCCACHE" = "1" ]; then \ - export RUSTC_WRAPPER=sccache \ - && sccache --show-stats; \ - fi \ - && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ - && bash tools/build_rust.sh \ - && test -x vllm/vllm-rs \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - sccache --show-stats; \ - fi - -# ----------------------- -# Shared vLLM build dependencies -# -# Install requirements once, before bringing source-specific inputs into either -# consumer. Native compilation and wheel packaging then share this expensive -# layer while retaining their independent source cache boundaries. -FROM base AS build_vllm_dependencies -ARG COMMON_WORKDIR -ENV VLLM_TARGET_DEVICE=rocm -WORKDIR ${COMMON_WORKDIR}/vllm - -COPY --from=fetch_vllm \ - ${COMMON_WORKDIR}/vllm/requirements/common.txt \ - ${COMMON_WORKDIR}/vllm/requirements/rocm.txt \ - requirements/ - -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -r requirements/rocm.txt - -# Fetch only the pinned ROCm Triton kernels package. Keeping this before the -# vLLM source COPYs makes the download reusable when native sources change. -FROM base AS rocm-triton-kernels -ARG ROCM_TRITON_KERNELS_COMMIT -ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \ - /tmp/rocm-triton.tar.gz -RUN mkdir -p /opt/rocm-triton-kernels \ - && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \ - --directory /opt/rocm-triton-kernels --strip-components=4 \ - --no-same-owner \ - "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \ - && test -f /opt/rocm-triton-kernels/__init__.py \ - && rm /tmp/rocm-triton.tar.gz - -# ----------------------- -# vLLM native build stages -# -# csrc-build intentionally copies only files that affect ROCm native extension -# compilation. That keeps unrelated CI/test/docs edits from invalidating the -# expensive HIP/C++ build layer. -FROM build_vllm_dependencies AS csrc-build -ARG ROCM_TRITON_KERNELS_COMMIT -COPY --from=rocm-triton-kernels \ - /opt/rocm-triton-kernels /opt/rocm-triton-kernels -ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels - -# pyproject.toml is bind-mounted so it is not persisted in the csrc-build image, -# but it still participates in the BuildKit cache key for this RUN step. -COPY setup.py CMakeLists.txt ./ -COPY tools/build_rust.py tools/build_rust.py -COPY cmake cmake/ -COPY csrc csrc/ -COPY vllm/envs.py vllm/envs.py -COPY vllm/__init__.py vllm/__init__.py - -ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" - -RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ - cmake/external_projects/triton_kernels.cmake \ - && export CCACHE_BASEDIR="$PWD" \ - && echo "=== ccache stats before ROCm native build ===" \ - && (ccache --show-stats || true) \ - && (ccache --zero-stats || true) \ - && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ - && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ - && PYTHONDONTWRITEBYTECODE=1 LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ - && test -d dist \ - && ls dist/*.whl >/dev/null \ - && echo "=== ccache stats after ROCm native build ===" \ - && (ccache --show-stats || true) - -# Build the full vLLM ROCm wheel by reusing the native extension wheel from -# csrc-build. This stage still rebuilds for Python/package changes, but skips -# the expensive dependency installation and HIP/C++ compilation when their -# respective inputs are unchanged. -FROM build_vllm_dependencies AS build_vllm -COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm -COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels -COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt - -# Drop the pre-built Rust artifacts into the source tree. setup.py detects -# them and ships them as-is, skipping the local Rust build. -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ - -RUN cd "${COMMON_WORKDIR}/vllm" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ - && export VLLM_USE_PRECOMPILED=1 \ - && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ - && export VLLM_DOCKER_BUILD_CONTEXT=1 \ - && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ - && PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=dist \ - && test -d dist \ - && ls dist/*.whl >/dev/null -FROM scratch AS export_vllm -ARG COMMON_WORKDIR -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools /tools -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72 /docker/ -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72_base /docker/ -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 - -# NIXL/UCX build stages -FROM base AS build_nixl -ARG NIXL_BRANCH="1683cf3b7f3d11674c03c5e861cea22339876c96" -ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git" -ARG UCX_BRANCH="2ff0b47ce7626584cb9b56c647d08f712c6c3edf" -ARG UCX_REPO="https://github.com/openucx/ucx.git" -ENV ROCM_PATH=/opt/rocm -ENV UCX_HOME=/usr/local/ucx -ENV NIXL_HOME=/usr/local/nixl - -# NIXL build system dependencies and RDMA support -RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ - libgrpc-dev \ - libgrpc++-dev \ - libprotobuf-dev \ - protobuf-compiler-grpc \ - libcpprest-dev \ - libaio-dev \ - librdmacm1 \ - librdmacm-dev \ - libibverbs1 \ - libibverbs-dev \ - ibverbs-utils \ - rdmacm-utils \ - ibverbs-providers \ - && rm -rf /var/lib/apt/lists/* - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system meson meson-python pyyaml types-PyYAML \ - auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0" - -RUN --mount=type=cache,target=/root/.cache/ccache \ - cd /usr/local/src && \ - git clone ${UCX_REPO} && \ - cd ucx && \ - git checkout ${UCX_BRANCH} && \ - ./autogen.sh && \ - mkdir build && cd build && \ - CC="ccache gcc" CXX="ccache g++" \ - ../configure \ - --prefix=/usr/local/ucx \ - --enable-shared \ - --disable-static \ - --disable-doxygen-doc \ - --enable-optimizations \ - --enable-devel-headers \ - --with-rocm=${ROCM_PATH} \ - --with-verbs \ - --with-dm \ - --enable-mt && \ - make -j$(nproc) && \ - make install - -ENV PATH=/usr/local/ucx/bin:$PATH -ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} - -RUN --mount=type=cache,target=/root/.cache/ccache \ - git clone ${NIXL_REPO} /opt/nixl && \ - cd /opt/nixl && \ - git checkout ${NIXL_BRANCH} && \ - CC="ccache gcc" CXX="ccache g++" \ - meson setup build --prefix=${NIXL_HOME} \ - -Ducx_path=${UCX_HOME} \ - -Dwheel_variant=rocm \ - -Dbuild_tests=false \ - -Dbuild_examples=false && \ - cd build && \ - ninja -j$(nproc) && \ - ninja install && \ - echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \ - > /etc/ld.so.conf.d/nixl.conf && \ - echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \ - >> /etc/ld.so.conf.d/nixl.conf && \ - ldconfig - -# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA, -# so configure the ROCm wheel variant directly through Meson. -# Exclude libcore and libpull from auditwheel: transitive dependencies -# that are not shipped in the wheel and vary across base images. -# Tag manylinux_2_35 to match the Ubuntu 22.04 base: NIXL >= 1.5 links -# GLIBCXX_3.4.30 symbols, which manylinux_2_34 does not allow. -RUN cd /opt/nixl && \ - ./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \ - CC="ccache gcc" CXX="ccache g++" \ - uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \ - --python ${PYTHON_VERSION} \ - -Csetup-args=-Ducx_path=${UCX_HOME} \ - -Csetup-args=-Dwheel_variant=rocm \ - -Csetup-args=-Dbuild_tests=false \ - -Csetup-args=-Dbuild_examples=false && \ - mkdir -p /tmp/nixl_wheels/repaired /app/install && \ - auditwheel repair \ - --exclude 'libamdhip64*' \ - --exclude 'libcore*' \ - --exclude 'libpull*' \ - /tmp/nixl_wheels/nixl_rocm*.whl \ - --plat manylinux_2_35_$(uname -m) \ - --wheel-dir /tmp/nixl_wheels/repaired && \ - ./contrib/wheel_add_ucx_plugins.py \ - --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ - --nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \ - /tmp/nixl_wheels/repaired/*.whl && \ - cp /tmp/nixl_wheels/repaired/*.whl /app/install - -# Fetch and verify pinned LMCache source. -FROM base AS lmcache_source -ARG LMCACHE_REPO="https://github.com/LMCache/LMCache.git" -ARG LMCACHE_REF="140819c9d57a975dbc5678a6459a218e544cb58b" -ARG LMCACHE_VERSION="0.5.3" - -# Remove NVIDIA-only dependencies and accept the release image's NumPy version. -RUN printf '%s\n' "${LMCACHE_REF}" | grep -Eq '^[0-9a-f]{40}$' \ - && git clone --depth=1 --branch "v${LMCACHE_VERSION}" \ - "${LMCACHE_REPO}" /app/lmcache \ - && test "$(git -C /app/lmcache rev-parse HEAD)" = "${LMCACHE_REF}" \ - && test "$(grep -Ec '^(cufile-python|nvtx)([<=>[:space:]]|$)' \ - /app/lmcache/requirements/common.txt)" -eq 2 \ - && test "$(grep -Fxc 'numpy<=2.2.6' \ - /app/lmcache/requirements/common.txt)" -eq 1 \ - && sed -i -E \ - -e '/^(cufile-python|nvtx)([<=>[:space:]]|$)/d' \ - -e 's/^numpy<=2\.2\.6$/numpy<=2.3.5/' \ - /app/lmcache/requirements/common.txt \ - && rm -rf /app/lmcache/.git - -FROM base AS build_lmcache -COPY --from=lmcache_source \ - /app/lmcache/requirements/build.txt /tmp/lmcache-build-requirements.txt -RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ - uv pip install --system -r /tmp/lmcache-build-requirements.txt - -ARG LMCACHE_VERSION="0.5.3" -ARG LMCACHE_ROCM_ARCH="gfx942;gfx950" -ARG USE_SCCACHE - -COPY --from=lmcache_source /app/lmcache /app/lmcache -WORKDIR /app/lmcache - -RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - set -eu; \ - printf '%s\n' "${LMCACHE_ROCM_ARCH}" \ - | grep -Eq '^gfx[0-9a-z]+(;gfx[0-9a-z]+)*$'; \ - offload_flags="$(printf '%s' "${LMCACHE_ROCM_ARCH}" \ - | sed 's/;/ --offload-arch=/g; s/^/--offload-arch=/')"; \ - if [ "${USE_SCCACHE:-0}" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers; \ - fi; \ - mkdir -p /app/install; \ - SETUPTOOLS_SCM_PRETEND_VERSION="${LMCACHE_VERSION}" \ - PYTORCH_ROCM_ARCH="${LMCACHE_ROCM_ARCH}" \ - HIPCC_COMPILE_FLAGS_APPEND="${offload_flags}" \ - BUILD_WITH_HIP=1 CXX=hipcc \ - python3 setup.py bdist_wheel --dist-dir=/app/install; \ - test "$(find /app/install -maxdepth 1 -name '*.whl' | wc -l)" -eq 1 - -RUN --mount=type=tmpfs,target=/tmp/lmcache-wheel \ - set -eu; \ - python3 -m zipfile -e /app/install/*.whl /tmp/lmcache-wheel; \ - c_ops="$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so')"; \ - test "$(find /tmp/lmcache-wheel -type f -name 'c_ops*.so' | wc -l)" -eq 1; \ - expected="$(printf '%s' "${LMCACHE_ROCM_ARCH}" | tr ';' '\n' | sort -u)"; \ - actual="$(/opt/rocm/llvm/bin/llvm-objdump --offloading "${c_ops}" \ - | grep -oE 'gfx[0-9a-z]+' | sort -u)"; \ - if [ "${actual}" != "${expected}" ]; then \ - printf 'LMCache code objects mismatch\nexpected:\n%s\nactual:\n%s\n' \ - "${expected}" "${actual}" >&2; \ - exit 1; \ - fi - -# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not -# invalidate the slow ROCShmem build. -FROM base AS build_rocshmem -ARG ROCSHMEM_BRANCH="f0acb0c6" -ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" -# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so -# it can be linked against DeepEP without arch mismatches. -ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" -ENV ROCM_PATH=/opt/rocm -ENV ROCSHMEM_DIR=/opt/rocshmem - -RUN --mount=type=cache,target=/root/.cache/ccache \ - git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ - && cd rocm-systems \ - && git sparse-checkout set --cone projects/rocshmem \ - && git checkout ${ROCSHMEM_BRANCH} \ - && mkdir -p projects/rocshmem/build \ - && cd projects/rocshmem/build \ - && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ - bash ../scripts/build_configs/all_backends \ - -DROCM_PATH=${ROCM_PATH} \ - -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ - -DUSE_EXTERNAL_MPI=OFF - -# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. -FROM build_rocshmem AS build_deepep -ARG DEEPEP_BRANCH="0f63d3e9" -ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" -ARG DEEPEP_NIC="cx7" - -# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. -# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. -# We pass --rocm-gfx942-fp8-fnuz-max 224 to align the DeepEP dispatch FP8 -# quantization bound with vLLM's own clamp -# (vllm/model_executor/layers/quantization/utils/fp8_utils.py): 224 for E4M3FNUZ -RUN --mount=type=cache,target=/root/.cache/ccache \ - export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ - && git clone ${DEEPEP_REPO} \ - && cd DeepEP \ - && git checkout ${DEEPEP_BRANCH} \ - && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} --rocm-gfx942-fp8-fnuz-max 224 bdist_wheel --dist-dir=/app/deep_install - -# MoRI runtime dependencies live in Dockerfile.rocm_72 so NIC backend changes do -# not force users to rebuild the long-lived Dockerfile.rocm_72_base image. -FROM base AS mori_base -ARG NIC_BACKEND -ARG AINIC_VERSION -ARG UBUNTU_CODENAME -RUN /bin/bash -lc 'set -euo pipefail; \ - \ - install_ainic() { \ - apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ - rm -rf /var/lib/apt/lists/*; \ - mkdir -p /etc/apt/keyrings; \ - curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ - echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ - > /etc/apt/sources.list.d/amdainic.list; \ - apt-get update && apt-get install -y --no-install-recommends \ - libionic-dev \ - ionic-common \ - ; \ - rm -rf /var/lib/apt/lists/*; \ - }; \ - \ - # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ - # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ - install_bnxt() { \ - install -m 0755 -d /etc/apt/keyrings; \ - curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ - -o /etc/apt/keyrings/broadcom-nic.asc; \ - chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ - echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ - > /etc/apt/sources.list.d/broadcom-nic.list; \ - apt-get update && apt-get install -y --no-install-recommends \ - bnxt-rocelib=235.2.86.0 \ - ; \ - cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ - ldconfig; \ - rm -rf /var/lib/apt/lists/*; \ - }; \ - \ - echo "[MORI] Install MoRI proxy deps"; \ - pip install --quiet --ignore-installed blinker && \ - pip install --quiet quart msgpack aiohttp pyzmq; \ - echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ - \ - # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ - # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ - case "${NIC_BACKEND}" in \ - none) ;; \ - all) install_ainic; install_bnxt ;; \ - ainic) install_ainic ;; \ - bnxt) install_bnxt ;; \ - *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ - esac' - -# ----------------------- -# vLLM wheel release build stage (for building distributable wheels) -# This stage pins dependencies to custom ROCm wheel versions and handles version detection -FROM fetch_vllm AS build_vllm_wheel_release - -ARG COMMON_WORKDIR -ARG ROCM_TRITON_KERNELS_COMMIT -COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt -COPY --from=rocm-triton-kernels \ - /opt/rocm-triton-kernels /opt/rocm-triton-kernels -ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels - -# Drop the pre-built Rust artifacts into the source tree. setup.py detects -# them and ships them as-is, skipping the local Rust build. -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs -COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ - -# Create /install directory for custom wheels -RUN mkdir -p /install - -# Copy custom ROCm wheels from docker/context if they exist -# COPY ensures Docker cache is invalidated when wheels change -# .keep file ensures directory always exists for COPY to work -COPY docker/context/base-wheels/ /tmp/base-wheels/ -# This is how we know if we are building for a wheel release or not. -# If there are not wheels found there, we are not building for a wheel release. -# So we exit with an error. To skip this stage. -RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ - echo "Found custom wheels - copying to /install"; \ - cp /tmp/base-wheels/*.whl /install/ && \ - echo "Copied custom wheels:"; \ - ls -lh /install/; \ - else \ - echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ - echo "Wheel releases require pre-built ROCm wheels."; \ - exit 1; \ - fi - -# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) -# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm -ARG GIT_REPO_CHECK=0 -RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ - echo "Running repository checks..."; \ - cd vllm && bash tools/check_repo.sh; \ - fi - -# Install the packaging helpers before modifying requirements/rocm.txt. -RUN --mount=type=cache,target=/root/.cache/uv \ - grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ - vllm/cmake/external_projects/triton_kernels.cmake \ - && cd vllm \ - && uv pip install --system setuptools_scm regex - -# Fail if git-based package dependencies are found in requirements files -# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) -# Extra notes: pip install is able to handle git+ URLs, but uv doesn't. -RUN echo "Checking for git-based packages in requirements files..." \ - && echo "Checking common.txt for git-based packages:" \ - && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ - echo "ERROR: Git-based packages found in common.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in common.txt"; \ - fi \ - && echo "Checking rocm.txt for git-based packages:" \ - && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ - echo "ERROR: Git-based packages found in rocm.txt:"; \ - grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ - echo "Please publish these packages to PyPI instead of using git dependencies."; \ - exit 1; \ - else \ - echo " ✓ No git-based packages found in rocm.txt"; \ - fi \ - && echo "All requirements files are clean - no git-based packages found" - -# Pin vLLM dependencies to exact versions of custom ROCm wheels -# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi -COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py -RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ - && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt - -# Install dependencies using custom wheels from /install -RUN --mount=type=cache,target=/root/.cache/uv \ - cd vllm \ - && echo "Building vLLM with custom wheels from /install" \ - && uv pip install --system --find-links /install -r requirements/rocm.txt - -# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt -# (setup.py auto-detects ccache/sccache in PATH) -RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - cd vllm \ - && export CCACHE_BASEDIR="$PWD" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ - && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ - && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist - -FROM scratch AS export_vllm_wheel_release -ARG COMMON_WORKDIR -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl / -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools /tools -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile /docker/Dockerfile -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.cpu /docker/Dockerfile.cpu -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72 /docker/ -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm_72_base /docker/ -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/ci-rocm.hcl /docker/ci-rocm.hcl -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/docker-bake.hcl -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml -COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 - -# ----------------------- -# CI base image (Tier 1) - stable, rarely changing CI dependencies. -# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers -# in the vLLM artifacts for the current commit. -FROM mori_base AS ci_base -ARG COMMON_WORKDIR - -ENV CARGO_HOME=/root/.cargo -ENV RUSTUP_HOME=/root/.rustup -ENV PATH=${CARGO_HOME}/bin:${PATH} - -# Bake the Rust toolchain into ci_base so Rust build/test steps do not -# need to fetch it on each per-commit image build. -COPY --from=rust-toolchain /root/.cargo /root/.cargo -COPY --from=rust-toolchain /root/.rustup /root/.rustup - -# Update rdma-core to support latest rocshmem. -ARG DEEPEP_NIC -RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ - git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ - cd /tmp/rdma-core && \ - mkdir -p build && cd build && \ - cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ - ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ -fi - -# Install NIXL + DeepEP wheels. -RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ - --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ - uv pip install --system /nixl_install/*.whl /deep_install/*.whl - -# Copy ROCShmem runtime libraries. -COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem - -# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. -RUN timeout --kill-after=30s 5m apt-get update -q -y --error-on=any \ - -o Acquire::Retries=3 \ - -o Acquire::http::Timeout=30 \ - -o Acquire::https::Timeout=30 \ - && apt-get install -q -y --no-install-recommends \ - -o Acquire::Retries=3 \ - -o Acquire::http::Timeout=30 \ - -o Acquire::https::Timeout=30 \ - librdmacm1 \ - libibverbs1 \ - ibverbs-providers \ - ibverbs-utils \ - util-linux \ - unzip \ - libsodium23 \ - pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ - libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ - && rm -rf /var/lib/apt/lists/* - -# Install torchcodec from source for ROCm/torch ABI compatibility. -COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=cache,target=/root/.cache/pip \ - --mount=type=cache,target=/root/.cache/torchcodec-wheels \ - bash /tmp/install_torchcodec.sh \ - && rm /tmp/install_torchcodec.sh \ - && apt-get clean && rm -rf /var/lib/apt/lists/* - -# Pre-install shared ROCm runtime dependencies. -COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ - && rm -rf /tmp/ci-base-requirements - -# Enable fast and less brittle model downloads in tests. -ENV HF_XET_HIGH_PERFORMANCE=1 -ENV HF_HUB_DOWNLOAD_TIMEOUT=60 - -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - -# Pre-install vLLM test dependencies. -COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r /tmp/rocm-test-reqs.txt - -# Rebuild fastsafetensors from source so its C++ extension is compiled with -# USE_ROCM and can detect libamdhip64.so at runtime. -RUN --mount=type=cache,target=/root/.cache/pip \ - FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ - && test -n "${FASTSAFETENSORS_REQ}" \ - && python3 -m pip install --force-reinstall --no-deps \ - --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ - && rm /tmp/rocm-test-reqs.txt - -# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. -# See: https://github.com/pytorch/pytorch/issues/169857 -ENV MIOPEN_DEBUG_CONV_DIRECT=0 -ENV MIOPEN_DEBUG_CONV_GEMM=0 - -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. -# See: https://github.com/ROCm/rocm-libraries/issues/6266 -ENV HSA_ENABLE_IPC_MODE_LEGACY=1 -ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf -ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf - -# ROCm profiler limits workaround. -RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf -ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" - -# Install vllm_test_utils in ci_base for ci_base + wheel parity. -COPY tests/vllm_test_utils /tmp/vllm_test_utils -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system /tmp/vllm_test_utils \ - && rm -rf /tmp/vllm_test_utils - -RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \ - --mount=type=bind,source=requirements/test/rocm.txt,target=/tmp/lmcache-constraints.txt \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --constraint /tmp/lmcache-constraints.txt \ - /lmcache_install/*.whl \ - && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" - -# ----------------------- -# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. -FROM ${CI_BASE_IMAGE} AS test -ARG COMMON_WORKDIR - -# Install the vLLM wheel (--no-deps: all deps already in ci_base). -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && uv pip install --system --no-deps *.whl - -# Store the vLLM wheel in the image for python-only install tests. -COPY --from=export_vllm /*.whl /opt/vllm-wheels/ - -WORKDIR /vllm-workspace -COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace - -# Copy in the v1 package (for python-only install test group). -COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 - -# Hide source under src/ so it won't shadow the installed package in tests. -RUN mkdir src && mv vllm src/vllm - -# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. -ENV VLLM_GPU_SYNC_CHECK=error - -# Run the structural smoke test inside BuildKit, where the test image layers -# are already available. Export only a build-scoped success marker. -FROM test AS test_smoke -ARG ROCM_SMOKE_ID -RUN --network=none \ - bash /vllm-workspace/.buildkite/scripts/rocm/smoke-test-image.sh --inside \ - && printf '%s\n' "${ROCM_SMOKE_ID:-local}" > /vllm-smoke-ok - -FROM scratch AS export_test_smoke -COPY --from=test_smoke /vllm-smoke-ok / - -# ----------------------- -# Final vLLM image -FROM mori_base AS final_common -ENV UV_COMPILE_BYTECODE=1 - -RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* - -# Clean up sccache from release image (not needed at runtime) -# This removes the binary and wrappers that may have been installed during build -RUN rm -f /usr/bin/sccache || true \ - && rm -rf /opt/sccache-wrappers || true - -# Unset sccache environment variables for the release image -# This prevents S3 bucket config from leaking into production images -ENV SCCACHE_BUCKET= -ENV SCCACHE_REGION= -ENV SCCACHE_ENDPOINT= -ENV SCCACHE_S3_NO_CREDENTIALS= -ENV SCCACHE_IDLE_TIMEOUT= - -# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. -# Manually remove it so that later steps of numpy upgrade can continue -RUN case "$(which python3)" in \ - *"/opt/conda/envs/py_3.9"*) \ - rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ - *) ;; esac - -RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system --upgrade huggingface-hub - -# Install vLLM using uv (inherited from base stage) -# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version -# Keep OpenTelemetry within LMCache's cap. -RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ - --mount=type=cache,target=/root/.cache/uv \ - cd /install \ - && echo "opentelemetry-api<=1.40.0" > /tmp/otel-constraints.txt \ - && uv pip install --system -c /tmp/otel-constraints.txt -r requirements/rocm.txt \ - && pip uninstall -y vllm \ - && uv pip install --system -c /tmp/otel-constraints.txt *.whl \ - && rm /tmp/otel-constraints.txt - -# Install NIXL ROCm wheel -RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \ - uv pip install --system /nixl_install/*.whl - -ARG COMMON_WORKDIR -ARG BASE_IMAGE -ARG NIC_BACKEND -ARG AINIC_VERSION - -# Copy over the benchmark scripts as well -COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks -COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples -COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker - -# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc -# See: https://github.com/ROCm/rocm-libraries/issues/6266 -ENV HSA_ENABLE_IPC_MODE_LEGACY=1 -ENV UCX_ROCM_COPY_SIGPOOL_MAX_ELEMS=inf -ENV UCX_ROCM_IPC_SIGPOOL_MAX_ELEMS=inf - -ENV TOKENIZERS_PARALLELISM=false - -# ENV that can improve safe tensor loading, and end-to-end time -ENV SAFETENSORS_FAST_GPU=1 - -# Performance environment variable. -ENV HIP_FORCE_DEV_KERNARG=1 - -# Keep torch.cuda.is_available() fork-safe (see vllm/env_override.py). -ENV PYTORCH_NVML_BASED_CUDA_CHECK=1 - -# Workaround for ROCm profiler limits -RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf -ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" -RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ - && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt - -CMD ["/bin/bash"] - -FROM final_common AS final_lmcache_false - -FROM final_common AS final_lmcache_true -RUN --mount=type=bind,from=build_lmcache,src=/app/install,target=/lmcache_install \ - --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system /lmcache_install/*.whl \ - && test "$(python3 -c 'import numpy; print(numpy.__version__)')" = "2.3.5" \ - && uv pip check --system \ - && python3 -c "import lmcache, lmcache.c_ops; print('lmcache', lmcache.__version__)" - -FROM final_lmcache_${INSTALL_LMCACHE} AS final - -#Set entrypoint for vllm-openai official images -FROM final AS vllm-openai -ENTRYPOINT ["vllm", "serve"] diff --git a/docker/Dockerfile.rocm_72_base b/docker/Dockerfile.rocm_72_base deleted file mode 100644 index 22b75d2941ea..000000000000 --- a/docker/Dockerfile.rocm_72_base +++ /dev/null @@ -1,479 +0,0 @@ -ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="669b31a" # release/internal/3.8.x as of 09/17 -ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="733fca1" # release/2.13 as of 09/17 -ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.28.0" -ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG PYTORCH_AUDIO_BRANCH="v2.11.0" -ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" -ARG FA_BRANCH="0e60e394" -ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.26" -ARG AITER_REPO="https://github.com/ROCm/aiter.git" -ARG MORI_BRANCH="v1.2.3.post1" -ARG MORI_REPO="https://github.com/ROCm/mori.git" -ARG ROCM_SYSTEMS_REPO="https://github.com/ROCm/rocm-systems.git" -ARG ROCPROFILER_SDK_COMMIT="2b22ab0195cc1461cd9abf3b969e9dd7c10af350" -ARG ROCM_RUNTIME_COMMIT="af70edc91f1dc5d228fd4c076df99e65e49c4521" - -# Sccache configuration (only used in release pipeline) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 - - -FROM ${BASE_IMAGE} AS base - -ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin -ENV ROCM_PATH=/opt/rocm -ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: - -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151 -ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} -ENV AITER_ROCM_ARCH=gfx942;gfx950 -# Note: Do not set MORI_GPU_ARCHS here, it is automatically inferred at runtime - -# Required for RCCL in ROCm7.1 -ENV HSA_NO_SCRATCH_RECLAIM=1 - -ARG PYTHON_VERSION=3.12 -ENV PYTHON_VERSION=${PYTHON_VERSION} - -RUN mkdir -p /app -WORKDIR /app -ENV DEBIAN_FRONTEND=noninteractive - -# Install Python and other dependencies (plus gcc-13, selected later in build_pytorch). -RUN apt-get update -y \ - && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ - && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \ - { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ - done \ - && apt-get update -y \ - && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 gcc-13 g++-13 \ - && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ - && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ - && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ - && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ - && python3 --version && python3 -m pip --version - -RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' 'pybind11<3.1.0' Cython -RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* - -# Install sccache if USE_SCCACHE is enabled (for release builds) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME -ARG SCCACHE_REGION_NAME -ARG SCCACHE_S3_NO_CREDENTIALS -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi - -# Setup sccache for HIP compilation via HIP_CLANG_PATH -# This creates wrapper scripts in a separate directory and points HIP to use them -# This avoids modifying the original ROCm binaries which can break detection -# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm_72) -# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache /opt/rocm/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ - fi - -# Set sccache environment variables only when USE_SCCACHE=1 -# This prevents S3 config from leaking into images when sccache is not used -ARG USE_SCCACHE -ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} -ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} -ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} -ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} - - -### -### Triton Build -### -FROM base AS build_triton -ARG TRITON_BRANCH -ARG TRITON_REPO -RUN git clone ${TRITON_REPO} -RUN cd triton \ - && git checkout ${TRITON_BRANCH} \ - && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ - && if [ ! -f setup.py ]; then cd python; fi \ - && MAX_JOBS=16 python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install - - -### -### AMD SMI Build -### -FROM base AS build_amdsmi -RUN cd /opt/rocm/share/amd_smi \ - && pip wheel . --wheel-dir=dist -RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install - - -### -### Pytorch build -### -FROM base AS build_pytorch -ARG PYTORCH_BRANCH -ARG PYTORCH_REPO -ARG USE_SCCACHE - -RUN apt-get update && apt-get install -y pkg-config liblzma-dev -# torch 2.12 emits float->__bf16 libcalls (__truncsfbf2) absent from gcc-11 libgcc; build with gcc-13. -RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 -RUN git clone ${PYTORCH_REPO} pytorch -RUN cd pytorch && git checkout ${PYTORCH_BRANCH} -RUN cd pytorch \ - && pip install -r requirements.txt && git submodule update --init --recursive -RUN cd pytorch && python3 tools/amd_build/build_amd.py \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache \ - && sccache --show-stats; \ - fi \ - && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install - - -### -### Shared PyTorch runtime for dependent builds -### -FROM base AS build_pytorch_runtime -RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install,ro \ - pip install /install/*.whl - - -### -### TorchVision build -### -FROM build_pytorch_runtime AS build_torchvision -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_VISION_REPO -ARG USE_SCCACHE - -RUN git clone ${PYTORCH_VISION_REPO} vision -RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/vision/dist/*.whl /app/install - - -### -### TorchAudio build -### -FROM build_pytorch_runtime AS build_torchaudio -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO -ARG USE_SCCACHE - -RUN git clone ${PYTORCH_AUDIO_REPO} audio -RUN cd audio && git checkout ${PYTORCH_AUDIO_BRANCH} \ - && git submodule update --init --recursive -RUN cd audio && pip install -r requirements.txt -RUN cd audio \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/audio/dist/*.whl /app/install - - -### -### MORI Build -### -FROM base AS build_mori -ARG MORI_BRANCH -ARG MORI_REPO -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${MORI_REPO} -RUN cd mori \ - && git checkout ${MORI_BRANCH} \ - && git submodule update --init --recursive \ - && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl -RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install - - -### -### FlashAttention Build -### -FROM base AS build_fa -ARG FA_BRANCH -ARG FA_REPO -ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone ${FA_REPO} -RUN cd flash-attention \ - && git checkout ${FA_BRANCH} \ - && git submodule update --init \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi -RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install - - -### -### AITER Build -### -FROM base AS build_aiter -ARG AITER_BRANCH -ARG AITER_REPO -ARG USE_SCCACHE -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - pip install /install/*.whl -RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} -RUN cd aiter \ - && git submodule update --init --recursive \ - && pip install -r requirements.txt -RUN pip install pyyaml && cd aiter \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ - && ls /app/aiter/dist/*.whl -RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install - - -### -### Final Build -### - -# Wheel release stage - -# only includes dependencies used by wheel release pipeline -FROM base AS debs_wheel_release -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs - -# Full debs stage - includes Mori (used by Docker releases) -FROM base AS debs -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchvision,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_torchaudio,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs - -### -### ROCr + CLR build -### -# Pull in updates to ROCr + clr from ROCm 7.2.4 and beyond. -FROM base AS build_rocm_runtime -ARG ROCM_SYSTEMS_REPO -ARG ROCM_RUNTIME_COMMIT -RUN apt-get update \ - && apt-get install -y --no-install-recommends g++ \ - libelf-dev libdrm-dev libnuma-dev libdw-dev xxd \ - && rm -rf /var/lib/apt/lists/* -RUN pip install --no-cache-dir CppHeaderParser -RUN git init -q /src && cd /src \ - && git remote add origin ${ROCM_SYSTEMS_REPO} \ - && git config core.sparseCheckout true \ - && git config remote.origin.promisor true \ - && git config remote.origin.partialclonefilter blob:none \ - && git sparse-checkout init --cone \ - && git sparse-checkout set projects/rocr-runtime projects/clr projects/hip shared cmake \ - && git fetch --filter=blob:none --depth 1 origin ${ROCM_RUNTIME_COMMIT} \ - && git checkout -q FETCH_HEAD -# ROCr's trap_handler and blit_shaders call find_package(Clang/LLVM REQUIRED) only to get -# IMPORTED executables. The ROCm dev image ships the binaries but no CMake package config, so -# supply one pointing at them. -RUN mkdir -p /opt/rocm-cmake-shim \ - && printf 'if(NOT TARGET clang)\n add_executable(clang IMPORTED GLOBAL)\n set_target_properties(clang PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/clang")\nendif()\nset(Clang_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/ClangConfig.cmake \ - && printf 'if(NOT TARGET llvm-objcopy)\n add_executable(llvm-objcopy IMPORTED GLOBAL)\n set_target_properties(llvm-objcopy PROPERTIES IMPORTED_LOCATION "/opt/rocm/llvm/bin/llvm-objcopy")\nendif()\nset(LLVM_FOUND TRUE)\nset(LLVM_PACKAGE_VERSION "rocm-image-shim")\n' > /opt/rocm-cmake-shim/LLVMConfig.cmake -# ROCr first, and installed into /opt/rocm, because the clr build below needs its headers -RUN cd /src/projects/rocr-runtime \ - && cmake -B build -G Ninja \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DClang_DIR=/opt/rocm-cmake-shim \ - -DLLVM_DIR=/opt/rocm-cmake-shim \ - && cmake --build build --parallel 16 \ - && cmake --install build \ - && cmake --install build --prefix /rocr-install --strip -RUN cd /src/projects/clr \ - && cmake -B build -G Ninja \ - -DCMAKE_BUILD_TYPE=Release \ - -DCLR_BUILD_HIP=ON \ - -DCLR_BUILD_OCL=OFF \ - -DHIP_COMMON_DIR=/src/projects/hip \ - -DROCM_PATH=/opt/rocm \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DLLVM_DIR=/opt/rocm-cmake-shim \ - -DHIP_LLVM_ROOT=/opt/rocm/llvm \ - && cmake --build build --parallel 16 \ - && cmake --install build --prefix /clr-install --strip -RUN mkdir -p /staging/lib \ - && cp -P /rocr-install/lib/libhsa-runtime64.so* /staging/lib/ \ - && cp -P /clr-install/lib/libamdhip64.so* /staging/lib/ \ - && ls -l /staging/lib - -# rocprofiler-sdk 1.3.2 fixes torch 2.12 profiler on ROCm 7.2.3 -FROM base AS build_rocprofiler_sdk -ARG ROCM_SYSTEMS_REPO -ARG ROCPROFILER_SDK_COMMIT -ARG USE_SCCACHE -RUN apt-get update \ - && apt-get install -y --no-install-recommends libdw-dev libelf-dev libsqlite3-dev libzstd-dev \ - && rm -rf /var/lib/apt/lists/* -RUN git init -q /src && cd /src \ - && git remote add origin ${ROCM_SYSTEMS_REPO} \ - && git config core.sparseCheckout true \ - && git config remote.origin.promisor true \ - && git config remote.origin.partialclonefilter blob:none \ - && git sparse-checkout init --cone \ - && git sparse-checkout set projects/rocprofiler-sdk shared cmake \ - && git fetch --filter=blob:none --depth 1 origin ${ROCPROFILER_SDK_COMMIT} \ - && git checkout -q FETCH_HEAD -# Patch the ROCm 7.2.3 build of rocprofiler-sdk to include queue-interposition fixes that are -# absent from the pinned 7.14 build. This resolves hangs on torch.profiler on ROCm. -# TODO: Drop these once ${ROCPROFILER_SDK_COMMIT} is bumped to a build that already -# contains them, such as ROCm 10.0. See https://github.com/vllm-project/vllm/issues/54087. -ARG ROCPROFILER_SDK_PR_7796=fae7fd57ffb8a76e811d6410b8a0cee1e94d710a -ARG ROCPROFILER_SDK_PR_7924=f8db22adbe17686a9b05ae03367dd1ec73011653 -RUN cd /src \ - && HSA_PATHS='projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/hsa/*' \ - && curl -fsSL -o /tmp/qi-7796.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7796}.patch" \ - && curl -fsSL -o /tmp/qi-7924.patch "${ROCM_SYSTEMS_REPO%.git}/commit/${ROCPROFILER_SDK_PR_7924}.patch" \ - && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7796.patch \ - && git apply -p1 --include="$HSA_PATHS" --verbose /tmp/qi-7924.patch -RUN cd /src/projects/rocprofiler-sdk \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export CMAKE_C_COMPILER_LAUNCHER=sccache \ - && export CMAKE_CXX_COMPILER_LAUNCHER=sccache; \ - fi \ - && cmake -B build -G Ninja \ - -DCMAKE_BUILD_TYPE=Release \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DROCPROFILER_BUILD_TESTS=OFF \ - -DROCPROFILER_BUILD_SAMPLES=OFF \ - -DROCPROFILER_BUILD_DOCS=OFF \ - && cmake --build build --parallel "$(nproc)" \ - && cmake --install build --prefix /staging --strip - -FROM base AS final -RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ - pip install /install/*.whl - -COPY --from=build_rocprofiler_sdk /staging/ /opt/rocm/ -ENV ROCPROFILER_QUEUE_INTERPOSITION=0 - -COPY --from=build_rocm_runtime /staging/ /opt/rocm/ -# Hard link the patched library onto every stock file name so that both the -# SONAME lookup and the hard coded paths resolve to it. -RUN set -eux; \ - cd /opt/rocm/lib; \ - ours_hsa="$(readlink -f libhsa-runtime64.so.1)"; \ - ours_hip="$(readlink -f libamdhip64.so.7)"; \ - for f in libhsa-runtime64.so.1.*; do \ - [ "$f" = "${ours_hsa##*/}" ] || ln -f "$ours_hsa" "$f"; \ - done; \ - for f in libamdhip64.so.7.*; do \ - [ "$f" = "${ours_hip##*/}" ] || ln -f "$ours_hip" "$f"; \ - done - -RUN ldconfig - -ARG BASE_IMAGE -ARG TRITON_BRANCH -ARG TRITON_REPO -ARG PYTORCH_BRANCH -ARG PYTORCH_VISION_BRANCH -ARG PYTORCH_REPO -ARG PYTORCH_VISION_REPO -ARG PYTORCH_AUDIO_BRANCH -ARG PYTORCH_AUDIO_REPO -ARG FA_BRANCH -ARG FA_REPO -ARG AITER_BRANCH -ARG AITER_REPO -ARG MORI_BRANCH -ARG MORI_REPO -ARG ROCM_SYSTEMS_REPO -ARG ROCPROFILER_SDK_COMMIT -ARG ROCM_RUNTIME_COMMIT -RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ - && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ - && echo "ROCPROFILER_SDK_COMMIT: ${ROCPROFILER_SDK_COMMIT}" >> /app/versions.txt \ - && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ - && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \ - && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_BRANCH: ${PYTORCH_AUDIO_BRANCH}" >> /app/versions.txt \ - && echo "PYTORCH_AUDIO_REPO: ${PYTORCH_AUDIO_REPO}" >> /app/versions.txt \ - && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ - && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ - && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ - && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ - && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ - && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 997af6183805..9561b4a8e24c 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,4 +1,6 @@ -ARG BASE_IMAGE=ubuntu:22.04 +ARG BASE_IMAGE=ubuntu:22.04@sha256:5ec03bb3441e8b0bf3b4f9cd4629a1ae763010dc3035bb8da3ae6cf026486401 +# Bump to rebuild the stack (changes its key and busts every layer); upstream drift alone never does. +ARG REFRESH=2026-10-12 ARG PYTHON_VERSION=3.12 ARG TORCH_VERSION=2.13.0+rocm10.0.0 ARG TORCHVISION_VERSION=0.28.0+rocm10.0.0 @@ -8,7 +10,7 @@ ARG TRITON_VERSION=3.8.0+git669b31ac ARG TRITON_REPO="https://github.com/ROCm/triton.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.26" +ARG AITER_BRANCH="v0.1.27" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.2.3.post1" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -26,14 +28,6 @@ ARG ROCM_SDK_VERSION=10.0.0 ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151" ARG AITER_ROCM_ARCH="gfx942;gfx950" -# Sccache configuration (only used in release pipeline) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME=vllm-build-sccache -ARG SCCACHE_REGION_NAME=us-west-2 -ARG SCCACHE_S3_NO_CREDENTIALS=0 - FROM ${BASE_IMAGE} AS base ARG PYTORCH_ROCM_ARCH @@ -52,6 +46,7 @@ ENV ROCM_PATH="$SDK_DEV" \ HIP_DEVICE_LIB_PATH="$SDK_CORE/lib/llvm/amdgcn/bitcode" ENV LD_LIBRARY_PATH="$SDK_DEV/lib:$SDK_DEV/lib/rocm_sysdeps/lib:$SDK_CORE/lib:$SDK_CORE/lib/rocm_sysdeps/lib" ENV PATH="$SDK_DEV/bin:$SDK_DEV/lib/llvm/bin:$SDK_CORE/bin:$PATH" +ENV ROCSHMEM_DIR="$SDK_DEV" # Required for RCCL in ROCm7.1 ENV HSA_NO_SCRATCH_RECLAIM=1 @@ -59,12 +54,16 @@ ENV HSA_NO_SCRATCH_RECLAIM=1 ARG PYTHON_VERSION ENV PYTHON_VERSION=${PYTHON_VERSION} -RUN mkdir -p /app +ARG REFRESH +RUN echo "$REFRESH" >/etc/vllm-rocm-refresh && mkdir -p /app WORKDIR /app ENV DEBIAN_FRONTEND=noninteractive # Install Python and other dependencies -RUN apt-get update -y \ +# pip runs with a cache mount, so its downloads (several GB of SDK +# and torch wheels) never land in an image layer. +RUN --mount=type=cache,target=/root/.cache/pip \ + apt-get update -y \ && apt-get install -y software-properties-common build-essential git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config libdrm-dev libnuma-dev kmod \ && for i in 1 2 3; do \ add-apt-repository -y ppa:deadsnakes/ppa && break || \ @@ -79,14 +78,16 @@ RUN apt-get update -y \ && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \ && python3 --version && python3 -m pip --version -RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11==3.0.4 Cython +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install -U packaging 'cmake<4' ninja wheel==0.48.0 'setuptools<80' pybind11==3.0.4 Cython RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* # Install ROCm SDK via TheRock ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL ARG ROCM_SDK_VERSION -RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ "rocm[libraries,devel]${ROCM_SDK_VERSION:+==${ROCM_SDK_VERSION}}" \ && rocm-sdk init @@ -95,7 +96,8 @@ RUN pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ ARG TORCH_VERSION ARG TORCHAUDIO_VERSION ARG TORCHVISION_VERSION -RUN pip install \ +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install \ --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ --extra-index-url ${ROCM_NIGHTLY_WHEELS_MULTIARCH_URL} \ "torch==${TORCH_VERSION}" \ @@ -112,7 +114,8 @@ RUN pip install \ # For each arch in PYTORCH_ROCM_ARCH, install the precompiled GPU code that # PyTorch needs to run on that specific chip. SHELL ["/bin/bash", "-c"] -RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \ +RUN --mount=type=cache,target=/root/.cache/pip \ + TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \ && for arch in $(echo "${PYTORCH_ROCM_ARCH}" | tr ';' ' '); do \ echo "Installing device packages for ${arch}" \ && pip install --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ @@ -126,7 +129,8 @@ RUN TORCH_VER=$(python3 -c "import torch; print(torch.__version__)") \ RUN rm -f "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack+.kpack" \ "${SITE_PACKAGES}/torch/.kpack/torch_gfx90a:xnack-.kpack" -RUN cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install . +RUN --mount=type=cache,target=/root/.cache/pip \ + cd $SITE_PACKAGES/_rocm_sdk_core/share/amd_smi && pip install . ENV PYTHONPATH=$SITE_PACKAGES/_rocm_sdk_core/share/amd_smi # TheRock runtime fixes @@ -156,61 +160,28 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ done \ && ldconfig -# Install sccache if USE_SCCACHE is enabled (for release builds) -ARG USE_SCCACHE -ARG SCCACHE_DOWNLOAD_URL -ARG SCCACHE_ENDPOINT -ARG SCCACHE_BUCKET_NAME -ARG SCCACHE_REGION_NAME -ARG SCCACHE_S3_NO_CREDENTIALS -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Installing sccache..." \ - && SCCACHE_ARCH="x86_64" \ - && SCCACHE_VERSION="v0.8.1" \ - && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ - && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ - && tar -xzf /tmp/sccache.tar.gz -C /tmp \ - && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ - && chmod +x /usr/bin/sccache \ - && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ - && sccache --version; \ - fi - -# Setup sccache for HIP compilation via HIP_CLANG_PATH -# This creates wrapper scripts in a separate directory and points HIP to use them -# This avoids modifying the original ROCm binaries which can break detection -# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) -# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 -RUN if [ "$USE_SCCACHE" = "1" ]; then \ - echo "Setting up sccache wrappers for HIP compilation..." \ - && mkdir -p /opt/sccache-wrappers \ - && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang++ "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang++ \ - && chmod +x /opt/sccache-wrappers/clang++ \ - && printf '#!/bin/bash\nexec sccache %s/lib/llvm/bin/clang "$@"\n' "${ROCM_PATH}" > /opt/sccache-wrappers/clang \ - && chmod +x /opt/sccache-wrappers/clang \ - && echo "sccache wrappers created in /opt/sccache-wrappers"; \ - fi - -ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} -ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} -ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} -ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} - ### ### MORI Build ### FROM base AS build_mori +# mori silently builds without UMBP when gRPC or Protobuf is missing. +RUN apt-get update \ + && apt-get install -y --no-install-recommends libgrpc++-dev libprotobuf-dev \ + protobuf-compiler protobuf-compiler-grpc \ + && rm -rf /var/lib/apt/lists/* ARG MORI_BRANCH ARG MORI_REPO RUN git clone ${MORI_REPO} -RUN cd mori \ +RUN --mount=type=cache,target=/root/.cache/pip cd mori \ && git checkout ${MORI_BRANCH} \ && git submodule update --init --recursive \ && ln -sfn ${ROCM_PATH} /opt/rocm \ && pip install "setuptools_scm>=6.2" \ && CMAKE_PREFIX_PATH="${ROCM_PATH}:${ROCM_PATH}/lib/rocm_sysdeps" \ - python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl + python3 setup.py bdist_wheel --dist-dir=dist \ + && grep -qx 'BUILD_UMBP:BOOL=ON' build/CMakeCache.txt \ + && ls /app/mori/dist/*.whl RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install ### @@ -219,17 +190,11 @@ RUN mkdir -p /app/install && cp /app/mori/dist/*.whl /app/install FROM base AS build_fa ARG FA_BRANCH ARG FA_REPO -ARG USE_SCCACHE RUN git clone ${FA_REPO} RUN cd flash-attention \ && git checkout ${FA_BRANCH} \ && git submodule update --init \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ - && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi + && GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install @@ -239,20 +204,14 @@ RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install FROM base AS build_aiter ARG AITER_BRANCH ARG AITER_REPO -ARG USE_SCCACHE RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} -RUN cd aiter \ +RUN --mount=type=cache,target=/root/.cache/pip cd aiter \ && git submodule update --init --recursive \ && pip install -r requirements.txt # AITER_USE_SYSTEM_TRITON keeps setup.py from resolving its own triton, which would # displace the pinned one already installed in base. -RUN pip install pyyaml && cd aiter \ - && if [ "$USE_SCCACHE" = "1" ]; then \ - export HIP_CLANG_PATH=/opt/sccache-wrappers \ - && sccache --show-stats; \ - fi \ +RUN --mount=type=cache,target=/root/.cache/pip pip install pyyaml && cd aiter \ && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ - && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install @@ -268,7 +227,10 @@ RUN cd triton \ && git checkout ${TRITON_BRANCH} \ && if [ ! -f setup.py ]; then cd python; fi \ && python3 setup.py bdist_wheel --dist-dir=dist \ - && mkdir -p /app/install && cp dist/*.whl /app/install + && python3 -m wheel unpack dist/*.whl -d /tmp/triton-wheel \ + && find /tmp/triton-wheel -name '*.so' -exec strip --strip-debug {} + \ + && mkdir -p /app/install \ + && python3 -m wheel pack /tmp/triton-wheel/* -d /app/install ### @@ -330,9 +292,8 @@ RUN mkdir -p /staging/lib \ ### Final Build ### -# Wheel release stage - -# only includes dependencies used by wheel release pipeline -FROM base AS debs_wheel_release +# The stack's half of the release wheel index. +FROM base AS wheels-build RUN mkdir /app/debs RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ cp /install/*.whl /app/debs @@ -348,7 +309,8 @@ ARG ROCM_RELEASE_WHEELS_MULTIARCH_URL ARG ROCM_NIGHTLY_WHEELS_MULTIARCH_URL COPY tools/vllm-rocm/therock_wheels.py /tmp/therock_wheels.py # regex and pybase64 aren't installed here; fall back to the stdlib modules. -RUN sed -i -e 's/import regex as re/import re/' \ +RUN --mount=type=cache,target=/root/.cache/pip \ + sed -i -e 's/import regex as re/import re/' \ -e 's/import pybase64 as base64/import base64/' /tmp/therock_wheels.py \ && pip download --no-deps --dest /tmp/torch-dl \ --index-url ${ROCM_RELEASE_WHEELS_MULTIARCH_URL} \ @@ -363,19 +325,12 @@ RUN sed -i -e 's/import regex as re/import re/' \ && rm -rf /tmp/torch-dl /tmp/amd_smi /tmp/therock_wheels.py \ && ls -l /app/debs -FROM base AS debs -RUN mkdir /app/debs -RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs -RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ - cp /install/*.whl /app/debs +FROM scratch AS wheels +COPY --from=wheels-build /app/debs / FROM base AS final ARG TRITON_VERSION +LABEL ai.vllm.rocm.variant=rocm100 # torch 2.12 implicitly attaches rocprofiler-sdk; its queue interposition hangs at HSA # teardown, so disable it. TODO: drop once ROCm's stock rocprofiler carries the fix. @@ -409,8 +364,17 @@ RUN TORCH_DIST_INFO=$(ls -d ${SITE_PACKAGES}/torch-*.dist-info) \ && sed -i '/^Requires-Dist: triton/d' "$TORCH_DIST_INFO/METADATA" \ && sed -i "s|^\(${TORCH_DIST_INFO##*/}/METADATA\),.*|\1,,|" "$TORCH_DIST_INFO/RECORD" -RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ - pip install /install/*.whl \ +# Runtime libraries of mori's UMBP extension. +RUN apt-get update \ + && apt-get install -y --no-install-recommends libgrpc++1 libprotobuf23 \ + && rm -rf /var/lib/apt/lists/* + +RUN --mount=type=bind,from=build_triton,src=/app/install,target=/i/triton \ + --mount=type=bind,from=build_fa,src=/app/install,target=/i/fa \ + --mount=type=bind,from=build_aiter,src=/app/install,target=/i/aiter \ + --mount=type=bind,from=build_mori,src=/app/install,target=/i/mori \ + --mount=type=cache,target=/root/.cache/pip \ + pip install /i/*/*.whl \ && ! pip check 2>&1 | grep -iE '^(torch|triton) ' \ && test "$(pip show triton | awk '/^Version:/ {print $2}')" = "${TRITON_VERSION}" \ # Create a constraints file to protect the versions we need, especially @@ -422,8 +386,16 @@ RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ && grep -q '^pybind11==3\.0\.' /etc/rocm-constraints.txt \ && python3 -c "import torch; assert torch.version.hip, f'non-ROCm torch: {torch.__version__}'; print(f'torch={torch.__version__}')" \ && python3 -c "import triton; print(f'triton={triton.__version__}')" \ + && python3 -c "from mori.cpp import UMBPClient" \ && pip list 2>/dev/null | grep -iE '^(torch|triton) ' +# Fail if a pip call above ran without the cache mount and left its downloads in the image. +RUN if [ -n "$(ls -A /root/.cache/pip 2>/dev/null)" ]; then \ + du -sh /root/.cache/pip; \ + echo "ERROR: /root/.cache/pip is not empty; add --mount=type=cache,target=/root/.cache/pip" >&2; \ + exit 1; \ + fi + ARG BASE_IMAGE ARG PYTHON_VERSION ARG ROCM_SDK_VERSION @@ -443,7 +415,6 @@ ARG MORI_REPO ARG ROCM_SYSTEMS_REPO ARG ROCM_RUNTIME_COMMIT ARG TRITON_VERSION -# Record hashes for Triton and pybind11 so we can verify later. RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ && echo "PYTHON_VERSION: ${PYTHON_VERSION}" >> /app/versions.txt \ && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ @@ -462,15 +433,5 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \ && echo "ROCM_SYSTEMS_REPO: ${ROCM_SYSTEMS_REPO}" >> /app/versions.txt \ && echo "ROCM_RUNTIME_COMMIT: ${ROCM_RUNTIME_COMMIT}" >> /app/versions.txt \ - # Fingerprint our triton so Dockerfile.rocm can prove it survived the build. && echo "TRITON_VERSION: $(pip show triton | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ - && echo "TRITON_SHA256: $(find $(python3 -c 'import triton, os; print(os.path.dirname(triton.__file__))') \ - -type f \( -name '*.py' -o -name '*.so' \) | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ - >> /app/versions.txt \ - && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt \ - && echo "PYBIND11_INTERNALS_VERSION: $(grep -rhE '^#[[:space:]]*define PYBIND11_INTERNALS_VERSION' \ - $(python3 -c 'import pybind11; print(pybind11.get_include())') | awk '{print $NF}')" \ - >> /app/versions.txt \ - && echo "PYBIND11_SHA256: $(find $(python3 -c 'import pybind11; print(pybind11.get_include())') \ - -type f -name '*.h' | sort | xargs sha256sum | sha256sum | cut -d' ' -f1)" \ - >> /app/versions.txt + && echo "PYBIND11_VERSION: $(pip show pybind11 | awk '/^Version:/ {print $2}')" >> /app/versions.txt diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index 634e522e0d62..dd22d0c27685 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -2,7 +2,7 @@ ARG REMOTE_VLLM="0" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base -ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base +ARG CI_BASE_IMAGE=ci_base # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 5e1e2984f4a4..fbde67b4e6d3 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -171,7 +171,7 @@ FROM vllm-base AS ucx-nixl-build # NIXL_VERSION (nixl contrib/Dockerfile.manylinux: ARG UCX_REF), so XPU runs the # same UCX commit as the released nixl-cu* wheels rather than an untested pairing. ARG UCX_VERSION=v1.23.x -ARG NIXL_VERSION=v1.5.0 +ARG NIXL_VERSION=v1.5.2 # Build-time only: compiler, autotools, and verbs dev headers RUN apt-get update -y && apt-get install -y --no-install-recommends \ diff --git a/docker/ci-rocm.hcl b/docker/ci-rocm.hcl deleted file mode 100644 index c251641452f6..000000000000 --- a/docker/ci-rocm.hcl +++ /dev/null @@ -1,411 +0,0 @@ -# ci-rocm.hcl - CI-specific configuration for vLLM ROCm Docker builds -# -# This file lives in the vLLM repo at docker/ci-rocm.hcl so ROCm Docker -# build mechanics can evolve with Dockerfile.rocm and docker-bake-rocm.hcl. -# Used with: docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci -# -# Registry cache: Docker Hub (rocm/vllm-ci-cache) is used exclusively. -# AMD build agents already have Docker Hub credentials (they push the test -# image to rocm/vllm-ci), so no additional credential setup is required. -# ROCm CI uses Docker Hub for BuildKit layer cache by default. A separate -# compiler cache can be enabled with USE_SCCACHE=1 when AMD provides a shared -# S3-compatible cache endpoint. - -# CI metadata - -variable "BUILDKITE_COMMIT" { - default = "" -} - -variable "BUILDKITE_BUILD_NUMBER" { - default = "" -} - -variable "BUILDKITE_BUILD_ID" { - default = "" -} - -variable "PARENT_COMMIT" { - default = "" -} - -# Merge-base of HEAD with main - provides a more stable cache fallback than -# parent commit for long-lived PRs. Mirrors the VLLM_MERGE_BASE_COMMIT -# pattern used in the shared ci.hcl file. Auto-computed by ci-bake-rocm.sh -# when unset. -variable "VLLM_MERGE_BASE_COMMIT" { - default = "" -} - -# Bridge to vLLM's COMMIT variable for OCI labels -variable "COMMIT" { - default = BUILDKITE_COMMIT -} - -# Image tags (set by CI) - -variable "IMAGE_TAG" { - default = "" -} - -variable "IMAGE_TAG_LATEST" { - default = "" -} - -# ROCm-specific GPU architecture targets - -variable "PYTORCH_ROCM_ARCH" { - default = "gfx90a;gfx942;gfx950" -} - -# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of -# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in -# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test -# stage and is irrelevant when building --target ci_base itself. -variable "CI_BASE_IMAGE" { - default = "rocm/vllm-dev:ci_base" -} - -# Leave CI_MAX_JOBS empty so the Dockerfile falls back to $(nproc) and uses -# the full builder parallelism. Operators can still override this per build. -variable "CI_MAX_JOBS" { - default = "" -} - -# Upstream dependency commit pins -- extracted from Dockerfile.rocm by -# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache -# functions produce no entries when the variable is empty. -variable "NIXL_BRANCH" { - default = "" -} - -variable "UCX_BRANCH" { - default = "" -} - -variable "DEEPEP_BRANCH" { - default = "" -} - -variable "NIXL_CACHE_KEY" { - default = "" -} - -variable "DEEPEP_CACHE_KEY" { - default = "" -} - -# Docker Hub registry cache for AMD builds. -# -# A separate repo (rocm/vllm-ci-cache) is used for BuildKit layer cache. -# Final-image cache exports use mode=min to reduce the volume of data pushed. -# Source-scoped csrc cache exports default to mode=max so fresh workers can -# recover more of the native build graph when ROCm extension inputs change. -# NOTE: mode=min still includes all layers referenced by the final image -# manifest, including inherited base layers (~7.25GB ROCm runtime). -# Docker Hub auto-creates the repo on first push. -# -# Final-image cache stays commit-scoped. Branch-to-branch reuse for the test -# image comes from importing the parent and merge-base commit cache refs. -# -# The source-scoped native cache is exported both per-commit and per-branch so -# ROCm extension rebuilds are shareable within the same commit reruns and across -# consecutive commits on the same branch without depending on a single global -# latest tag. - -variable "DOCKERHUB_CACHE_REPO" { - default = "rocm/vllm-ci-cache" -} - -variable "DOCKERHUB_CACHE_TO" { - default = "" -} - -variable "ROCM_CACHE_BRANCH_TAG" { - default = "" -} - -variable "ROCM_CACHE_UPSTREAM_BRANCH_TAG" { - default = "" -} - -variable "ROCM_CSRC_CACHE_TO_MODE" { - default = "max" -} - -variable "ROCM_RUST_CACHE_TO_MODE" { - default = "max" -} - -variable "ROCM_FINAL_CACHE_TO_MODE" { - default = "min" -} - -# Functions - -function "get_cache_from_rocm" { - params = [] - result = compact([ - # Exact commit hit - fastest cache on re-runs of the same commit - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT}" : "", - # Parent commit - useful cache for incremental changes - PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${PARENT_COMMIT}" : "", - # Merge-base with main - stable fallback for long-lived or rebased PRs; - # maps to a real main-branch commit whose cache layers are likely warm - VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${VLLM_MERGE_BASE_COMMIT}" : "", - # Import the source-scoped native build cache as well so builds whose - # Python/package layers changed can still reuse compiled ROCm objects. - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "", - PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "", - VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", - ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", - # Import the source-scoped Rust frontend cache so non-Rust changes do not - # force a fresh cargo release build. - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT}" : "", - PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${PARENT_COMMIT}" : "", - VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", - ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", - # Branch-scoped full image cache - fallback when parent-commit cache is evicted - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", - ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", - ]) -} - -function "get_cache_to_rocm" { - params = [] - result = compact([ - # Commit-scoped cache for exact re-runs. - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-${BUILDKITE_COMMIT},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "", - # Branch-scoped cache so later commits on the same branch can reuse the full - # image layers when the parent-commit cache is evicted. Unlike the old - # rocm-latest tag (which caused duplicate exporter 400s), this is per-branch. - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_FINAL_CACHE_TO_MODE}" : "", - ]) -} - -function "get_cache_from_rocm_csrc" { - params = [] - result = compact([ - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT}" : "", - PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${PARENT_COMMIT}" : "", - VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", - ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", - ]) -} - -function "get_cache_to_rocm_csrc" { - params = [] - result = compact([ - # Export the exact-commit native cache for same-commit reruns. - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-${BUILDKITE_COMMIT},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "", - # Export the branch-scoped native cache so later commits on the same branch - # can reuse compiled ROCm objects even when the exact parent cache is absent. - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:csrc-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_CSRC_CACHE_TO_MODE}" : "", - ]) -} - -function "get_cache_from_rocm_rust" { - params = [] - result = compact([ - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT}" : "", - PARENT_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${PARENT_COMMIT}" : "", - VLLM_MERGE_BASE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${VLLM_MERGE_BASE_COMMIT}" : "", - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG}" : "", - ROCM_CACHE_UPSTREAM_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_UPSTREAM_BRANCH_TAG}" : "", - ]) -} - -function "get_cache_to_rocm_rust" { - params = [] - result = compact([ - # Export exact-commit and branch-scoped Rust caches. A content-addressed - # cache ref is appended by ci-bake-rocm.sh when that wrapper is used. - BUILDKITE_COMMIT != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-${BUILDKITE_COMMIT},mode=${ROCM_RUST_CACHE_TO_MODE}" : "", - ROCM_CACHE_BRANCH_TAG != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rust-rocm-branch-${ROCM_CACHE_BRANCH_TAG},mode=${ROCM_RUST_CACHE_TO_MODE}" : "", - ]) -} - -# Cache functions for upstream dependency stages (NIXL/UCX, DeepEP). -# These stages are pinned to specific upstream commit hashes, so cache keys use -# those hashes rather than the Buildkite commit. This means the cache persists -# across all vLLM commits as long as the upstream dependency pins don't change. - -function "get_cache_from_rocm_deps" { - params = [] - result = compact([ - NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""), - DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}" : ""), - ]) -} - -function "get_cache_to_rocm_nixl" { - params = [] - result = compact([ - NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""), - ]) -} - -function "get_cache_to_rocm_deepep" { - params = [] - result = compact([ - DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY},mode=min" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH},mode=min" : ""), - ]) -} - -# CI targets - -target "_ci-rocm" { - annotations = [ - "manifest:vllm.buildkite.build_number=${BUILDKITE_BUILD_NUMBER}", - "manifest:vllm.buildkite.build_id=${BUILDKITE_BUILD_ID}", - ] - args = { - ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH - CI_BASE_IMAGE = CI_BASE_IMAGE - ROCM_SMOKE_ID = BUILDKITE_BUILD_ID - max_jobs = CI_MAX_JOBS - } -} - -target "test-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm", "_labels"] - target = "test" - cache-from = get_cache_from_rocm() - cache-to = get_cache_to_rocm() - tags = compact([ - IMAGE_TAG, - IMAGE_TAG_LATEST, - ]) - output = ["type=registry"] -} - -# Validate the test image in the shared BuildKit graph and export only the -# success marker. This avoids pulling the multi-GB image into the host daemon. -target "smoke-test-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "export_test_smoke" - cache-from = get_cache_from_rocm() - output = ["type=local,dest=./build/rocm-smoke-export"] -} - -# Cache-only target for the source-scoped ROCm native build stage. -# This persists the csrc-build stage in the registry cache even though the -# final test image only consumes it indirectly while packaging the wheel. -target "csrc-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "csrc-build" - cache-from = get_cache_from_rocm_csrc() - cache-to = get_cache_to_rocm_csrc() - output = ["type=cacheonly"] -} - -# Cache-only target for the Rust frontend build stage. Final-image cache -# exports use mode=min and do not reliably persist intermediate cargo layers, -# so Rust gets its own source-scoped cache target. -target "rust-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "rust-build" - cache-from = get_cache_from_rocm_rust() - cache-to = get_cache_to_rocm_rust() - output = ["type=cacheonly"] -} - -# Keep wheel export on the same CI graph as the test image build so the -# shared build_vllm/export_vllm stages resolve identically within one bake -# invocation. Without this, export-wheel-rocm uses the plain local target -# args while test-rocm-ci uses CI-only args, which can lead to separate -# cache lineages and inconsistent export_vllm results. -target "export-wheel-rocm" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "export_vllm" - cache-from = get_cache_from_rocm() - cache-to = get_cache_to_rocm() - output = ["type=local,dest=./wheel-export"] -} - -# Artifact-only vLLM build. GPU test jobs consume this artifact on top of -# ci_base, avoiding a per-commit multi-GB image push/pull. -group "test-rocm-ci-with-artifacts" { - targets = ["rust-rocm-ci", "csrc-rocm-ci", "export-wheel-rocm"] -} - -# Full test image + wheel export. Kept for fallback/debugging when a pushed, -# build-scoped image is useful. -group "test-rocm-ci-with-wheel" { - targets = [ - "rust-rocm-ci", - "csrc-rocm-ci", - "test-rocm-ci", - "smoke-test-rocm-ci", - "export-wheel-rocm", - ] -} - -# Primary output tag for the ci_base build. In Buildkite this is a unique, -# build-scoped tag; ci-bake-rocm.sh validates it before creating content and -# stable aliases. -variable "CI_BASE_IMAGE_TAG" { - default = "rocm/vllm-dev:ci_base" -} - -# Versioned, content-addressed trusted ref. Preview builds import it read-only; -# the wrapper creates their source-scoped content alias after validation. -variable "CI_BASE_TRUSTED_CONTENT_REF" { - default = "" -} - -# Cache-only targets for upstream dependency stages. These persist each stage -# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds -# (e.g., requirements change), these stages are cache hits if their upstream -# pins haven't changed -- saving ~35min of compilation. -target "nixl-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "build_nixl" - cache-from = get_cache_from_rocm_deps() - cache-to = get_cache_to_rocm_nixl() - output = ["type=cacheonly"] -} - -target "deepep-rocm-ci" { - inherits = ["_common-rocm", "_ci-rocm"] - target = "build_deepep" - cache-from = get_cache_from_rocm_deps() - cache-to = get_cache_to_rocm_deepep() - output = ["type=cacheonly"] -} - -# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.) -# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting -# files drifts from the remote image label. Per-PR builds then pull the result -# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit. -# Uses inline cache metadata on the ci_base image itself instead of exporting a -# separate registry cache artifact. -target "ci-base-rocm-ci" { - inherits = ["_common-rocm"] - target = "ci_base" - args = { - max_jobs = CI_MAX_JOBS - } - cache-from = concat( - compact([ - CI_BASE_IMAGE_TAG != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG}" : "", - CI_BASE_TRUSTED_CONTENT_REF != "" ? "type=registry,ref=${CI_BASE_TRUSTED_CONTENT_REF}" : "", - ]), - # Import upstream dependency caches so NIXL/DeepEP stages - # are cache hits even when ci_base itself needs rebuilding. - get_cache_from_rocm_deps(), - ) - cache-to = ["type=inline"] - tags = compact([CI_BASE_IMAGE_TAG]) - attest = ["type=provenance,disabled=true"] - output = ["type=registry"] -} - -# Group for ci_base builds -- exports dependency stage caches alongside the -# ci_base image so future rebuilds can reuse them independently. -group "ci-base-rocm-ci-with-deps" { - targets = ["nixl-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"] -} diff --git a/docker/context/base-wheels/.keep b/docker/context/base-wheels/.keep new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl index 2eb1d5a43076..75c0ebe0cb02 100644 --- a/docker/docker-bake-rocm.hcl +++ b/docker/docker-bake-rocm.hcl @@ -1,143 +1,162 @@ -# docker-bake-rocm.hcl - vLLM ROCm Docker build configuration -# -# This file lives in the vLLM repo at docker/docker-bake-rocm.hcl -# Equivalent of docker-bake.hcl for ROCm builds. -# -# Usage: -# docker buildx bake -f docker/docker-bake-rocm.hcl # Build test (default) -# docker buildx bake -f docker/docker-bake-rocm.hcl final-rocm # Build final image -# docker buildx bake -f docker/docker-bake-rocm.hcl --print # Show resolved config -# -# CI usage (with the vLLM-owned CI overlay): -# docker buildx bake -f docker/docker-bake-rocm.hcl -f docker/ci-rocm.hcl test-rocm-ci - -variable "MAX_JOBS" { - # Empty string lets the Dockerfile fall back to $(nproc) via - # MAX_JOBS="${MAX_JOBS:-$(nproc)}" in each RUN step, which uses all - # available cores on whatever machine the build runs on. - # Override with --set '*.args.max_jobs=8' for local builds on small machines. - default = "" -} - -variable "PYTORCH_ROCM_ARCH" { - default = "gfx90a;gfx942;gfx950" -} - -variable "COMMIT" { - default = "" -} - -# Content hash of ci_base-affecting files. Computed by ci-bake-rocm.sh and -# embedded as a label so future builds can compare without rebuilding. -variable "CI_BASE_CONTENT_HASH" { - default = "" -} - -variable "CI_BASE_DOCKERFILE" { - default = "docker/Dockerfile.rocm" -} - -# REMOTE_VLLM=0: use local source via Docker build context (ONBUILD COPY ./ vllm/) -# REMOTE_VLLM=1: clone from GitHub at VLLM_BRANCH (standalone builds without local source) -variable "REMOTE_VLLM" { - default = "0" -} - -variable "VLLM_BRANCH" { - default = "main" -} - -# CI_BASE_IMAGE: pre-built ci_base image for per-PR test builds. -# Defaults to the local "ci_base" stage for standalone/local builds. -# CI overrides this to "rocm/vllm-dev:ci_base" via environment variable. -variable "CI_BASE_IMAGE" { - default = "rocm/vllm-dev:ci_base" -} - -# Upstream dependency commit pins. Plain local bake builds use the Dockerfile -# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env -# overrides) and writes a small HCL override before invoking CI targets. -variable "NIXL_BRANCH" { - default = "" -} - -variable "UCX_BRANCH" { - default = "" -} - -variable "DEEPEP_BRANCH" { - default = "" -} - -group "default" { - targets = ["test-rocm"] -} - -target "_common-rocm" { - dockerfile = CI_BASE_DOCKERFILE - context = "." - args = { - max_jobs = MAX_JOBS - ARG_PYTORCH_ROCM_ARCH = PYTORCH_ROCM_ARCH - REMOTE_VLLM = REMOTE_VLLM - VLLM_BRANCH = VLLM_BRANCH - CI_BASE_IMAGE = CI_BASE_IMAGE - } -} - -target "_labels" { - labels = { - "org.opencontainers.image.source" = "https://github.com/vllm-project/vllm" - "org.opencontainers.image.vendor" = "vLLM" - "org.opencontainers.image.title" = "vLLM ROCm" - "org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs (ROCm)" - "org.opencontainers.image.licenses" = "Apache-2.0" - "org.opencontainers.image.revision" = COMMIT - } - annotations = [ - "manifest:org.opencontainers.image.revision=${COMMIT}", - ] -} - -target "test-rocm" { - inherits = ["_common-rocm", "_labels"] +# ROCm images (docker/Dockerfile.rocm). CI/release drive this via .buildkite/scripts/rocm/build.sh. Locally: +# docker buildx bake -f docker/docker-bake-rocm.hcl # test image, everything from source +# STACK_IMAGE=vllm/vllm-openai-rocm:base-nightly RUNTIME=vllm/vllm-openai-rocm:base-nightly \ +# docker buildx bake -f docker/docker-bake-rocm.hcl vllm-openai # compile only vLLM +variable "STACK" { default = "rocm" } # rocm: TheRock ROCm 10 +variable "STACK_IMAGE" { default = "" } # prebuilt stack ref@digest; empty: build it in this graph +variable "STACK_WHEELS" { default = "" } # that stack's -wheels image; empty: build it in this graph +variable "RUNTIME" { default = "" } # prebuilt runtime ref@digest; empty: build it in this graph +variable "TEST_DEPS" { default = "" } # prebuilt test-deps ref@digest; empty: build it in this graph +variable "CI_ARCH" { default = "gfx90a;gfx942;gfx950" } # MI250, MI300/325, MI355: what CI can test +variable "TAG" { default = "vllm-rocm:dev" } +variable "PUSH" { default = false } +variable "CACHE" { default = "" } # trusted cache repo (rocm/vllm-ci-cache): every CI build reads it +variable "CACHE_TO" { default = "" } # ref prefix this build writes ("" = none); build.sh sets it per repo +variable "AITER_SPEC" { default = "" } # pins the AITER nightly wheel (pip spec) +variable "AITER_NIGHTLY" { default = "" } # "1": test the newest AITER nightly wheel +variable "VLLM_VERSION" { default = "" } # release: image version label (build.sh reuse check) +variable "CACHE_CI" { default = true } # false (PR builds): skip the full ci cache; compile-cache still exports +variable "BUILDKITE_COMMIT" { default = "" } # the commit the kernel symbol map is stamped with +variable "CONTEXT" { default = "." } # build.sh: .rocm-src, an export of HEAD +variable "SCCACHE_BUCKET_NAME" { default = "" } # release: S3 sccache for the csrc compile (release-pipeline.yaml) + +function "cfrom" { # trusted cache (CI only), plus this build's own cache when CACHE_TO points elsewhere + # (without its "ci" when CACHE_CI is false: nothing writes that) + params = [names] + result = flatten([for p in distinct(compact([CACHE == "" ? "" : "${CACHE}:", CACHE_TO])) : + [for n in names : "type=registry,ref=${p}${STACK}-${n}" if CACHE_CI || n != "ci" || p != CACHE_TO]]) +} +function "cto" { + params = [name] + result = CACHE_TO == "" ? [] : ["type=registry,ref=${CACHE_TO}${STACK}-${name},mode=max,image-manifest=true,ignore-error=true,compression=zstd,compression-level=3,force-compression=true"] +} +# Images and cache layers are zstd: on the CI fleet, unpacking dominates pull time (vllm#60258, +# vllm#28656). Pulling needs Docker >= 23 or containerd >= 1.5. force-compression keeps every layer +# zstd: without it, a builder exports whichever copy of a layer it holds, e.g. a gzip one imported +# from an older cache. File mtimes stay as built: clamping them would leave every timestamp-based +# .pyc (pip's, and those written by imports during the build) stale in each new container. +function "out" { + params = [] + result = PUSH ? "type=registry,compression=zstd,compression-level=3,force-compression=true,oci-mediatypes=true" : "type=docker" +} +function "ctx" { # named contexts: prebuilt refs when given, else in-graph targets + params = [] + result = merge( + { stack = STACK_IMAGE == "" ? "target:stack-${STACK}" : "docker-image://${STACK_IMAGE}" }, + RUNTIME == "" ? {} : { runtime = "docker-image://${RUNTIME}" }, + TEST_DEPS == "" ? {} : { test-deps = "docker-image://${TEST_DEPS}" }) +} + +target "_rocm" { + context = CONTEXT + args = { SOURCE_DATE_EPOCH = "0", STACK = STACK } + attest = ["type=provenance,disabled=true"] # provenance would change the pushed digest on every build + output = [out()] +} +target "_stack" { + inherits = ["_rocm"] + target = "final" + tags = [TAG] + cache-from = cfrom(["stack"]) + cache-to = cto("stack") +} +# literal dockerfile paths: ci-infra ci_selector (images.py) only follows `dockerfile = "..."` literals +target "stack-rocm" { + inherits = ["_stack"] + dockerfile = "docker/Dockerfile.rocm_base" +} +target "stack-wheels" { + inherits = ["stack-${STACK}"] + target = "wheels" + tags = ["${TAG}-wheels"] + cache-to = [] +} +target "runtime" { + inherits = ["_rocm"] + dockerfile = "docker/Dockerfile.rocm" + target = "runtime" + contexts = ctx() + tags = [TAG] + cache-from = cfrom(["runtime"]) + cache-to = cto("runtime") +} +target "test-deps" { # content-keyed like the runtime (build.sh base mode) + inherits = ["_rocm"] + dockerfile = "docker/Dockerfile.rocm" + target = "test-deps" + contexts = ctx() + tags = [TAG] + cache-from = cfrom(["test-deps", "runtime", "ci"]) + cache-to = cto("test-deps") +} +target "_app" { + inherits = ["_rocm"] + dockerfile = "docker/Dockerfile.rocm" + contexts = ctx() + args = { ARCH = "", AITER_SPEC = AITER_SPEC, AITER_NIGHTLY = AITER_NIGHTLY, VLLM_VERSION = VLLM_VERSION, + SCCACHE_BUCKET_NAME = SCCACHE_BUCKET_NAME, CI_ARCH = CI_ARCH } + # dev/csrc records are shared. Only release (CACHE = "") writes "release"; only CI "ci" and "compile". + cache-from = cfrom(CACHE == "" ? ["runtime", "release"] : ["runtime", "ci", "compile"]) +} +target "test" { + inherits = ["_app"] target = "test" - tags = ["rocm/vllm:test"] - output = ["type=docker"] -} - -# CI base image target - builds only the ci_base stage (NIXL, DeepEP, -# torchcodec, requirements, etc.). Used by the weekly scheduled build and -# the auto-rebuild trigger when requirements change in a PR. -target "ci-base-rocm" { - inherits = ["_common-rocm", "_labels"] - target = "ci_base" - labels = { - "vllm.ci_base.content_hash" = CI_BASE_CONTENT_HASH - } - tags = ["rocm/vllm-dev:ci_base"] - output = ["type=docker"] -} - -# Wheel export target - extracts the built vLLM wheel + test workspace -# to local disk. Used by CI to upload the wheel as a Buildkite artifact -# so test jobs can assemble images locally from ci_base + wheel instead -# of pulling the full large image from Docker Hub. -# -# Usage: -# docker buildx bake -f docker/docker-bake-rocm.hcl export-wheel-rocm -# # Creates ./wheel-export/*.whl, ./wheel-export/requirements/, etc. -# -# After a full bake build, BuildKit cache makes this nearly instant. -target "export-wheel-rocm" { - inherits = ["_common-rocm"] - target = "export_vllm" - output = ["type=local,dest=./wheel-export"] -} - -target "final-rocm" { - inherits = ["_common-rocm", "_labels"] - target = "final" - tags = ["rocm/vllm:latest"] - output = ["type=docker"] -} + args = { ARCH = CI_ARCH } + tags = [TAG] + cache-to = CACHE_CI ? cto("ci") : [] + no-cache-filter = AITER_SPEC == "" && AITER_NIGHTLY != "1" ? [] : ["test-base"] # AITER nightly: always fetch the newest wheel +} +target "test-smoke" { + inherits = ["test"] + target = "test-smoke" + output = ["type=cacheonly"] + cache-to = [] +} +target "compile-cache" { # csrc + rust-deps on their own: a fresh builder (or a new PR) reuses them + inherits = ["_app"] + target = "compile" + args = { ARCH = CI_ARCH } + output = ["type=cacheonly"] + cache-to = cto("compile") +} +target "vllm-wheel" { # the python-only install test's wheel, uploaded as a build artifact + inherits = ["test"] + target = "vllm-wheel" + output = ["type=local,dest=.rocm-wheel"] + cache-to = [] +} +target "kernel-symbol-map" { # CI test selection (build.sh ci, VLLM_KERNEL_SYMBOL_MAP=1): from test's csrc + inherits = ["test"] + target = "kernel-symbol-map" + args = { BUILDKITE_COMMIT = BUILDKITE_COMMIT } + tags = [] + output = ["type=local,dest=.rocm-kernel-map"] + cache-to = [] +} +target "vllm-openai" { + inherits = ["_app"] + target = "vllm-openai" + tags = [TAG] +} +# cacheonly checks +target "openai-smoke" { + inherits = ["_app"] + target = "openai-smoke" + output = ["type=cacheonly"] +} +target "wheels" { + inherits = ["_app"] + target = "wheels" + tags = ["${TAG}-wheels"] + contexts = merge(ctx(), { stack-wheels = STACK_WHEELS == "" ? "target:stack-wheels" : "docker-image://${STACK_WHEELS}" }) + cache-to = cto("release") # the all-arch compile, for the next release +} +target "version" { + inherits = ["_app"] + target = "version" + output = ["type=local,dest=.rocm-version"] +} + +group "default" { targets = ["test"] } +group "base" { targets = ["stack-${STACK}", "stack-wheels"] } # build.sh pushes runtime separately (own key) +group "ci" { targets = ["test", "test-smoke", "compile-cache", "vllm-wheel"] } diff --git a/docker/versions.json b/docker/versions.json index 75c21afafc69..84e4ac65d915 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -14,7 +14,7 @@ "default": "false" }, "NCCL_VERSION": { - "default": "2.30.7" + "default": "2.31.2" }, "BUILD_BASE_IMAGE": { "default": "pytorch/manylinux2_28-builder:cuda13.0-78e737ad29420ffc4800e677c51e2a852caf8359" @@ -37,9 +37,6 @@ "INSTALL_KV_CONNECTORS": { "default": "false" }, - "BUILD_NIXL": { - "default": "false" - }, "SCCACHE_BUCKET_NAME": { "default": "vllm-build-sccache" }, @@ -83,7 +80,7 @@ "default": ">=1.0.17" }, "RUNAI_MODEL_STREAMER_VERSION": { - "default": ">=0.15.7" + "default": ">=0.17.1" }, "GGUF_VERSION": { "default": ">=0.19.0" diff --git a/docs/.nav.yml b/docs/.nav.yml index 783c40bb5c6d..0b470f200019 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -30,7 +30,7 @@ nav: - models/extensions - Hardware Supported Models: - models/hardware_supported_models/* - - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models_features/ + - TPU: https://docs.vllm.ai/projects/tpu/en/latest/recommended_models/ - Features: features - Developer Guide: - contributing/README.md diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index a38ef0ab66af..533b89816971 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -9,6 +9,48 @@ It's a living document, updated as new features and datasets become available. For benchmarking production vLLM servers, we recommend [GuideLLM](https://github.com/vllm-project/guidellm), an established performance benchmarking framework with live progress updates and automatic report generation. It is also more flexible than `vllm bench serve` in terms of dataset loading, request formatting, and workload patterns. +## Model Forward Benchmark + +For exact-batch model timing, including steady-state pipeline-parallel forwards, +see [vllm bench forward](../cli/README.md#forward). + +For skewed synthetic MoE expert loads, use `--routing zipf`: + +```bash +vllm bench forward --model Qwen/Qwen3-30B-A3B \ + --routing zipf --case decode=6*1@8192 --case prefill=2048 +``` + +Each token independently selects distinct experts without replacement, weighted +by popularity rank to the power of `-1.25`. Expert popularity is randomly +permuted per layer to avoid systematically assigning hot experts to the lowest +EP ranks. +The fixture is seeded by `--seed`, fixed across timing replays, and generated +outside the timed forward. Sequence-parallel TP ranks use their corresponding +token rows. Zipf allows EP load imbalance; `--routing balanced` gives equal +load per EP rank. Neither synthetic mode supports grouped top-k routing; +Zipf also does not support prefill context parallelism. + +The default `--routing auto` selects Zipf with dummy weights and real routing +otherwise. When the model or parallel layout does not support the synthetic +fixture, `auto` warns and falls back to real routing. Use `--routing balanced` +to explicitly request equal EP load. Zipf is a synthetic workload, not a fit +to a particular model's routing or speculative-token correlations. + +## Attention Kernel Benchmark + +Run the attention kernel smoke benchmark from the repository root using the +exact backend names: + +```bash +python3 benchmarks/attention_benchmarks/benchmark.py \ + --backends FLASH_ATTN FLASHINFER --batch-specs "8q1s1k" +``` + +If any configuration fails, the benchmark exits with status `1` after reporting +results and writing any requested CSV or JSON output. Successful runs exit with +status `0`. Profiling runs also propagate failures, but do not save timing metrics. + ## Dataset Overview