diff --git a/.buildkite/intel_jobs/entrypoints_intel.yaml b/.buildkite/intel_jobs/entrypoints_intel.yaml index efed04c1b383..e8006bac8f66 100644 --- a/.buildkite/intel_jobs/entrypoints_intel.yaml +++ b/.buildkite/intel_jobs/entrypoints_intel.yaml @@ -134,7 +134,8 @@ steps: 'cd tests && pytest -v -s tool_use --deselect="tests/tool_use/mistral/test_mistral_tool_calls.py::test_tool_call_none_with_tools[ministral-3b]" - --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" && + --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[granite-3.0-8b]" + --deselect="tests/tool_use/test_chat_completions.py::test_chat_completion_with_tools[llama3.2]" && pytest -v -s entrypoints/tool_parsers && pytest -v -s entrypoints/generate && pytest -v -s entrypoints/anthropic && diff --git a/.buildkite/intel_jobs/lm_eval_intel.yaml b/.buildkite/intel_jobs/lm_eval_intel.yaml index 11d3edc90e10..1ef00a47945a 100644 --- a/.buildkite/intel_jobs/lm_eval_intel.yaml +++ b/.buildkite/intel_jobs/lm_eval_intel.yaml @@ -23,5 +23,6 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && + 'uv pip install "gpt-oss[eval]==0.0.5" && + cd tests && pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-xpu.txt' diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index a76331f4d5c4..ef2b4ba07396 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -97,7 +97,7 @@ steps: commands: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh - 'cd tests && + '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 && 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" && diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 15f8274fc62b..6c5f3308b785 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -540,6 +540,26 @@ prepare_ci_build_context() { 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}}" + if [[ " ${stages} " != *" build_rocshmem "* ]]; then + unset ROCSHMEM_BRANCH ROCSHMEM_CACHE_KEY DEEPEP_CACHE_KEY + fi +} + compose_dependency_cache_key() { local prefix="$1" local material="$2" @@ -569,11 +589,11 @@ hash_dockerfile_stages() { } emit = 0 } - $1 == "FROM" { + toupper($1) == "FROM" { stage = "" for (idx = 1; idx <= NF; idx++) { if (tolower($idx) == "as" && idx < NF) { - stage = $(idx + 1) + stage = tolower($(idx + 1)) } } emit = (stage in wanted) @@ -613,11 +633,11 @@ discover_dockerfile_stage_args() { } { line = $0 - if ($1 == "FROM") { + if (toupper($1) == "FROM") { stage = "" for (idx = 1; idx <= NF; idx++) { if (tolower($idx) == "as" && idx < NF) { - stage = $(idx + 1) + stage = tolower($(idx + 1)) } } emit = (stage in wanted) @@ -630,7 +650,7 @@ discover_dockerfile_stage_args() { for (idx = 1; idx <= line_count; idx++) { line = lines[idx] arg_name = line - sub(/^[[:space:]]*ARG[[:space:]]+/, "", arg_name) + 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_]*$/) { @@ -716,7 +736,7 @@ extract_dockerfile_arg_default() { if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" fi - sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ "${physical_dockerfile}" | head -1 } @@ -1000,7 +1020,21 @@ create_and_bootstrap_builder() { } 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}}}" @@ -1010,10 +1044,12 @@ init_config() { 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}}" - CI_BASE_DOCKERFILE_STAGES="${CI_BASE_DOCKERFILE_STAGES:-${DEFAULT_CI_BASE_DOCKERFILE_STAGES}}" + 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 + 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" @@ -1042,6 +1078,11 @@ print_header() { } 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" @@ -1069,6 +1110,11 @@ load_ci_hcl() { 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() { @@ -1804,7 +1850,7 @@ compute_rocm_csrc_content_hash() { local -a content_args=() bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" read -r -a content_paths <<< "${content_files}" mapfile -t content_args < <( get_content_arg_names "${dockerfile_rocm}" "${stages}" "${ROCM_CSRC_CONTENT_ARGS:-}" @@ -1853,7 +1899,7 @@ compute_rocm_rust_content_hash() { local -a content_args=() bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" read -r -a content_paths <<< "${content_files}" mapfile -t content_args < <( get_content_arg_names \ @@ -1943,7 +1989,7 @@ write_rocm_build_arg_override() { local arg_value="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + dockerfile_rocm="${CI_BASE_DOCKERFILE:-${bake_dir}/Dockerfile.rocm}" mapfile -t arg_names < <( { get_content_arg_names \ @@ -2259,7 +2305,7 @@ extract_dependency_pins() { local val="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + 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}" @@ -2275,12 +2321,12 @@ extract_dependency_pins() { fi val=$( - sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + 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" + echo "Extracted ${var}=${val} from ${dockerfile_rocm}" fi done } @@ -2297,7 +2343,7 @@ compute_dependency_cache_keys() { local deepep_material="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") - dockerfile_rocm="${bake_dir}/Dockerfile.rocm" + 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") rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH") @@ -2687,6 +2733,7 @@ main() { 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 @@ -2697,6 +2744,7 @@ main() { # 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 diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index c9897e732e51..7f1b700fe6bd 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -65,6 +65,39 @@ _ci_otel_python() { fi } +# The Docker test commands run as root, while /workdir is owned by the +# Buildkite checkout user. Run the GPU sampler as that unprivileged identity so +# any future project-code import cannot leave root-owned files in the checkout. +_ci_otel_prepare_gpu_user() { + _CI_INFRA_GPU_UID="" + _CI_INFRA_GPU_GID="" + [ "$(id -u 2>/dev/null || :)" = "0" ] || return 0 + command -v setpriv >/dev/null 2>&1 || return 1 + + _CI_INFRA_OTEL_WORKSPACE_DIR="${CI_INFRA_OTEL_WORKSPACE_DIR:-/workdir}" + [ -d "${_CI_INFRA_OTEL_WORKSPACE_DIR}" ] || return 1 + # The owner fields are numeric and whitespace-free on GNU stat. + # shellcheck disable=SC2046 + set -- $(stat -c '%u %g' -- "${_CI_INFRA_OTEL_WORKSPACE_DIR}" 2>/dev/null || :) + [ "$#" -eq 2 ] || return 1 + case "$1:$2" in + *[!0-9:]* | 0:* | *:0) return 1 ;; + esac + + _CI_INFRA_GPU_UID="$1" + _CI_INFRA_GPU_GID="$2" + chown "${_CI_INFRA_GPU_UID}:${_CI_INFRA_GPU_GID}" \ + "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + chmod 0700 "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then + # Permit traversal to the private child spool, but not directory listing. + chmod 0711 "${CI_INFRA_OTEL_RUNTIME_DIR}" || return 1 + fi + setpriv --reuid="${_CI_INFRA_GPU_UID}" \ + --regid="${_CI_INFRA_GPU_GID}" --clear-groups --no-new-privs \ + -- test -w "${CI_INFRA_OTEL_SPOOL_DIR}" +} + ci_otel_start() { _CI_INFRA_OTEL_COMMAND_INDEX="$1" _CI_INFRA_OTEL_COMMAND_LABEL="${2:-command ${_CI_INFRA_OTEL_COMMAND_INDEX}}" @@ -87,9 +120,23 @@ ci_otel_start() { if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && command -v nvidia-smi >/dev/null 2>&1; then - PYTHONDONTWRITEBYTECODE=1 \ - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" &2 || : + fi fi } diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ae51010f9eeb..c66530989f97 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -987,6 +987,85 @@ def upload(spans, timeout_seconds): assert [len(batch[0].events) for batch in batches] == [1, 1, 1] +def _write_executable(path: Path, body: str) -> None: + path.write_text(f"#!/bin/sh\n{body}") + path.chmod(0o755) + + +@pytest.mark.parametrize( + "workspace_owner,expect_sampler", + [("997 996", True), ("0 0", False)], +) +def test_root_shell_runs_gpu_sampler_as_checkout_owner( + tmp_path, workspace_owner, expect_sampler +): + bin_dir = tmp_path / "bin" + workspace = tmp_path / "workdir" + runtime = tmp_path / "runtime" + log = tmp_path / "calls.log" + bin_dir.mkdir() + workspace.mkdir() + + _write_executable( + bin_dir / "id", + '[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n', + ) + _write_executable( + bin_dir / "stat", + f"printf '%s\\n' '{workspace_owner}'\n", + ) + for command in ("chown", "chmod"): + _write_executable(bin_dir / command, "exit 0\n") + _write_executable( + bin_dir / "setpriv", + 'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n' + 'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n' + '[ "$#" -gt 0 ] && shift\n' + 'exec "$@"\n', + ) + _write_executable( + bin_dir / "python3", + '[ "$1" = "-c" ] && exit 0\n' + '[ "$2" = "new-context" ] && {\n' + ' echo "01010101010101010101010101010101 0202020202020202 - 1"\n' + " exit 0\n" + "}\n" + 'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n' + "exit 0\n", + ) + _write_executable(bin_dir / "nvidia-smi", "exit 0\n") + shell = ( + f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' + f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" + ) + + result = subprocess.run( + ["/bin/sh", "-c", shell], + check=False, + capture_output=True, + text=True, + env={ + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), + "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), + "CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace), + "CI_OTEL_TEST_LOG": str(log), + }, + ) + + assert result.returncode == 0, result.stderr + calls = log.read_text() if log.exists() else "" + if expect_sampler: + expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --" + assert calls.count("setpriv") == 2 + assert calls.count(expected) == 2 + assert "gpu" in calls + else: + assert calls == "" + assert "non-root GPU sampler unavailable" in result.stderr + + @pytest.mark.parametrize("sampling", ["1", "0"]) def test_gpu_sampler_stops_with_command_and_preserves_failure_status( tmp_path, sampling diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 62fda27bf0be..541a8f55258c 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -34,6 +34,9 @@ ############################################################################### 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}" @@ -106,6 +109,10 @@ clear_ci_orchestration_env() { 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 \ @@ -225,6 +232,9 @@ prepare_artifact_image() { 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" diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 052063f2f69e..77e0f3e07bad 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -112,7 +112,9 @@ print_packages export CONTAINER_NAME export -f kernel_tests model_tests serving_tests case "$SHARD_ID" in - 0) timeout 30m bash -c kernel_tests ;; + # The kernel shard contains several large parametrized suites and can exceed + # 30 minutes on otherwise healthy Arm runners. + 0) timeout 60m bash -c kernel_tests ;; 1) timeout 30m bash -c model_tests ;; 2) timeout 30m bash -c serving_tests ;; esac diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 74c7749d50ef..e8d0a24bc7d1 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -59,7 +59,7 @@ docker pull ubuntu:25.04 || true # building the docker image echo "--- :docker: Building Docker image" -BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:' +BUILD_RETRY_PATTERN='dial tcp|i/o timeout|failed to authorize|TLS handshake timeout|connection reset|PROTOCOL_ERROR|Could not resolve host|Temporary failure in name resolution|dns error: failed to lookup address information|client error \(Connect\)|error sending request for url|Failed to fetch:|The read operation timed out|BrokenPipeError:.*Broken pipe|HTTP/[0-9.]+ stream [0-9]+ was not closed cleanly|fetch-pack: unexpected disconnect|fatal: early EOF|invalid index-pack output' BUILD_MAX_ATTEMPTS=4 # 1 initial + 3 retries BUILD_RETRY_WAITS=(10 20 40) # seconds to wait before retry 1/2/3 build_log="$(mktemp)" diff --git a/.buildkite/scripts/install-kv-offload.sh b/.buildkite/scripts/install-kv-offload.sh new file mode 100644 index 000000000000..46d756b621b1 --- /dev/null +++ b/.buildkite/scripts/install-kv-offload.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +set -euo pipefail + +if [ -z "${CUDA_VERSION:-}" ]; then + exit 0 +fi + +# TODO: Use the PyPI wheel once it is published. +uv pip install --system \ + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0" + +# Keep only the NIXL wheel matching the CI image's CUDA runtime. +NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p') +uv pip uninstall --system nixl-cu12 nixl-cu13 +uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}" diff --git a/.buildkite/scripts/rocm/build-config.sh b/.buildkite/scripts/rocm/build-config.sh new file mode 100644 index 000000000000..ad0179b16359 --- /dev/null +++ b/.buildkite/scripts/rocm/build-config.sh @@ -0,0 +1,85 @@ +#!/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 Rock, set CI_ROCM_DOCKERFILE_BASE=docker/Dockerfile.rock_base and +# CI_ROCM_DOCKERFILE=docker/Dockerfile.rock. +# 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/refresh-base-image.sh b/.buildkite/scripts/rocm/refresh-base-image.sh index 0e10b98c9eb2..10e594577a87 100755 --- a/.buildkite/scripts/rocm/refresh-base-image.sh +++ b/.buildkite/scripts/rocm/refresh-base-image.sh @@ -4,6 +4,10 @@ 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}}" @@ -12,7 +16,12 @@ DEFAULT_ROCM_BASE_METADATA_VERSION="2" DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}" ROCM_BASE_LAYER_CACHE_REF="" -ROCM_BASE_TRUSTED_LAYER_CACHE_REF="${CACHE_REPO}:rocm-base-main" +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=() @@ -120,7 +129,7 @@ configure_rocm_base_layer_cache() { fi scope=$(rocm_base_layer_cache_scope) - ROCM_BASE_LAYER_CACHE_REF="${CACHE_REPO}:rocm-base-${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}" ) @@ -139,7 +148,7 @@ configure_rocm_base_layer_cache() { extract_arg_default() { local arg_name="$1" - sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ + sed -n -E "s/^[[:space:]]*[Aa][Rr][Gg][[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ "${DOCKERFILE}" | head -1 } @@ -419,11 +428,16 @@ build_base_image() { 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="$(rocm_version_from_base_image "${base_image_arg}")" + 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)" @@ -569,6 +583,13 @@ build_base_image() { } 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 diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 5ab0e836ebd8..226d7ac94ae7 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -287,6 +287,7 @@ steps: PREFILL_BLOCK_SIZE: "256" DECODE_BLOCK_SIZE: "256" MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" + VLLM_ENGINE_READY_TIMEOUT_S: "900" 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: "" diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d9afefc82e4f..ec45f7e69e73 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -116,6 +116,24 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt +- label: ":nvidia: (B200) LM Eval HiSparse Nightly" + key: lm-eval-hisparse-nightly-4xb200 + timeout_in_minutes: 60 + device: b200-k8s + num_devices: 4 + optional: true + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/v1/hisparse/ + - vllm/distributed/kv_transfer/kv_connector/v1/hisparse/ + - vllm/v1/attention/backends/mla/ + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - vllm/v1/worker/gpu/ + - csrc/libtorch_stable/hisparse_kernels.cu + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt + - label: ":nvidia: (B200) LM Eval PCP" key: lm-eval-pcp-4xb200 timeout_in_minutes: 360 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index b45b26b2f221..940447bf1b52 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -130,6 +130,7 @@ steps: source_file_dependencies: *v1-core-kv-metrics-deps commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh - pytest -v -s v1/kv_offload - pytest -v -s v1/simple_kv_offload mirror: @@ -141,19 +142,20 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors" +- label: ":nvidia: (H200 MIG 35GB) V1 KV Connectors Shard %N" device: h200_35gb key: v1-kv-connectors timeout_in_minutes: 45 + parallelism: 4 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 + - 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 mirror: amd: - label: ":amd: (MI300) V1 KV Connectors" + label: ":amd: (MI300) V1 KV Connectors Shard %N" dind: false device: mi300_1 timeout_in_minutes: 45 diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml index fda4add1d726..dd695d28d529 100644 --- a/.github/workflows/buf.yml +++ b/.github/workflows/buf.yml @@ -5,12 +5,6 @@ on: paths: - "rust/proto/**" - ".github/workflows/buf.yml" - push: - tags: - - "v*" - schedule: - - cron: "0 8 * * *" - workflow_dispatch: permissions: contents: read @@ -21,37 +15,14 @@ concurrency: jobs: proto: - if: github.repository_owner == 'vllm-project' runs-on: ubuntu-latest steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 - with: - fetch-depth: 0 - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0 with: version: 1.72.0 input: rust/proto - token: ${{ github.repository == 'vllm-project/vllm' && github.event_name != 'pull_request' && secrets.BUF_TOKEN || '' }} format: false breaking: false pr_comment: false push: false - - name: Publish nightly schema - if: github.repository == 'vllm-project/vllm' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') - run: >- - buf push rust/proto - --create - --create-default-label main - --create-visibility public - --label nightly - --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" - - name: Publish release schema - if: github.repository == 'vllm-project/vllm' && github.event_name == 'push' - run: >- - buf push rust/proto - --create - --create-default-label main - --create-visibility public - --label main - --label "${GITHUB_REF_NAME}" - --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" diff --git a/benchmarks/kernels/benchmark_persistent_topk.py b/benchmarks/kernels/benchmark_persistent_topk.py new file mode 100644 index 000000000000..a77a04cfd625 --- /dev/null +++ b/benchmarks/kernels/benchmark_persistent_topk.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark sparse-indexer decode top-k on variable-length FP32 rows. + +Run on a reserved GPU, for example: + chg run -- .venv/bin/python benchmarks/kernels/benchmark_persistent_topk.py + +Use --capacity to emulate graphs captured at a larger maximum context length. +Run the same command on the base and candidate builds to compare JSON results. +""" + +import argparse +import itertools +import json +import statistics +from pathlib import Path + +import torch +from flashinfer.testing import bench_gpu_time_with_cupti + +import vllm._custom_ops # noqa: F401 +from vllm.platforms import current_platform + + +def benchmark(rows: int, width: int, k: int, args) -> dict: + capacity = args.capacity or width + if capacity < width: + raise ValueError("Capacity must be at least the valid-length upper bound") + torch.manual_seed(args.seed) + logits = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) + if args.full_length: + lengths = torch.full((rows,), width, device="cuda", dtype=torch.int32) + else: + lengths = torch.randint( + int(width * 0.8), width + 1, (rows,), device="cuda", dtype=torch.int32 + ) + invalid = torch.arange(capacity, device="cuda")[None] >= lengths[:, None] + logits.masked_fill_(invalid, float("nan")) + out = torch.empty((rows, k), device="cuda", dtype=torch.int32) + workspace = torch.empty(1024 * 1024, device="cuda", dtype=torch.uint8) + cooperative = ( + args.backend == "auto" + and rows <= 64 + and capacity % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + op = torch.ops._C.cooperative_topk if cooperative else torch.ops._C.persistent_topk + + def run(): + op(logits, lengths, out, workspace, k, capacity) + + run() + valid = torch.arange(k, device="cuda")[None] < lengths[:, None] + assert torch.all(out[~valid] == -1) + assert torch.all(((out >= 0) & (out < lengths[:, None])) == valid) + sorted_indices = out.sort(dim=1).values + assert torch.all( + (sorted_indices[:, 1:] != sorted_indices[:, :-1]) + | (sorted_indices[:, 1:] == -1) + ) + selected = logits.gather(1, out.clamp_min(0).long()).masked_fill( + ~valid, -float("inf") + ) + reference = logits.masked_fill(invalid, -float("inf")).topk(k, dim=1).values + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, reference, atol=0, rtol=0 + ) + trials = [ + statistics.median( + bench_gpu_time_with_cupti( + run, + use_cuda_graph=True, + cold_l2_cache=not args.warm, + dry_run_iters=5, + repeat_iters=args.repeat, + ) + ) + * 1000 + for _ in range(args.trials) + ] + us = statistics.median(trials) + return dict( + rows=rows, + width=width, + capacity=capacity, + k=k, + backend="cooperative" if cooperative else "persistent", + us=us, + trials_us=trials, + # Minimum useful traffic: one valid-score read and one index write. + effective_gbps=(int(lengths.sum()) * 4 + out.numel() * 4) / (us * 1000), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 32, 128, 256, 1024]) + parser.add_argument("--widths", type=int, nargs="+", default=[8192, 65536, 163840]) + parser.add_argument("--top-k", type=int, nargs="+", default=[512, 2048]) + parser.add_argument("--capacity", type=int) + parser.add_argument( + "--full-length", action="store_true", help="Use exactly --widths valid scores" + ) + parser.add_argument("--backend", choices=["auto", "persistent"], default="auto") + parser.add_argument( + "--warm", action="store_true", help="Keep L2 warm between replays" + ) + parser.add_argument("--repeat", type=int, default=60) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--seed", type=int, default=123) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + metadata = dict( + gpu=str(torch.cuda.get_device_properties(0)), + torch=torch.__version__, + cuda=torch.version.cuda, + args=vars(args), + ) + print(json.dumps(metadata, default=str)) + results = [] + for rows, width, k in itertools.product(args.rows, args.widths, args.top_k): + result = benchmark(rows, width, k, args) + results.append(result) + print(json.dumps(result), flush=True) + if args.output: + args.output.write_text( + json.dumps(dict(metadata=metadata, results=results), indent=2, default=str) + + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 72f7f5eb66a8..2ff69bd8b6b3 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,11 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # Pinned to the tip of the nv_dev branch (SM120 support). - set(_DEEPGEMM_UPSTREAM_TAG "8b1392b978f5a03c828dd1711090d7fb50958b8a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") + # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 + # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU + # alpha/beta, and SiTU for FP8/FP4 Mega MoE. + set(_DEEPGEMM_UPSTREAM_TAG "ad1f1726aa540a76c1d26d6a120effb8de21eaa4") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -51,7 +53,7 @@ else() BINARY_DIR "${_deepgemm_bin}" GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}" GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}" - GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_SUBMODULES "third-party/cutlass" "third-party/deep_jit" GIT_PROGRESS TRUE ) endif() @@ -119,13 +121,15 @@ if(DEEPGEMM_ARCHS) message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}") # add_custom_command does no implicit header scanning; glob explicitly so - # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild. + # header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild. file(GLOB_RECURSE _dg_headers "${deepgemm_SOURCE_DIR}/csrc/*.h" "${deepgemm_SOURCE_DIR}/csrc/*.hpp" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp" - "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh") + "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp") set(_dg_markers) set(_dg_seen_soabis) diff --git a/cmake/external_projects/flashmla.cmake b/cmake/external_projects/flashmla.cmake index 382fd0289a2a..cfa8243f726c 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 6bc49418c5ead572ff0339191ddf3b155749e183 + GIT_TAG 0eee43b12f034b657133cf2afca6a72ebb6efccf GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" @@ -76,43 +76,77 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_SOURCES ${flashmla_SOURCE_DIR}/csrc/api/api.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_prefill.cpp + ${flashmla_SOURCE_DIR}/csrc/api/sparse_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/dense_decode.cpp + ${flashmla_SOURCE_DIR}/csrc/api/fused_norm_rope_attn_rope_cast_fwd.cpp # Misc kernels for decoding - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu - ${flashmla_SOURCE_DIR}/csrc/smxx/decode/combine/combine.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/smxx/decode/combine/combine.cu # sm90 dense decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/fp16.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/dense/instantiations/bf16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/fp16.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/dense/instantiations/bf16.cu # sm90 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/model1_persistent_h128.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h64.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/decode/sparse_fp8/instantiations/v32_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v4_persistent_h128.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/decode/sparse/instantiations/v32_persistent_h128.cu # sm90 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/fwd.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k512_topklen.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm90/prefill/sparse/instantiations/phase1_k576_topklen.cu - # sm100 dense prefill & backward - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu + # sm100 dense prefill (inference-only: dense backward stays out) + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu # sm100 sparse prefill - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_prefill_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head64/instantiations/phase1_h64_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd/head128/instantiations/phase1_k576.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_k512.cu + + # sm100 fused norm + rope + attn + rope + cast + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_prefill_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h64_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_norm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/core_attn/instantiations/v41fp4_h128_decode_nonorm.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_q_b_proj/kernel.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fused_norm_rope_attn_rope_cast_fwd/permute_wv_proj/kernel.cu # sm100 sparse decode - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/model1.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/decode/head64/instantiations/v32_nvfp4_fp8rope.cu - ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v32_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/head64/instantiations/v41fp4_h64_no_split.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/decode/sparse/nvfp4_head64/instantiations/v32_nvfp4_fp8rope.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41_splitkv.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4.cu + ${flashmla_SOURCE_DIR}/csrc/kernels/sm100/prefill/sparse/fwd_for_small_topk/head128/instantiations/phase1_decode_k512_v41fp4_splitkv.cu ) set(FlashMLA_Extension_SOURCES @@ -125,7 +159,6 @@ if(FLASH_MLA_ARCHS) set(FlashMLA_INCLUDES ${flashmla_SOURCE_DIR}/csrc ${flashmla_SOURCE_DIR}/csrc/kerutils/include - ${flashmla_SOURCE_DIR}/csrc/sm90 ${flashmla_SOURCE_DIR}/csrc/cutlass/include ${flashmla_SOURCE_DIR}/csrc/cutlass/tools/util/include ) diff --git a/cmake/external_projects/fmha_sm100.cmake b/cmake/external_projects/fmha_sm100.cmake index 442097dec7d3..ae5bfec7b47d 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 087c161814d4d9c735b46c21212a09e5f8eb92fa + GIT_TAG f355c37eb4e1413f21ee2ad8bbad25079e6bef9d GIT_PROGRESS TRUE CONFIGURE_COMMAND "" BUILD_COMMAND "" diff --git a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu index 06fe7108bd61..344b82fbda2c 100644 --- a/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu +++ b/csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu @@ -23,9 +23,15 @@ * cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin * second along last dim; each half is rope_dim/2 = 32 values) * - * Cache layout per paged-cache block (block_size tokens): - * [0, bs*576): token data, 448 fp8 + 128 bf16 each - * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * Cache layout per paged-cache block (block_size tokens), selected by the + * `kv_mxfp8` argument: + * V4 (584 B/token, kv_mxfp8=false) + * [0, bs*576): token data, 448 fp8 + 128 bf16 each + * [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token + * V4.1 (528 B/token, kv_mxfp8=true) -- RoPE is quantized as well, so the + * whole row is MXFP8 (fp8 e4m3 with one UE8M0 scale per 32 dims) + * [0, bs*512): token data, 512 fp8 each + * [bs*512, bs*512 + bs*16): UE8M0 scales, 16 per token */ #include "torch_utils.h" @@ -91,6 +97,10 @@ constexpr int kQuantBlock = 64; constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad) constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576 +// V4.1 MXFP8 record: every dim quantized, one UE8M0 scale per 32 dims. +constexpr int kMxfp8QuantBlock = 32; +constexpr int kMxfp8NumQuantBlocks = kHeadDim / kMxfp8QuantBlock; // 16 +constexpr int kMxfp8TokenDataBytes = kHeadDim; // 512 // FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw // 240.0) to match the rest of vLLM's FNUZ pipeline. #if defined(USE_ROCM) && defined(__gfx942__) @@ -141,6 +151,11 @@ __device__ __forceinline__ uint4 packFp8E4M3x16(float const* values, // ──────────────────────────────────────────────────────────────────────────── // Small inline helpers // ──────────────────────────────────────────────────────────────────────────── +__device__ __forceinline__ float warp2MaxAbs(float val) { + // Reduce absolute max across 2 consecutive lanes (32 dims = one MXFP8 block). + return fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 1)); +} + __device__ __forceinline__ float warp4MaxAbs(float val) { // Reduce absolute max across 4 consecutive lanes (lane id & 3 group). float peer = __shfl_xor_sync(FINAL_MASK, val, 1); @@ -176,6 +191,10 @@ __device__ __forceinline__ float warpSum(float val) { // slot == kNumHeadsQPadded → KV (RoPE + UE8M0 quant // + paged-cache // insert) +// +// `kv_mxfp8` is grid-uniform and only steers the KV branch, so it stays a +// runtime argument rather than doubling every (head count, q-norm) +// instantiation of a kernel whose hot path is the Q branch. template __device__ __forceinline__ void processDeepseekV4Slot( uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx, @@ -184,7 +203,7 @@ __device__ __forceinline__ void processDeepseekV4Slot( uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping, int64_t const* __restrict__ position_ids, float const* __restrict__ cos_sin_cache, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { using Converter = vllm::_typeConvert; bool const isKV = (slotIdx == kNumHeadsQPadded); bool const isPadQ = !isKV && (slotIdx >= num_heads_q); @@ -302,19 +321,24 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(dst) = out0; *reinterpret_cast(dst + 8) = out1; } else { - // ── KV: FP8 quant on NoPE + bf16 store on RoPE + cache insert. + // ── KV: FP8 quant + cache insert. V4 keeps the RoPE dims in bf16; + // V4.1 (kv_mxfp8) quantizes them too. int64_t const slot_id = slot_mapping[tokenIdx]; if (slot_id >= 0) { + int const token_data_bytes = + kv_mxfp8 ? kMxfp8TokenDataBytes : kTokenDataBytes; + int const scale_bytes = + kv_mxfp8 ? kMxfp8NumQuantBlocks : kScaleBytesPerToken; int64_t const block_idx = slot_id / cache_block_size; int64_t const pos_in_block = slot_id % cache_block_size; uint8_t* block_base = k_cache + block_idx * static_cast(kv_block_stride); - uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; + uint8_t* token_fp8_ptr = block_base + pos_in_block * token_data_bytes; uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim; uint8_t* token_scale_ptr = block_base + - static_cast(cache_block_size) * kTokenDataBytes + - pos_in_block * kScaleBytesPerToken; + static_cast(cache_block_size) * token_data_bytes + + pos_in_block * scale_bytes; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -326,11 +350,13 @@ __device__ __forceinline__ void processDeepseekV4Slot( for (int i = 0; i < kElemsPerLane; i++) { local_absmax = fmaxf(local_absmax, fabsf(elements[i])); } - float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f); + float const absmax = fmaxf( + kv_mxfp8 ? warp2MaxAbs(local_absmax) : warp4MaxAbs(local_absmax), + 1e-4f); float const exponent = ceilf(log2f(absmax / kFp8Max)); float const inv_scale = exp2f(-exponent); - if (!is_rope_lane) { + if (kv_mxfp8 || !is_rope_lane) { uint8_t out_bytes[kElemsPerLane]; #pragma unroll for (int i = 0; i < kElemsPerLane; i++) { @@ -347,13 +373,19 @@ __device__ __forceinline__ void processDeepseekV4Slot( *reinterpret_cast(token_fp8_ptr + dim_base) = *reinterpret_cast(out_bytes); - if ((laneId & 3) == 0) { - int const q_block_idx = laneId >> 2; - float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); - token_scale_ptr[q_block_idx] = static_cast(encoded); - } - if (laneId == 0) { - token_scale_ptr[kNumQuantBlocks] = 0; + float const encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f); + if (kv_mxfp8) { + // One scale per 32 dims = one per lane pair. + if ((laneId & 1) == 0) { + token_scale_ptr[laneId >> 1] = static_cast(encoded); + } + } else { + if ((laneId & 3) == 0) { + token_scale_ptr[laneId >> 2] = static_cast(encoded); + } + if (laneId == 0) { + token_scale_ptr[kNumQuantBlocks] = 0; + } } } else { uint4 out0, out1; @@ -421,7 +453,8 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full int const num_heads_q, // live Q heads (input layout) int const cache_block_size, // tokens per paged-cache block - int const kv_block_stride) { // bytes per paged-cache block + int const kv_block_stride, // bytes per paged-cache block + bool const kv_mxfp8) { // V4.1 all-dims MXFP8 record #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) // BF16 _typeConvert specialization is unavailable on pre-Ampere. The // DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a @@ -477,7 +510,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel( processDeepseekV4Slot( v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaTriggerProgrammaticLaunchCompletion(); @@ -506,7 +539,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( float const* __restrict__ cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const cache_block_size, - int const kv_block_stride) { + int const kv_block_stride, bool const kv_mxfp8) { #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) if constexpr (std::is_same_v) { return; @@ -564,7 +597,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid( processDeepseekV4Slot( v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId, num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids, - cos_sin_cache, cache_block_size, kv_block_stride); + cos_sin_cache, cache_block_size, kv_block_stride, kv_mxfp8); // ── Buffer rotation: hand the prefetched LDGs to the next iter. v0_curr = v0_next; @@ -590,7 +623,7 @@ static void launchFusedDeepseekV4Templated( uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids, float const* cos_sin_cache, float const eps, int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, - int const cache_block_size, int const kv_block_stride, + int const cache_block_size, int const kv_block_stride, bool const kv_mxfp8, cudaStream_t stream) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; @@ -624,13 +657,19 @@ static void launchFusedDeepseekV4Templated( config.attrs = attrs; config.numAttrs = (sm_version >= 90) ? 1 : 0; - if (num_tokens_full < NUM_TOKEN_CUTOFF) { + // The reduced grid pays off by iterating a token's Q heads inside one CTA. + // With zero query heads there is nothing to iterate, so one warp per block + // would do all the work and the other seven idle; keep the warp-per-slot + // grid instead. Only reachable above NUM_TOKEN_CUTOFF tokens in a single + // insert, where it is worth ~2x. + if (kNumHeadsQPadded == 0 || num_tokens_full < NUM_TOKEN_CUTOFF) { cudaLaunchKernelEx(&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel< scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, - num_heads_q, cache_block_size, kv_block_stride); + num_heads_q, cache_block_size, kv_block_stride, + kv_mxfp8); } else { config.gridDim = dim3(num_tokens_full); cudaLaunchKernelEx( @@ -639,7 +678,7 @@ static void launchFusedDeepseekV4Templated( scalar_t_in, kNumHeadsQPadded, APPLY_Q_NORM>, q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size, - kv_block_stride); + kv_block_stride, kv_mxfp8); } #else // ROCm: use standard kernel launch syntax (no PDL/stream serialization) @@ -649,7 +688,7 @@ static void launchFusedDeepseekV4Templated( <<>>( q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, - cache_block_size, kv_block_stride); + cache_block_size, kv_block_stride, kv_mxfp8); #endif } @@ -663,19 +702,19 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert( int const num_tokens_full, int const num_tokens_insert, int const num_heads_q, int const num_heads_q_padded, int const cache_block_size, int const kv_block_stride, - bool const apply_q_norm, cudaStream_t stream) { + bool const apply_q_norm, bool const kv_mxfp8, cudaStream_t stream) { #define DISPATCH(N) \ case N: \ if (apply_q_norm) { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } else { \ launchFusedDeepseekV4Templated( \ q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, \ cos_sin_cache, eps, num_tokens_full, num_tokens_insert, \ - num_heads_q, cache_block_size, kv_block_stride, stream); \ + num_heads_q, cache_block_size, kv_block_stride, kv_mxfp8, stream); \ } \ return; @@ -959,6 +998,86 @@ static void launchFullCacheKernel( // ──────────────────────────────────────────────────────────────────────────── // Torch op wrapper // ──────────────────────────────────────────────────────────────────────────── +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8) { + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous() && + kv.dim() == 2 && kv.size(1) == 512 && + kv.scalar_type() == ScalarType::BFloat16, + "kv must be contiguous CUDA bfloat16 [tokens, 512]"); + auto const device = kv.get_device_index(); + STD_TORCH_CHECK(k_cache.device().is_cuda() && + k_cache.get_device_index() == device && + k_cache.dim() == 3 && cache_block_size > 0 && + k_cache.size(1) == cache_block_size && + k_cache.stride(2) == 1, + "k_cache must be a paged cache on the KV device"); + for (auto const* tensor : {&slot_mapping, &position_ids}) { + STD_TORCH_CHECK(tensor->device().is_cuda() && + tensor->get_device_index() == device && + tensor->scalar_type() == ScalarType::Long && + tensor->dim() == 1 && tensor->is_contiguous(), + "slots and positions must be contiguous int64 on the KV device"); + } + STD_TORCH_CHECK(position_ids.size(0) == kv.size(0) && + slot_mapping.size(0) <= kv.size(0), + "positions must match KV rows; slots may omit padded rows"); + STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() && + cos_sin_cache.get_device_index() == device && + cos_sin_cache.scalar_type() == ScalarType::Float && + cos_sin_cache.dim() == 2 && + cos_sin_cache.size(1) == 64 && cos_sin_cache.is_contiguous(), + "cos_sin_cache must be contiguous float32 [max_pos, 64]"); + bool const packed = k_cache.scalar_type() == ScalarType::Byte; + bool const fp8 = k_cache.scalar_type() == ScalarType::Float8_e4m3fn; + STD_TORCH_CHECK(packed || fp8 || + k_cache.scalar_type() == ScalarType::BFloat16, + "cache must be uint8, bfloat16 or float8_e4m3fn"); + STD_TORCH_CHECK(k_cache.size(2) == (packed ? (kv_mxfp8 ? 528 : 584) : 512), + "cache row width does not match its dtype"); + if (fp8) { + STD_TORCH_CHECK(fp8_scale.has_value() && fp8_scale->device().is_cuda() && + fp8_scale->get_device_index() == device && + fp8_scale->scalar_type() == ScalarType::Float && + fp8_scale->numel() == 1, + "fp8 cache requires one float32 KV scale on the KV device"); + } + int const num_tokens = static_cast(slot_mapping.size(0)); + if (num_tokens == 0) return; + const torch::stable::accelerator::DeviceGuard device_guard(device); + const cudaStream_t stream = get_current_cuda_stream(device); + using scalar_t = c10::BFloat16; + auto const* input = reinterpret_cast(kv.const_data_ptr()); + auto* cache = reinterpret_cast(k_cache.mutable_data_ptr()); + auto const* slots = slot_mapping.const_data_ptr(); + auto const* positions = position_ids.const_data_ptr(); + auto const* cos_sin = cos_sin_cache.const_data_ptr(); + // Zero query heads schedules only the existing KV branch, preserving its + // RoPE rounding and quantization without allocating any query tensors. + if (packed) { + vllm::deepseek_v4_fused_ops::launchFusedDeepseekV4Templated( + nullptr, nullptr, input, cache, slots, positions, cos_sin, 0.0f, + num_tokens, num_tokens, 0, static_cast(cache_block_size), + static_cast(k_cache.stride(0)), kv_mxfp8, stream); + } else if (fp8) { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, + fp8_scale->const_data_ptr(), nullptr, 0.0f, num_tokens, num_tokens, + 0, static_cast(cache_block_size), k_cache.stride(0), + k_cache.stride(1), false, "fused_deepseek_v4_kv_rope_insert", stream); + } else { + vllm::deepseek_v4_fused_ops::launchFullCacheKernel( + nullptr, nullptr, 0, 0, input, cache, slots, positions, cos_sin, nullptr, + nullptr, 0.0f, num_tokens, num_tokens, 0, static_cast(cache_block_size), + k_cache.stride(0) * 2, k_cache.stride(1) * 2, false, + "fused_deepseek_v4_kv_rope_insert", stream); + } +} + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only) @@ -967,7 +1086,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& position_ids, // [N] int64 torch::stable::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16 int64_t q_head_padded, // padded Q head count for output - double eps, int64_t cache_block_size, bool apply_q_norm) { + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8) { STD_TORCH_CHECK(q_in.device().is_cuda() && q_in.is_contiguous(), "q_in must be contiguous CUDA"); STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(), @@ -1036,7 +1155,7 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( cos_sin_cache.const_data_ptr(), static_cast(eps), num_tokens_full, num_tokens_insert, num_heads_q, num_heads_q_padded, cache_block_size_i, kv_block_stride, - apply_q_norm, stream); + apply_q_norm, kv_mxfp8, stream); }); return q_out; } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 68a24ca4340a..bc08d90a4836 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -259,12 +259,19 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, torch::stable::Tensor& position_ids, int64_t forced_token_heads_per_warp); +void fused_deepseek_v4_kv_rope_insert( + torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, + torch::stable::Tensor const& slot_mapping, + torch::stable::Tensor const& position_ids, + torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size, + std::optional fp8_scale, bool kv_mxfp8); + torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping, torch::stable::Tensor const& position_ids, torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded, - double eps, int64_t cache_block_size, bool apply_q_norm); + double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8); void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( torch::stable::Tensor& q, torch::stable::Tensor const& kv, diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 9ae0fd4bcbaa..bf5b46ab1875 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -1030,41 +1030,34 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB -/*! - * \brief Filtered Top-K kernel for ragged sequences. - * - * \tparam DType Data type (float, half, nv_bfloat16) - * \tparam IdType Index type (int32_t) - * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) - */ -template -__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) - FilteredTopKUnifiedKernel(const DType* __restrict__ input, - IdType* __restrict__ output, - const IdType* __restrict__ lengths, - uint32_t num_rows, uint32_t top_k, - uint32_t max_len) { +template +struct FilteredTopKStorage { + alignas(128) int histogram[2][256 + 128]; + alignas(128) int counter; + alignas(128) int threshold_bin; + alignas(128) int num_input[2]; + alignas(128) int indices[MAX_K]; + int last_remain; +}; + +// With CheckOverflow, return false before using a truncated stash so the caller +// can retry with exact full-row selection. +template +__device__ __forceinline__ bool filtered_topk_row( + const DType* score, IdType* dst, int length, uint32_t top_k, + FilteredTopKStorage& storage) { constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; constexpr int RADIX = 256; constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; - - const uint32_t bid = blockIdx.x; const int tx = threadIdx.x; - if (bid >= num_rows) return; - - const int length = - (lengths != nullptr) ? lengths[bid] : static_cast(max_len); - const DType* score = input + bid * max_len; - IdType* dst = output + bid * top_k; - // Trivial case: length <= top_k if (length <= static_cast(top_k)) { for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { dst[i] = (i < length) ? static_cast(i) : static_cast(-1); } - return; + return true; } // Short path @@ -1077,16 +1070,14 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) hist4096::histogram_4096_topk(score, dst, length, _smem_reg); } - return; + return true; } - // Static shared memory - alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; - alignas(128) __shared__ int s_counter; - alignas(128) __shared__ int s_threshold_bin_id; - alignas(128) __shared__ int s_num_input[2]; - alignas(128) __shared__ int s_indices[MAX_K]; - + auto& s_histogram_buf = storage.histogram; + auto& s_counter = storage.counter; + 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 @@ -1213,10 +1204,13 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) // Stage 2: refine with 8bit radix passes #pragma unroll for (int round = 0; round < NUM_ROUNDS; ++round) { - __shared__ int s_last_remain; + auto& s_last_remain = storage.last_remain; const auto r_idx = round % 2; const auto _raw_num_input = s_num_input[r_idx]; + if constexpr (CheckOverflow) { + if (_raw_num_input > SMEM_INPUT_SIZE) return false; + } const auto num_input = (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; @@ -1284,6 +1278,30 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const int idx = s_indices[base]; dst[base] = static_cast(idx); } + return true; +} + +/*! + * \brief Filtered Top-K kernel for ragged sequences. + * + * \tparam DType Data type (float, half, nv_bfloat16) + * \tparam IdType Index type (int32_t) + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) + */ +template +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) + FilteredTopKUnifiedKernel(const DType* __restrict__ input, + IdType* __restrict__ output, + const IdType* __restrict__ lengths, + uint32_t num_rows, uint32_t top_k, + uint32_t max_len) { + 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); } // Helper to compute GCD for VEC_SIZE selection diff --git a/csrc/libtorch_stable/sampled_topk.cuh b/csrc/libtorch_stable/sampled_topk.cuh new file mode 100644 index 000000000000..867abbd63542 --- /dev/null +++ b/csrc/libtorch_stable/sampled_topk.cuh @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +// Sampled top-k for long FP32 sparse-indexer rows, inspired by DeepSelect: +// https://github.com/deepseek-ai/DeepSelect +// A sample estimates a cutoff, then survivors are compacted in shared memory +// for exact FP32 selection. Too few or too many candidates trigger an exact +// full-row fallback. + +#pragma once + +#include +#include +#include + +#include "persistent_topk.cuh" + +namespace vllm::sampled_topk { + +constexpr int kThreads = 1024; +// Conservative crossover bounds from the B300 batch/length sweep. +template +constexpr int kMinSampledLength = K == 512 ? 98304 : 65536; +// Target half the buffer to leave room for sampling error. +constexpr int kSample = 4096; +constexpr int kCapacity = 8192; + +struct Storage { + uint2 candidates[kCapacity]; + int histogram[2048]; + typename cub::BlockScan::TempStorage scan; + int count; + int remaining; + int bin; + uint32_t prefix; +}; + +__device__ __forceinline__ uint32_t ordered(float value) { + const uint32_t bits = __float_as_uint(value); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ uint32_t coarse(uint32_t key) { + const uint32_t bits = (key & 0x80000000u) ? (key ^ 0x80000000u) : ~key; + return topk_histogram_4096::extract_coarse_bin_N<11>(__uint_as_float(bits)); +} + +// Select the bin containing the rank and update the rank within that bin. +template +__device__ void select_bin(Storage& s) { + constexpr int kItems = (Bins + kThreads - 1) / kThreads; + int sums[kItems]; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + sums[j] = bin >= 0 ? s.histogram[bin] : 0; + } + cub::BlockScan(s.scan).InclusiveSum(sums, sums); + const int remaining = s.remaining; + __syncthreads(); +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int bin = Bins - 1 - (threadIdx.x * kItems + j); + if (bin >= 0 && sums[j] >= remaining && + sums[j] - s.histogram[bin] < remaining) { + s.bin = bin; + s.remaining = remaining - (sums[j] - s.histogram[bin]); + } + } + __syncthreads(); +} + +// 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) { + const int tid = threadIdx.x; + if (tid == 0) { + s.remaining = rank; + s.prefix = 0; + } + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + atomicAdd(&s.histogram[coarse(key)], 1); + } + __syncthreads(); + select_bin<2048>(s); + const int coarse_bin = s.bin; + uint32_t mask = 0; + for (int shift = 24; shift >= 0; shift -= 8) { + if (tid < 256) s.histogram[tid] = 0; + __syncthreads(); + const uint32_t prefix = s.prefix; + for (int i = tid; i < count; i += kThreads) { + const uint32_t key = Buffered ? s.candidates[i].y : ordered(row[i]); + if (coarse(key) == coarse_bin && (key & mask) == prefix) { + atomicAdd(&s.histogram[(key >> shift) & 255], 1); + } + } + __syncthreads(); + select_bin<256>(s); + if (tid == 0) s.prefix = prefix | (static_cast(s.bin) << shift); + mask |= 255u << shift; + __syncthreads(); + } + return s.prefix; +} + +__device__ __forceinline__ int reserve(bool hit, int* counter) { + const uint32_t mask = __ballot_sync(0xffffffffu, hit); + if (mask == 0) return 0; + const int lane = threadIdx.x & 31; + int base = 0; + if (lane == 0) base = atomicAdd(counter, __popc(mask)); + base = __shfl_sync(0xffffffffu, base, 0); + return base + __popc(mask & ((1u << lane) - 1)); +} + +template +__device__ void emit(Storage& s, const float* row, int32_t* dst, int count, + uint32_t cutoff) { + const int tid = threadIdx.x; + if (tid == 0) s.count = 0; + __syncthreads(); + 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, ordered(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; + } + const int offset = reserve(keep, &s.count); + if (keep) dst[offset] = pair.x; + } +} + +template +__global__ void __launch_bounds__(kThreads) + sampled_topk_kernel(const float* __restrict__ input, + const int32_t* __restrict__ lengths, + int32_t* __restrict__ output, int64_t stride, + int max_length) { + extern __shared__ __align__(16) unsigned char smem[]; + 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; + int32_t* dst = output + blockIdx.x * K; + // Graph allocations can hold rows much shorter than the host length bound. + if (length < kMinSampledLength) { + __shared__ filtered_topk::FilteredTopKStorage fallback_storage; + bool selected; + if ((stride & 3) == 0 && (reinterpret_cast(row) & 15) == 0) { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } else { + selected = + filtered_topk::filtered_topk_row( + row, dst, length, K, fallback_storage); + } + if (!selected) { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } + return; + } + if (tid == 0) s.remaining = max(1, (kCapacity / 2) * kSample / length); + for (int i = tid; i < 2048; i += kThreads) s.histogram[i] = 0; + __syncthreads(); + // 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(ordered(row[index]))], 1); + } + __syncthreads(); + select_bin<2048>(s); + const uint16_t half_key = s.bin << 5; + const uint16_t half_bits = + (half_key & 0x8000u) ? (half_key ^ 0x8000u) : ~half_key; + const uint32_t cutoff = ordered(__half2float(__ushort_as_half(half_bits))); + if (tid == 0) s.count = 0; + __syncthreads(); + constexpr int kItems = 16; + // Compact survivors with one reservation per warp per tile. Keep counting + // beyond capacity so overflow triggers fallback instead of using a subset. + for (int base = 0; base < length; base += kItems * kThreads) { + uint32_t keys[kItems]; + uint32_t hit_mask = 0; + if ((reinterpret_cast(row) & 15) == 0 && + base + kItems * kThreads <= length) { +#pragma unroll + for (int j = 0; j < kItems / 4; ++j) { + const float4 values = + reinterpret_cast(row + base)[tid + j * kThreads]; + keys[4 * j] = ordered(values.x); + keys[4 * j + 1] = ordered(values.y); + keys[4 * j + 2] = ordered(values.z); + keys[4 * j + 3] = ordered(values.w); + } + } else { +#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 ? ordered(row[i]) : 0; + } + } +#pragma unroll + for (int j = 0; j < kItems; ++j) { + const int i = base + tid * 4 + (j / 4) * kThreads * 4 + j % 4; + if (i < length && keys[j] >= cutoff) hit_mask |= 1u << j; + } + const int lane = tid & 31; + const int hits = __popc(hit_mask); + int inclusive = hits; +#pragma unroll + for (int delta = 1; delta < 32; delta *= 2) { + const int previous = __shfl_up_sync(0xffffffffu, inclusive, delta); + if (lane >= delta) inclusive += previous; + } + const int total = __shfl_sync(0xffffffffu, inclusive, 31); + int warp_base = 0; + if (lane == 31 && total != 0) warp_base = atomicAdd(&s.count, total); + warp_base = __shfl_sync(0xffffffffu, warp_base, 31); + int offset = warp_base + inclusive - hits; +#pragma unroll + for (int j = 0; j < kItems; ++j) { + if (hit_mask & (1u << j)) { + if (offset < kCapacity) { + s.candidates[offset] = make_uint2( + base + tid * 4 + (j / 4) * kThreads * 4 + j % 4, keys[j]); + } + ++offset; + } + } + } + __syncthreads(); + const int count = s.count; + if (count >= K && count <= kCapacity) { + const uint32_t exact_cutoff = threshold(s, row, count, K); + emit(s, row, dst, count, exact_cutoff); + } else { + const uint32_t exact_cutoff = threshold(s, row, length, K); + emit(s, row, dst, length, exact_cutoff); + } +} + +} // namespace vllm::sampled_topk diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index ada1c4ba39c0..b8d64f7af130 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -154,7 +154,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads, } template __device__ bool processHistogramStep( const int* indices, const float* logits, int rowEnd, uint32_t& logitPattern, @@ -180,10 +181,15 @@ __device__ bool processHistogramStep( << patternShift; } + int negativeInfinityCount = 0; auto distributeToBins = [&](float logit, int /* idx */ = 0) { if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - atomicAdd(&smemFinal.histo.data[binIdx], 1); + if (useTopK512Optimization && __float_as_uint(logit) == 0xff800000u) { + ++negativeInfinityCount; + } else { + atomicAdd(&smemFinal.histo.data[binIdx], 1); + } } }; @@ -198,6 +204,12 @@ __device__ bool processHistogramStep( distributeToBins(logit, idx); } } + if constexpr (useTopK512Optimization) { + if (negativeInfinityCount > 0) { + const auto bin = extractBinIdx(-INFINITY); + atomicAdd(&smemFinal.histo.data[bin], negativeInfinityCount); + } + } // Make sure the histogram is ready. __syncthreads(); @@ -254,6 +266,12 @@ __device__ bool processHistogramStep( // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; + // Resolve a -inf cutoff with exact tie emission, keeping sort padding out. + const bool finalBinFits = + smemFinalBinSize[0] <= kNumFinalItems && + !(useTopK512Optimization && step < 3 && + isPartialMatch(-INFINITY, logitPattern) && + thresholdBinIdx == extractBinIdx(-INFINITY)); auto processBins = [&](float logit, int idx) { if (isPartialMatch(logit, logitPattern)) { @@ -262,8 +280,7 @@ __device__ bool processHistogramStep( // 1. This is step 0 and the threshold bin is small enough (no step 1) // 2. This is step >= 1 (where pattern matching filters correctly) // This prevents duplicates when step 0 and step 1 both run. - bool shouldWriteDirectly = - (step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1); + bool shouldWriteDirectly = (step == 0 && finalBinFits) || (step >= 1); if (binIdx < thresholdBinIdx && shouldWriteDirectly) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -279,8 +296,7 @@ __device__ bool processHistogramStep( } if constexpr (step < 3) { // Only fill the final items for sorting if the threshold bin fits - if (binIdx == thresholdBinIdx && - smemFinalBinSize[0] <= kNumFinalItems) { + if (binIdx == thresholdBinIdx && finalBinFits) { int dstIdx = atomicAdd(&smemFinalDstIdx[0], 1); smemFinal.items.logits[dstIdx] = logit; if constexpr (mergeBlocks) { @@ -325,20 +341,21 @@ __device__ bool processHistogramStep( __syncthreads(); // Check if we should continue to next step - return smemFinalBinSize[0] > kNumFinalItems; + return !finalBinFits; } // Follows half - 11 - 11 - 10 bit iterations // Keep the adaptive kernel's device instantiation separate from legacy calls. template + bool deviceLengthAware = false, bool useTopK512Optimization = false> static __device__ void topKPerRowJob(const int* indices, const float* logits, int rowStart, int rowEnd, int* outIndices, float* outLogits, int stride1, int topK) { static_assert(!deviceLengthAware || multipleBlocksPerRow != mergeBlocks); + static_assert(!useTopK512Optimization || kNumThreadsPerBlock == 1024); // The number of slots for the final pass. - static constexpr int kNumFinalItems = 2048; + static constexpr int kNumFinalItems = useTopK512Optimization ? 1024 : 2048; // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; @@ -404,7 +421,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, rowIt += kNumThreadsPerBlock) { outIndices[rowIt] = -1; if constexpr (multipleBlocksPerRow) { - outLogits[rowIt] = -FLT_MAX; + outLogits[rowIt] = useTopK512Optimization ? -INFINITY : -FLT_MAX; } } @@ -422,7 +439,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 0: Process first 11 bits of half representation bool continueToNextStep = processHistogramStep<0, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -431,7 +449,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 1: Process next 11 bits continueToNextStep = processHistogramStep<1, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -441,7 +460,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // Step 2: Process next 11 bits continueToNextStep = processHistogramStep<2, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -450,7 +470,8 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, if (continueToNextStep) { // Step 3: Process last 10 bits processHistogramStep<3, kNumThreadsPerBlock, kNumBins, kNumFinalItems, - multipleBlocksPerRow, mergeBlocks>( + multipleBlocksPerRow, mergeBlocks, + useTopK512Optimization>( indices, logits, rowEnd, logitPattern, thresholdBinIdx, smemOutput, smemThresholdBinIdx, smemFinalDstIdx, smemFinalBinSize, smemFoundTopKValues, smemFinal, stride1, rowStart, topK); @@ -460,50 +481,7 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // The histogram did not proceed to the final 10 bits, therefore we need to // sort the final items The logits of the elements to be sorted in the final // pass. - if constexpr (useRadixSort) { - // Sorting with radix sort - float finalLogits[kNumFinalItemsPerThread]; - // The indices of the elements to be sorted in the final pass. - int finalIndices[kNumFinalItemsPerThread]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - finalLogits[ii] = -FLT_MAX; - } - - // Read the elements from SMEM. -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < smemFinalDstIdx[0]) { - finalLogits[ii] = smemFinal.items.logits[srcIdx]; - finalIndices[ii] = smemFinal.items.indices[srcIdx]; - } - } - // Make sure the shared memory has been read. - __syncthreads(); - - // Sort the elements. - FinalSort(smemFinal.finalSort) - .SortDescendingBlockedToStriped(finalLogits, finalIndices); - - // Copy the data back to the shared memory storage. - int baseIdx = smemFoundTopKValues[0]; - -#pragma unroll - for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { - int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - int dstIdx = baseIdx + srcIdx; - - if (dstIdx < topK) { - smemOutput[dstIdx] = finalIndices[ii]; - if constexpr (multipleBlocksPerRow) { - reinterpret_cast(smemOutput + topK)[dstIdx] = - finalLogits[ii]; - } - } - } - } else { + auto insertionSort = [&]() { // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; for (int i = threadIdx.x; i < smemFinalDstIdx[0]; @@ -525,6 +503,57 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, } } } + }; + if constexpr (useRadixSort) { + if (useTopK512Optimization && smemFinalDstIdx[0] <= 128) { + insertionSort(); + } else { + // Sorting with radix sort + float finalLogits[kNumFinalItemsPerThread]; + // The indices of the elements to be sorted in the final pass. + int finalIndices[kNumFinalItemsPerThread]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + finalLogits[ii] = useTopK512Optimization ? -INFINITY : -FLT_MAX; + if constexpr (useTopK512Optimization) finalIndices[ii] = -1; + } + + // Read the elements from SMEM. +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + if (srcIdx < smemFinalDstIdx[0]) { + finalLogits[ii] = smemFinal.items.logits[srcIdx]; + finalIndices[ii] = smemFinal.items.indices[srcIdx]; + } + } + // Make sure the shared memory has been read. + __syncthreads(); + + // Sort the elements. + FinalSort(smemFinal.finalSort) + .SortDescendingBlockedToStriped(finalLogits, finalIndices); + + // Copy the data back to the shared memory storage. + int baseIdx = smemFoundTopKValues[0]; + +#pragma unroll + for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { + int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; + int dstIdx = baseIdx + srcIdx; + + if (dstIdx < topK) { + smemOutput[dstIdx] = finalIndices[ii]; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[dstIdx] = + finalLogits[ii]; + } + } + } + } + } else { + insertionSort(); } __syncthreads(); } @@ -669,6 +698,70 @@ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecodeDeviceLengthAware( stride1, topK); } +#ifdef USE_ROCM +// The grid and workspace stay fixed while graph replays change device lengths. +__device__ int gfx950TopK512ActiveBlocks(int rowLength, int numRows, + int maxBlocks) { + int activeBlocks = 1; + if (numRows <= 32 && rowLength > 16 * 1024) { + activeBlocks = rowLength <= 128 * 1024 ? 8 + : rowLength <= 256 * 1024 ? 10 + : 16; + } else if (numRows <= 128 && rowLength > 64 * 1024) { + activeBlocks = rowLength <= 512 * 1024 ? 4 : 10; + } else if (numRows > 128 && numRows <= 256 && rowLength > 256 * 1024) { + activeBlocks = numRows <= 160 ? 3 : 2; + } + // Every block in a merged row must contain at least k real indices. + return min(min(activeBlocks, maxBlocks), max(1, rowLength / 512)); +} + +template +static __global__ +__launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode512DeviceLengthAware( + const float* logits, const int* seqLens, int* outIndices, + int* outIndicesAux, float* outLogitsAux, int stride0, int next_n, + int seqLensIs2D, int maxBlocks) { + #if defined(__gfx950__) + constexpr int topK = 512; + constexpr int kNumBins = 2048; + const int rowIdx = blockIdx.x; + const int rowEnd = + seqLensIs2D + ? max(0, seqLens[rowIdx]) + : max(0, seqLens[rowIdx / next_n] - next_n + rowIdx % next_n + 1); + const int activeBlocks = + gfx950TopK512ActiveBlocks(rowEnd, gridDim.x, maxBlocks); + outIndices += static_cast(rowIdx) * topK; + if constexpr (mergeBlocks) { + if (activeBlocks == 1) return; + const int64_t offset = static_cast(rowIdx) * maxBlocks * topK; + topKPerRowJob(outIndicesAux + offset, outLogitsAux + offset, 0, + activeBlocks * topK, outIndices, nullptr, 1, topK); + } else { + if (blockIdx.y >= activeBlocks) return; + logits += static_cast(rowIdx) * stride0; + if (activeBlocks == 1) { + topKPerRowJob(nullptr, logits, 0, rowEnd, outIndices, nullptr, 1, + topK); + return; + } + const int blockSize = rowEnd / activeBlocks; + const int rowStart = blockSize * blockIdx.y; + const int blockRowEnd = + blockIdx.y + 1 == activeBlocks ? rowEnd : rowStart + blockSize; + const int64_t offset = + (static_cast(rowIdx) * maxBlocks + blockIdx.y) * topK; + topKPerRowJob(nullptr, logits, rowStart, blockRowEnd, + outIndicesAux + offset, outLogitsAux + offset, 1, topK); + } + #endif +} +#endif + } // namespace vllm void apply_repetition_penalties_( @@ -731,6 +824,60 @@ void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, // the same seq_len and the kernel computes the per-row offset itself. int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0; +#ifdef USE_ROCM + const bool useGfx950DeviceLengthAwareTopK512 = + topK == 512 && stride1 == 1 && numRows > 0 && numRows <= 384 && + numColumns <= 1024 * 1024 && + std::strncmp(get_device_prop()->gcnArchName, "gfx950", 6) == 0; + if (useGfx950DeviceLengthAwareTopK512) { + int multipleBlocksPerRowConfig = 1; + if (numRows <= 32 && numColumns > 16 * 1024) { + multipleBlocksPerRowConfig = 16; + } else if (numRows <= 128 && numColumns > 64 * 1024) { + multipleBlocksPerRowConfig = 10; + } else if (numRows <= 160 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 3; + } else if (numRows <= 256 && numColumns > 256 * 1024) { + multipleBlocksPerRowConfig = 2; + } + + constexpr int kNumThreadsPerSplitBlock = 1024; + if (multipleBlocksPerRowConfig == 1) { + vllm::topKPerRowDecode512DeviceLengthAware + <<>>(logits.const_data_ptr(), + seqLens.const_data_ptr(), + indices.mutable_data_ptr(), nullptr, nullptr, + static_cast(stride0), static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } + const auto outIndicesAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Int, std::nullopt, logits.device()); + const auto outLogitsAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Float, std::nullopt, logits.device()); + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + logits.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), static_cast(stride0), + static_cast(next_n), seqLensIs2D, multipleBlocksPerRowConfig); + constexpr int kNumThreadsPerBlockMerge = 1024; + vllm::topKPerRowDecode512DeviceLengthAware + <<>>( + nullptr, seqLens.const_data_ptr(), + indices.mutable_data_ptr(), + outIndicesAux.mutable_data_ptr(), + outLogitsAux.mutable_data_ptr(), 0, static_cast(next_n), + seqLensIs2D, multipleBlocksPerRowConfig); + return; + } +#endif + if (numColumns < kSortingAlgorithmThreshold) { // Use insertion sort vllm::topKPerRowDecode diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index b3e1a8623f3e..fb3d946b96b0 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -9,6 +9,7 @@ #ifndef USE_ROCM #include "persistent_topk.cuh" + #include "sampled_topk.cuh" #endif namespace { @@ -36,7 +37,22 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, max_smem_per_block = device_prop->sharedMemPerBlockOptin; } - if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + // 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; + constexpr size_t smem_size = + vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC; + cudaError_t status = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + STD_TORCH_CHECK(status == cudaSuccess, + "sampled_topk smem failed: ", cudaGetErrorString(status)); + kernel<<>>( + logits.const_data_ptr(), 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(), diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 110f8595ee99..3918ab8498cd 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -425,12 +425,18 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + ops.def( + "fused_deepseek_v4_kv_rope_insert(" + "Tensor kv, Tensor! k_cache, Tensor slot_mapping, Tensor position_ids, " + "Tensor cos_sin_cache, int cache_block_size, Tensor? fp8_scale=None, " + "bool kv_mxfp8=False) -> " + "()"); ops.def( "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(" "Tensor q_in, Tensor kv, Tensor! k_cache, " "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, " "int q_head_padded, float eps, int cache_block_size, " - "bool apply_q_norm=True) -> Tensor"); + "bool apply_q_norm=True, bool kv_mxfp8=False) -> Tensor"); // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate // FP8 tensor, and KV into a contiguous 512-wide token-strided cache. @@ -807,6 +813,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // Positional encoding kernels (shared CUDA/ROCm) ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + ops.impl("fused_deepseek_v4_kv_rope_insert", + TORCH_BOX(&fused_deepseek_v4_kv_rope_insert)); ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert)); ops.impl( diff --git a/docker/Dockerfile b/docker/Dockerfile index 9b344552c5ae..d1a7afcbfb70 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -131,6 +131,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ git \ curl \ sudo \ + elfutils-devel \ rdma-core-devel \ && dnf clean all \ && rm -rf /var/cache/dnf @@ -640,7 +641,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ python3 use_existing_torch.py --prefix; \ fi \ && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\])[[:space:]]*==/Id' \ + sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \ requirements/cuda.txt; \ fi diff --git a/docker/Dockerfile.rock b/docker/Dockerfile.rock index c3ce785c679d..14ba523b072f 100644 --- a/docker/Dockerfile.rock +++ b/docker/Dockerfile.rock @@ -875,6 +875,15 @@ 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 diff --git a/docker/Dockerfile.rock_base b/docker/Dockerfile.rock_base index 1b1d08988123..80b6aaca2964 100644 --- a/docker/Dockerfile.rock_base +++ b/docker/Dockerfile.rock_base @@ -146,6 +146,15 @@ RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ ln -sfn "${ROCM_PATH}/lib/${bn}" "$f"; \ fi; \ done \ + && for f in librocprofiler-sdk.so.1 librocprofiler-register.so.0 \ + librocprofiler-sdk-roctx.so.1; do \ + if [ -e "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}" ] \ + && [ -e "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" ]; then \ + rm -f "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + ln -sfn "${SITE_PACKAGES}/_rocm_sdk_devel/lib/${f}" \ + "${SITE_PACKAGES}/_rocm_sdk_core/lib/${f}"; \ + fi; \ + done \ && ldconfig # Install sccache if USE_SCCACHE is enabled (for release builds) diff --git a/docker/docker-bake-rocm.hcl b/docker/docker-bake-rocm.hcl index 9d73a6505eab..ab523a296a93 100644 --- a/docker/docker-bake-rocm.hcl +++ b/docker/docker-bake-rocm.hcl @@ -33,6 +33,10 @@ 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" { @@ -74,7 +78,7 @@ group "default" { } target "_common-rocm" { - dockerfile = "docker/Dockerfile.rocm" + dockerfile = CI_BASE_DOCKERFILE context = "." args = { max_jobs = MAX_JOBS diff --git a/docs/design/metrics.md b/docs/design/metrics.md index c7bc85756e0d..3d108890d702 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -33,7 +33,8 @@ In v1, an extensive set of metrics are exposed via a Prometheus-compatible `/met - `vllm:request_prompt_tokens` (Histogram) - Histogram of input prompt token counts. - `vllm:request_generation_tokens` (Histogram) - Histogram of generation token counts. - `vllm:time_to_first_token_seconds` (Histogram) - Time to first token (TTFT). -- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency. +- `vllm:inter_token_latency_seconds` (Histogram) - Inter-token latency (time between consecutive streamed outputs). +- `vllm:request_time_per_output_token_seconds` (Histogram) - Per-request Time Per Output Token (TPOT). - `vllm:e2e_request_latency_seconds` (Histogram) - End-to-end request latency. - `vllm:request_prefill_time_seconds` (Histogram) - Request prefill time. - `vllm:request_decode_time_seconds` (Histogram) - Request decode time. @@ -49,7 +50,8 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of - `vllm:e2e_request_latency_seconds_bucket` - End to end request latency measured in seconds. - `vllm:prompt_tokens` - Prompt tokens. - `vllm:generation_tokens` - Generation tokens. -- `vllm:inter_token_latency_seconds` - Inter-token latency (Time Per Output Token, TPOT) in seconds. +- `vllm:inter_token_latency_seconds` - Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains exactly one token, but differs from request-level TPOT when an output contains multiple tokens or when aggregation weights differ. +- `vllm:request_time_per_output_token_seconds` - Per-request Time Per Output Token (TPOT) in seconds, computed as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. It is recorded as zero for requests that generate no more than one token. - `vllm:time_to_first_token_seconds` - Time to First Token (TTFT) latency in seconds. - `vllm:num_requests_running` (also, `_swapped` and `_waiting`) - Number of requests in the RUNNING, WAITING, and SWAPPED states. - `vllm:kv_cache_usage_perc` - Percentage of used cache blocks by vLLM. @@ -63,6 +65,20 @@ The subset of metrics exposed in the Grafana dashboard gives us an indication of See [the PR which added this Dashboard](https://github.com/vllm-project/vllm/pull/2316) for interesting and useful background on the choices made here. +`vllm:inter_token_latency_seconds` records one sample per streamed output event +(the wall-clock gap between successive outputs), whereas +`vllm:request_time_per_output_token_seconds` is recorded once per finished request +as `(end-to-end latency - TTFT) / (number of output tokens - 1)`. These two +metrics differ when an output bundles multiple tokens (e.g. speculative +decoding) or when mean ITL and mean per-request TPOT are aggregated over +differing weights. See [Understanding the latency metrics](../benchmarking/cli.md#understanding-the-latency-metrics) +for the benchmark definitions. Requests that generate no more than one token +are recorded with a TPOT of zero, while `vllm bench serve` excludes them from +its TPOT statistics. Therefore, the Prometheus histogram and benchmark TPOT +statistics can differ. Use `vllm:request_time_per_output_token_seconds` for +request-level TPOT and `vllm:inter_token_latency_seconds` when you specifically +want the inter-output latency distribution. + ### Prometheus Client Library Prometheus support was initially added [using the aioprometheus library](https://github.com/vllm-project/vllm/pull/1890), but a switch was made quickly to [prometheus_client](https://github.com/vllm-project/vllm/pull/2730). The rationale is discussed in both linked PRs. @@ -238,8 +254,7 @@ statistics relating to that iteration: iteration. - The prefill intervals for any requests that completed prefill in this iteration. -- The inter-token intervals (Time Per Output Token, TPOT), for all - requests included in this iteration. +- The inter-token intervals, for all requests included in this iteration. - The Time-To-First-Token (TTFT) for any requests that completed prefill in this iteration. However, we calculate this interval relative to when the request was first received by the frontend diff --git a/docs/design/nixl_kv_push_connector.md b/docs/design/nixl_kv_push_connector.md index 6f66be5eb9f5..a6debd7dba81 100644 --- a/docs/design/nixl_kv_push_connector.md +++ b/docs/design/nixl_kv_push_connector.md @@ -181,6 +181,50 @@ The completion notif sent from P to D after a WRITE is the existing is D's own request id, taken from the registration), so the D-side accounting code is unchanged. +## Pipeline parallelism and hybrid KV caches + +The pull connector requires the local and remote workers to expose a +congruent list of KV regions — region *i* here corresponds to region +*i* there. That assumption breaks under pipeline parallelism (PP) +combined with a hybrid (HMA) KV layout: + +* a PP-sharded prefiller (**P**) holds only a slice of the model's + layers while the `PP=1` decoder (**D**) holds them all, so region + counts differ; +* with HMA, several layer names are pooled into one region and the layer + that represents a pooled region can differ between P and D. + +`NixlPushConnector` handles this for PP-sharded producers by routing +**by layer-name (member) identity** instead of by region index. Each +worker advertises which layer names back each of its NIXL regions in the +handshake metadata (`NixlAgentMetadata.region_members`). A producer that +needs member routing derives its member-major layout once, when it +registers its KV caches, and every transfer it issues uses that order. +`add_remote_agent` then selects exactly the remote regions this stage +owns and reorders them to match, so both sides stay paired regardless of +how each remote rank happens to order its metadata. + +Addresses, block lengths, strides, and per-region capacities follow the same +member order. Descriptor offsets use each member's region capacity, so P and +D need not allocate the same number of blocks. Physical allocations are still +registered once, even when multiple members share them. + +Invariants enforced when the remote regions are aligned: + +* every locally owned member must be advertised exactly once by the + remote; a missing member fails the handshake rather than silently + leaving that layer's KV stale, and remote-only members (owned by other + PP stages) are ignored; +* a remote that omits member metadata while the local layout requires + member routing fails the handshake instead of falling back to + region-index routing; +* member order is a property of the local layout alone, so the same local + source descriptors serve every remote engine and TP rank; only the + remote descriptor list is rebuilt per rank. + +Decode-side PP is unsupported because completions are counted per +consumer rank. Mamba/SSM hybrids are unsupported under PP. + ## Scheduler-side responsibilities `NixlPushConnectorScheduler` extends the base scheduler with: diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 54a54d89887d..22a2350d8095 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -22,6 +22,35 @@ flowchart LR CPU <--> SN["..."] ``` +## Per-request load control + +Individual requests can cap how many tokens are loaded from offloaded storage +by setting `max_load_tokens` in `kv_transfer_params`. The cap applies to tokens +beyond those already available in the GPU prefix cache. Set it to `0` to +disable external loading for the request: + +```json +{ + "kv_transfer_params": { + "max_load_tokens": 0 + } +} +``` + +GPU prefix-cache reuse remains enabled, and tokens beyond the aligned load cap +are recomputed. Omitting the field leaves loading uncapped. The value must be a +non-negative integer; invalid values are ignored. Positive caps are rounded +down to a boundary supported by the configured KV cache groups. Offloading +remains enabled unless it is controlled separately with `max_offload_tokens`. + +!!! warning + `max_load_tokens` is experimental and subject to change. + +`kv_load_tiers` continues to select secondary tiers. CPU is always included: +it can satisfy a resident hit directly and is the required staging tier for +secondary-to-GPU loads. Consequently, an empty tier list allows CPU hits but +does not query any secondary tier. + ## Terminology: Chunks The unit of operation is a **chunk** — a fixed-size piece of KV data covering a group of tokens. By default, a chunk maps to a single accelerator block. A configurable `blocks_per_chunk` parameter allows larger chunks, yielding larger I/Os to the host and secondary tiers. @@ -203,9 +232,31 @@ Block content hashes must match across instances for peers to exchange blocks (s | `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. | | `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. | | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. | +| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. | The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process. +A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives +within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier +slots. Raise it when prefills legitimately take longer than the default: + +```bash +vllm serve \ + --kv-transfer-config '{ + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "spec_name": "TieringOffloadingSpec", + "secondary_tiers": [ + { + "type": "p2p", + "unbound_store_timeout_s": 180 + } + ] + } + }' +``` + #### Environment Variables Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence. diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 031fba5b2f04..a41d5868b169 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -104,6 +104,24 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting. - Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA. +### Pipeline parallelism + +The default (pull) NixlConnector does **not** support +`pipeline-parallel-size > 1` together with a hybrid KV cache (HMA) +layout: region indices are not a stable identity across the +prefill/decode layer split, so the connector raises at startup. Use the +push connector (`NixlPushConnector`) for pipeline parallelism with +hybrid KV caches — it routes transfers by layer-name (member) identity, +so a PP-sharded prefiller can write into a `PP=1` decoder. See +[NIXL push-mode KV transfer](../design/nixl_kv_push_connector.md). + +Current push PP + HMA limitations: + +- Only the prefiller (producer) may be PP-sharded; decode-side PP is not supported. +- Hybrid SSM/Mamba layouts are not supported under PP. +- HMA requires the same block size on P and D. +- Attention-HMA member routing requires decode TP to be no greater than prefill TP. + ### Quantized KV cache [Quantized KV cache](quantization/quantized_kvcache.md) (e.g., FP8) requires both P and D instances to use the **same** `cache_dtype`. Mismatched cache dtypes will fail the compatibility hash check during handshake. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d55348ad724a..ddc50ad41cdb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -602,6 +602,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MossTranscribeDiarizeForConditionalGeneration` | MOSS-Transcribe-Diarize | T + A | `OpenMOSS-Team/MOSS-Transcribe-Diarize` | | ✅︎ | | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ | | `MuseGlimmerForCausalLM`, `MuseGlimmerForConditionalGeneration` | Muse Glimmer | T + I+ + V+ | `meta-models/Muse-Glimmer-30B` | ✅︎ | ✅︎ | +| `NemotronH_Nano_Omni_Reasoning_V3`, `NemotronH_Nano_VL_V2` | Nemotron Nano 3 Omni, Nemotron Nano V2 VL | T + IE+ + V+ + A* | `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` | ✅︎ | | | `NVLM_D_Model` | NVLM-D 1.0 | T + I+ | `nvidia/NVLM-D-72B`, etc. | | ✅︎ | | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + IE+ | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ | | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + IE+ + VE+ | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ | @@ -701,6 +702,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor `meta-models/Muse-Glimmer-30B-assistant` checkpoint, which vLLM serves through its [DFlash](../features/speculative_decoding/README.md) path. +!!! note + Both Nemotron Nano architecture names map to the same vLLM implementation. + Audio inputs are only supported by the Omni variants; `NemotronH_Nano_VL_V2` + checkpoints such as `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` accept text, + image and video only. + !!! note The official `openbmb/MiniCPM-V-2` doesn't work yet, so we need to use a fork (`HwwwH/MiniCPM-V-2`) for now. For more details, please see: diff --git a/docs/serving/context_parallel_deployment.md b/docs/serving/context_parallel_deployment.md index 3df2b06635f5..cb9a1ece71a4 100644 --- a/docs/serving/context_parallel_deployment.md +++ b/docs/serving/context_parallel_deployment.md @@ -20,6 +20,15 @@ Both approaches are under active development. Due to the auto-regressive nature of decoding, every decoding step needs to compute a small amount of query tokens w.r.t. a large number of key/value tokens stored in the paged KV cache. The core of decode context parallel is how to shard the KV cache across GPUs. +`engine_client.vllm_config.cache_config.effective_attention_block_size` reports the +initialized full-attention block size in tokens, including DCP: a physical block of +16 tokens with DCP=4 represents 64 tokens. The same value is available through +gRPC `Control.GetServerInfo.effective_attention_block_size`. Physical block-size +fields retain their existing meaning. The new value is unavailable (`None` in +Python, absent in gRPC) if there is no common full-attention size, the scheduler +does not expose it, or an older engine omits it. gRPC also omits the value when +engine replicas disagree. Partial cache events carry their own actual size. + For a model with `H` kv-heads, a request with `T` tokens in the context needs to store `H * T` key/value tensors in the KV cache. 1. If one GPU can hold them all, and the performance is good enough, then no parallelization is needed. diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 1f5d4bdb6a93..a292a4a3a436 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -130,7 +130,12 @@ Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` ### Renderer APIs -For further details on renderer APIs, please refer to [this page](renderer.md). +Renderer APIs preprocess completion, chat, and Responses requests without running +inference. They handle tokenization, model-specific prompt formatting, and +multimodal preprocessing, returning prompt token IDs, sampling parameters, and +any processed multimodal inputs for generation. + +See the [renderer guide](renderer.md) for setup instructions and examples. - [Completions Render API](renderer.md) (`/v1/completions/render`) - Render completion requests diff --git a/docs/serving/online_serving/openai_compatible_server.md b/docs/serving/online_serving/openai_compatible_server.md index 22e2e9a60844..c50910040d32 100644 --- a/docs/serving/online_serving/openai_compatible_server.md +++ b/docs/serving/online_serving/openai_compatible_server.md @@ -199,6 +199,11 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai Code example: [examples/tool_calling/openai_responses_client_with_tools.py](../../../examples/tool_calling/openai_responses_client_with_tools.py) +To get prompt token IDs before generation, use +[`/v1/responses/render`](renderer.md#get-responses-prompt-token-ids). It preprocesses +a self-contained Responses request without inference and returns the prompt +token IDs, so a separate `/tokenize` call is not needed. + #### Extra parameters The following extra parameters in the request object are supported: diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 9be23142c49c..9d494df22efe 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -26,27 +26,58 @@ vllm serve --enable-scale-out - [Responses Render API](renderer.md) (`/v1/responses/render`) - Render a self-contained Responses request +## Get Responses prompt token IDs + +Use `/v1/responses/render` to get prompt token IDs before choosing a model +replica. Rendering applies prompt construction and preprocessing without running +inference. + The Responses render endpoint uses the same prompt construction as `/v1/responses` and returns one token-in `GenerateRequest`. It is stateless: inline history is supported, but `previous_response_id` is not. Callers must resolve stored response state and include the resulting history in the request before rendering. +Configure the renderer and generation workers with the same model, tokenizer, +chat template, and preprocessing options. Render the full request, including +instructions, history, tools, and any template or truncation options, so the +returned IDs reflect the prompt the model will receive. + For multimodal requests, the `GenerateRequest` contains the model-processed multimodal payload, which can be substantially larger than the source image or video. The caller must forward that payload unchanged to the generation service and provision transport limits and memory accordingly. +For example, start a standard inference server with scale-out endpoints enabled: + +```bash +VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + vllm serve meta-llama/Llama-3.1-8B-Instruct +``` + +Send a Responses request and use `jq` to extract the `token_ids` field: + ```bash -curl http://localhost:8000/v1/responses/render \ +curl --fail --silent --show-error http://localhost:8000/v1/responses/render \ -H "Content-Type: application/json" \ -d '{ "model": "meta-llama/Llama-3.1-8B-Instruct", "input": "Explain prefix caching in one sentence.", "max_output_tokens": 32 - }' + }' | jq '.token_ids' ``` +To count the rendered prompt tokens for this text request, replace the `jq` +filter with `'.token_ids | length'`. The endpoint returns the full +`GenerateRequest`; `jq` filters the response on the client. Keep the complete +response if you will forward it to `/inference/v1/generate`, including any +multimodal features. + +If the server has `--api-key` or `VLLM_API_KEY` configured, add +`-H "Authorization: Bearer "` to the request. See +[API key authentication limitations](../../usage/security.md#api-key-authentication-limitations) +for the existing authentication boundaries. + For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md). ## Multimodal Render Features diff --git a/docs/usage/security.md b/docs/usage/security.md index 3c7f1e3f9ba8..5d4d429c6fed 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -170,7 +170,7 @@ When `--api-key` is configured, the following endpoints require Bearer token aut - `/v1/messages/render` - Render Anthropic-compatible messages (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render`) - `/v1/messages/count_tokens` - Count tokens for Anthropic messages - `/v1/responses` - Create a response -- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1`, or on `vllm launch render` unless explicitly disabled) +- `/v1/responses/render` - Render a self-contained response request (available on `vllm serve` only when `--enable-scale-out` is set, or on `vllm launch render` unless explicitly disabled) - `/v1/responses/{response_id}` - Retrieve a response - `/v1/responses/{response_id}/cancel` - Cancel a response - `/v1/score` - Scoring API diff --git a/examples/observability/prometheus_grafana/grafana.json b/examples/observability/prometheus_grafana/grafana.json index 1c89d4593830..6144f2bfd3ad 100644 --- a/examples/observability/prometheus_grafana/grafana.json +++ b/examples/observability/prometheus_grafana/grafana.json @@ -318,7 +318,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "description": "Inter token latency in seconds.", + "description": "Inter-token latency in seconds, measured between consecutive streamed outputs. It approximates TPOT when each output contains one token, but differs from request-level TPOT under speculative decoding or varying aggregation weights.", "fieldConfig": { "defaults": { "color": { @@ -479,6 +479,172 @@ "title": "Inter Token Latency", "type": "timeseries" }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Per-request Time Per Output Token (TPOT) in seconds, computed as (end-to-end latency - TTFT) divided by the number of output tokens minus one. Requests with no more than one output token are recorded as zero; vllm bench serve excludes these requests from its TPOT statistics.", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + }, + "unit": "s" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 48 + }, + "id": 17, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.99, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.95, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P95", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.9, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "disableTextWrap": false, + "editorMode": "builder", + "expr": "histogram_quantile(0.5, sum by(le) (rate(vllm:request_time_per_output_token_seconds_bucket{model_name=\"$model_name\"}[$__rate_interval])))", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": false, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "D", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "rate(vllm:request_time_per_output_token_seconds_sum{model_name=\"$model_name\"}[$__rate_interval])\n/\nrate(vllm:request_time_per_output_token_seconds_count{model_name=\"$model_name\"}[$__rate_interval])", + "hide": false, + "instant": false, + "legendFormat": "Mean", + "range": true, + "refId": "E" + } + ], + "title": "Time Per Output Token (TPOT)", + "type": "timeseries" + }, { "datasource": { "type": "prometheus", diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt index 18a644f6b1b4..8b9f7d31550b 100644 --- a/requirements/rubin-prerelease.txt +++ b/requirements/rubin-prerelease.txt @@ -2,15 +2,25 @@ # Installed when INSTALL_RUBIN_PRERELEASE=true (see docker/Dockerfile). # TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels # (vllm-project/vllm#49735). ---pre --extra-index-url https://download.pytorch.org/whl/nightly/cu134 +--extra-index-url https://flashinfer.ai/whl/cu134 +--extra-index-url https://flashinfer.ai/whl -torch==2.15.0.dev20260903+cu134 -torchvision==0.30.0.dev20260903+cu134 -torchaudio==2.11.0.dev20260903+cu134 +# Match each architecture's torch pin to the latest torchvision wheel. +torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64" +torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64" +torchvision==0.30.0.dev20260909+cu134 +torchaudio==2.11.0.dev20260911+cu134 -# docker/Dockerfile strips these two from requirements/cuda.txt whenever +# docker/Dockerfile strips these packages from requirements/cuda.txt whenever # INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to # still get installed. nvidia-cutlass-dsl[cu13]==4.8.0.dev0 -flashinfer-python==0.6.18 +flashinfer-python==0.6.18.post1 +quack-kernels==0.6.5 + +cuda-python==13.4.1 +cuda-bindings==13.4.1 +flashinfer-jit-cache==0.6.18.post1+cu134 +flashinfer-cubin==0.6.18.post1 +nixl==1.4.1 diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index b7cf49649b7e..2e2fbd88ede6 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -257,7 +257,7 @@ googleapis-common-protos==1.75.1 # google-api-core # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http -gpt-oss==0.0.9 +gpt-oss==0.0.8 # via -r requirements/test/xpu.in graphql-core==3.2.11 # via hypothesis-graphql @@ -406,10 +406,6 @@ jsonschema-rs==0.49.9 # via schemathesis jsonschema-specifications==2025.9.1 # via jsonschema -jupyter-client==8.10.0 - # via gpt-oss -jupyter-core==5.9.1 - # via jupyter-client kiwisolver==1.5.1 # via matplotlib lark==1.2.2 @@ -672,9 +668,7 @@ pillow==12.3.0 # mistral-common # torchvision platformdirs==4.11.3 - # via - # jupyter-core - # pooch + # via pooch pluggy==1.6.0 # via # pytest @@ -813,7 +807,6 @@ pytest-timeout==2.4.0 python-dateutil==2.9.0.post0 # via # botocore - # jupyter-client # matplotlib # pandas # typepy @@ -849,7 +842,6 @@ pyzmq==27.1.0 # via # -c requirements/common.txt # -r requirements/test/../common.txt - # jupyter-client rapidfuzz==3.14.5 # via # -r requirements/test/xpu.in @@ -1074,8 +1066,6 @@ torchvision==0.28.0+xpu # via # open-clip-torch # timm -tornado==6.5.8 - # via jupyter-client tqdm==4.70.0 # via # -r requirements/test/../common.txt @@ -1091,10 +1081,6 @@ tqdm==4.70.0 # pqdm # sentence-transformers # transformers -traitlets==5.16.1 - # via - # jupyter-client - # jupyter-core transformers==5.16.1 # via # -c requirements/common.txt @@ -1146,7 +1132,6 @@ typing-extensions==4.16.0 # grpcio # httpx2 # huggingface-hub - # jupyter-client # lm-eval # mcp # mcp-types diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 1b7dae65685b..54a1640926fb 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5813,6 +5813,7 @@ dependencies = [ "vllm-engine-core-client", "vllm-managed-engine", "vllm-server", + "vllm-text", "vllm-tracing", ] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index bd9d99de623c..dde065e79044 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -147,7 +147,7 @@ vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } vllm-parser = { path = "src/parser" } -vllm-proto = { version = "0.2.0", path = "proto" } +vllm-proto = { path = "proto" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } diff --git a/rust/proto/README.md b/rust/proto/README.md index 62b134280bd9..b8a20b0c349b 100644 --- a/rust/proto/README.md +++ b/rust/proto/README.md @@ -2,14 +2,13 @@ This directory is the canonical source for vLLM's gRPC schema. -The schema is published to `buf.build/vllm-project/vllm`: +See [context parallel deployment](../../docs/serving/context_parallel_deployment.md) +for effective attention block-size metadata in Python and `Control.GetServerInfo`. -- A daily workflow publishes the latest Git `main` schema to the `nightly` label. -- The workflow can be run manually to retry nightly publication. -- Tags matching `v*` update the Buf `main` label and publish the corresponding release label. -- Buf commits and generated SDK versions are immutable and can be pinned by consumers. - -Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically. +Schema updates are no longer published to the Buf Schema Registry. Rust consumers +should use `vllm-proto` from crates.io; consumers in other languages can generate +bindings from the `.proto` files in this directory. Buf still builds and lints +the schemas on pull requests, including on forks. ## Rust crate @@ -38,8 +37,7 @@ On pull requests and releases, `cargo-semver-checks` compares the crate with its latest published version. Include any required version bump in the protocol change PR. This check becomes available after the first manual publication. -1. Update the crate version and the `vllm-proto` workspace dependency together, - and update `rust/Cargo.lock`. +1. Update the crate version in `rust/proto/Cargo.toml` and update `rust/Cargo.lock`. 2. Run `cargo publish --manifest-path rust/proto/Cargo.toml --locked --dry-run` and the frontend gRPC tests. Record the tested vLLM releases or revisions in the release notes; matching crate versions alone do not establish runtime compatibility. @@ -47,8 +45,7 @@ change PR. This check becomes available after the first manual publication. The `proto-crate.yml` workflow checks that the tag matches the crate version, verifies the package, and publishes it to crates.io. -Use `proto-v*` tags so they do not trigger Buf's `v*` release workflow. There is -no scheduled crate publication. Pull requests and manual workflow runs only +There is no scheduled crate publication. Pull requests and manual workflow runs only verify packaging, including on forks. Retry a failed release by rerunning its tag workflow; crates.io versions cannot be overwritten. diff --git a/rust/proto/buf.md b/rust/proto/buf.md deleted file mode 100644 index 065124feee65..000000000000 --- a/rust/proto/buf.md +++ /dev/null @@ -1,8 +0,0 @@ -# vLLM gRPC API - -Protocol definitions for the optional gRPC server exposed by the vLLM Rust frontend. - -- `Inference` provides unary and streaming generation. -- `Control` provides discovery and request control. - -The `nightly` label is updated daily from vLLM's `main` branch. The `main` label tracks the latest vLLM release, and version labels match tags such as `v0.27.0`. Pin a Buf commit or generated SDK version for reproducible builds. diff --git a/rust/proto/buf.yaml b/rust/proto/buf.yaml index 64764f561ef3..d30dc3297d76 100644 --- a/rust/proto/buf.yaml +++ b/rust/proto/buf.yaml @@ -1,7 +1,6 @@ version: v2 modules: - path: . - name: buf.build/vllm-project/vllm lint: use: - MINIMAL diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 3050035fdcf6..d2816f68d347 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -50,6 +50,8 @@ message ServerInfo { uint64 max_batched_tokens = 9; uint32 max_loras = 10; RlCapabilities rl_capabilities = 11; + // Full-attention block size in tokens after initialization; absent if unavailable. + optional uint64 effective_attention_block_size = 14; } message RlCapabilities { diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md index d4e0c69b80b0..c8724d6cd675 100644 --- a/rust/src/bench/README.md +++ b/rust/src/bench/README.md @@ -667,6 +667,7 @@ With `--multi-turn`, `--num-prompts` controls the number of **conversations**, n - `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`. - `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns. +- `--dataset-name hf` — downloads a ShareGPT-format HuggingFace config, then loads it like `sharegpt`. **Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0. diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs index 3a0936b5f4c1..b5a24d942f50 100644 --- a/rust/src/bench/src/config.rs +++ b/rust/src/bench/src/config.rs @@ -549,13 +549,15 @@ impl BenchConfig { )); } - // Normalize and validate min/max turns. ShareGPT only consumes max_turns - // (the loader walks all available turns up to the cap), so the - // min/num/max coupling used for synthetic generation does not apply. - if args.dataset_name == DatasetName::ShareGpt { + // Normalize and validate min/max turns. Conversation datasets only + // consume max_turns (the loader walks all available turns up to the + // cap), so the min/num/max coupling used for synthetic generation + // does not apply. + if matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) { if args.multi_turn_max_turns == 1 { return Err(BenchError::Config( - "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(), + "--multi-turn-max-turns must be at least 2 for ShareGPT-format multi-turn" + .into(), )); } } else { @@ -722,7 +724,7 @@ impl BenchConfig { multi_turn_min_turns, multi_turn_max_turns, sharegpt_multi_turn_max_turns: if args.multi_turn - && args.dataset_name == DatasetName::ShareGpt + && matches!(args.dataset_name, DatasetName::ShareGpt | DatasetName::Hf) && args.multi_turn_max_turns != 0 { Some(args.multi_turn_max_turns) @@ -1067,6 +1069,30 @@ mod tests { assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20)); } + #[test] + fn test_hf_multi_turn_uses_conversation_turn_cap() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "hf", + "--dataset-path", + "org/sharegpt-dataset", + "--hf-subset", + "sharegpt", + "--multi-turn-max-turns", + "2", + ]; + let args = parse_args(args); + let config = BenchConfig::from_args(&args).unwrap(); + + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2)); + } + #[test] fn test_sweep_summary_percentiles_default_empty() { let args = base_multi_turn_args(); diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs index 97ec522e29ad..48ff11ff7bf6 100644 --- a/rust/src/bench/src/datasets/multi_turn.rs +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -350,6 +350,27 @@ pub fn load_sharegpt_multi_turn( .as_array() .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?; + load_sharegpt_multi_turn_from_rows( + tokenizer, + entries, + num_conversations, + output_len_override, + max_turns, + seed, + request_id_prefix, + ) +} + +/// Load multi-turn conversations from deserialized ShareGPT rows. +pub fn load_sharegpt_multi_turn_from_rows( + tokenizer: &TokenizerKind, + entries: &[serde_json::Value], + num_conversations: usize, + output_len_override: Option, + max_turns: Option, + seed: u64, + request_id_prefix: &str, +) -> Result> { // Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant) let mut filtered: Vec<&serde_json::Value> = entries .iter() @@ -507,6 +528,42 @@ fn gen_prompt_to_target_len( mod tests { use super::*; + #[test] + fn test_load_sharegpt_multi_turn_from_hf_rows() { + let tokenizer = + TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap()); + let rows = vec![serde_json::json!({ + "conversations": [ + {"from": "human", "value": "hello"}, + {"from": "gpt", "value": "hi"}, + {"from": "human", "value": "how are you?"}, + {"from": "gpt", "value": "great"} + ] + })]; + + let conversations = + load_sharegpt_multi_turn_from_rows(&tokenizer, &rows, 1, Some(7), None, 42, "hf-") + .unwrap(); + let actual = conversations + .iter() + .map(|conversation| { + ( + conversation.conversation_id.as_str(), + conversation + .turns + .iter() + .map(|turn| (turn.user_message.as_ref(), turn.expected_output_len)) + .collect::>(), + ) + }) + .collect::>(); + + assert_eq!( + actual, + vec![("hf-conv-0", vec![("hello", 7), ("how are you?", 7)])] + ); + } + fn common_prefix_bytes(strings: &[&str]) -> usize { if strings.is_empty() { return 0; diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs index e0a31381705f..366fd4f6fbc1 100644 --- a/rust/src/bench/src/multi_turn.rs +++ b/rust/src/bench/src/multi_turn.rs @@ -186,9 +186,30 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { - return Err(BenchError::Config( - "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), - )); + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?; + let dataset_id = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config("--dataset-path is required for --dataset-name hf".into()) + })?; + let (rows, _config, _split) = crate::datasets::hf_dataset::download_hf_dataset( + dataset_id, + config.hf_subset.as_deref(), + config.hf_split.as_deref(), + config.num_prompts, + config.seed, + config.disable_shuffle, + ) + .await?; + crate::datasets::multi_turn::load_sharegpt_multi_turn_from_rows( + tok, + &rows, + config.num_prompts, + config.hf_output_len, + config.sharegpt_multi_turn_max_turns, + config.seed, + &config.request_id_prefix, + )? } }; diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index 37212a96505c..5b5b8772c55f 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -129,7 +129,8 @@ pub(super) async fn load_model_backends( model_id: &str, options: LoadModelBackendsOptions, ) -> Result { - let files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; + let mut files = ResolvedModelFiles::new(model_id, options.revision.as_deref()).await?; + files.apply_overrides(&options.hf_overrides)?; let text_backend = HfTextBackend::from_resolved_model_files( files.clone(), model_id.to_string(), @@ -166,13 +167,11 @@ fn resolve_multimodal_render_info( #[cfg(test)] mod tests { use std::collections::HashMap; - use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; use thiserror_ext::AsReport as _; use vllm_text::Prompt; - use vllm_text::backend::hf::TokenizerSource; use vllm_text::tokenizer::DynTokenizer; use vllm_tokenizer::test_utils::TestTokenizer; @@ -206,16 +205,12 @@ mod tests { write_json(&config_path, config_json); write_json(&tokenizer_config_path, tokenizer_config_json); - vllm_text::backend::hf::ResolvedModelFiles { - tokenizer: TokenizerSource::HuggingFace(PathBuf::from("/tmp/unused-tokenizer.json")), - tokenizer_config_path: Some(tokenizer_config_path), - generation_config_path: None, - preprocessor_config_path: None, - video_preprocessor_config_path: None, - processor_config_path: None, - chat_template_path: None, - config_path: Some(config_path), - } + write_json(&root.join("tokenizer.json"), "{}"); + futures::executor::block_on(vllm_text::backend::hf::ResolvedModelFiles::new( + root.to_str().unwrap(), + None, + )) + .unwrap() } fn test_tokenizer() -> DynTokenizer { @@ -232,6 +227,7 @@ mod tests { "test-model".to_string(), LoadModelBackendsOptions { revision: None, + hf_overrides: Default::default(), generation_config: Default::default(), renderer, language_model_only: false, diff --git a/rust/src/chat/src/backend/mod.rs b/rust/src/chat/src/backend/mod.rs index 028c0c40c8ec..eed4e95cbe0c 100644 --- a/rust/src/chat/src/backend/mod.rs +++ b/rust/src/chat/src/backend/mod.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use std::sync::Arc; use serde_json::Value; +use vllm_text::backend::hf::HfOverrides; use vllm_text::{DynTextBackend, GenerationConfigMode, TextBackend}; use crate::error::Result; @@ -63,6 +64,8 @@ pub type DynChatTextBackend = Arc; pub struct LoadModelBackendsOptions { /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). pub revision: Option, + /// JSON Merge Patch applied to the model config before loading any backend. + pub hf_overrides: HfOverrides, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Which chat renderer implementation to use. diff --git a/rust/src/chat/src/renderer/deepseek.rs b/rust/src/chat/src/renderer/deepseek.rs index 5925c5fd1655..01576a7c4341 100644 --- a/rust/src/chat/src/renderer/deepseek.rs +++ b/rust/src/chat/src/renderer/deepseek.rs @@ -224,21 +224,30 @@ pub(super) fn render_request_with_media_order( } } - if (is_user_like_entry(message, current_render_index as usize, dialect) - || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. }))) - && next_rendered_entry_is_assistant_or_end( + if is_user_like_entry(message, current_render_index as usize, dialect) + || (dialect == DsDialect::V4 && matches!(message, ChatMessage::System { .. })) + { + // Skip the opener only at the last rendered entry when + // add_generation_prompt is off. Inspect the next rendered + // entry, not the raw index: V4 merges trailing tool results. + let should_write_transition = match next_rendered_entry( request.messages.as_slice(), message_index, drop_historical_developers, last_user_like_message_index, - ) - { - write_assistant_transition( - &mut out, - thinking_mode, - drop_thinking, - current_render_index >= last_user_render_index, - ); + ) { + Some(ChatMessage::Assistant { .. }) => true, + None => request.chat_options.add_generation_prompt(), + Some(_) => false, + }; + if should_write_transition { + write_assistant_transition( + &mut out, + thinking_mode, + drop_thinking, + current_render_index >= last_user_render_index, + ); + } } } @@ -420,14 +429,17 @@ fn is_historical_developer( && last_user_like_message_index.is_some_and(|last_index| message_index < last_index) } -/// Return whether the next rendered entry is assistant, or there is no next -/// entry. -fn next_rendered_entry_is_assistant_or_end( +/// Return the next rendered message after `message_index`. +/// +/// Skips merged user-content tails and dropped historical developers so the +/// caller can tell a following assistant turn apart from the end of the +/// conversation. +fn next_rendered_entry( messages: &[ChatMessage], message_index: usize, drop_historical_developers: bool, last_user_like_message_index: Option, -) -> bool { +) -> Option<&ChatMessage> { let mut next_index = message_index + 1; while next_index < messages.len() && (is_following_user_content(messages, next_index) @@ -441,10 +453,7 @@ fn next_rendered_entry_is_assistant_or_end( next_index += 1; } - messages - .get(next_index) - .map(|message| matches!(message, ChatMessage::Assistant { .. })) - .unwrap_or(true) + messages.get(next_index) } /// Render the tool preamble shown to the model for one DeepSeek dialect. diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs index 3ccf4d11db7f..4fc5f26530f8 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs @@ -61,6 +61,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { find_last_user_render_index(request.messages.as_slice(), render_offset); let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice()); let continue_final_message = request.chat_options.continue_final_message(); + let add_generation_prompt = request.chat_options.add_generation_prompt(); let mut prompt = String::from(BOS_TOKEN); if !initial_tools.is_empty() { @@ -79,6 +80,7 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { thinking_mode, drop_thinking, continue_final_message, + add_generation_prompt, )?; } @@ -110,11 +112,13 @@ fn render_message( thinking_mode: ThinkingMode, drop_thinking: bool, continue_final_message: bool, + add_generation_prompt: bool, ) -> Result<()> { let render_index = message_index as isize + render_offset; let opens_thinking = render_index == last_user_render_index; let after_last_user_turn = render_index > last_user_render_index; let after_or_at_last_user_turn = render_index >= last_user_render_index; + let add_generation_prompt = message_index + 1 != messages.len() || add_generation_prompt; match message { ChatMessage::System { content } => render_system_message(out, Some(content), &[]), @@ -123,11 +127,13 @@ fn render_message( content, tools.as_deref().unwrap_or(&[]), thinking_mode == ThinkingMode::Thinking && opens_thinking, + add_generation_prompt, ), ChatMessage::User { content } => render_user_message( out, content, thinking_mode == ThinkingMode::Thinking && opens_thinking, + add_generation_prompt, ), ChatMessage::Assistant { content } => render_assistant_message( out, @@ -195,6 +201,7 @@ fn render_developer_message( content: &ChatContent, tools: &[ChatTool], opens_thinking: bool, + add_generation_prompt: bool, ) -> Result<()> { if content.is_empty() { return Err(Error::ChatTemplate( @@ -209,7 +216,7 @@ fn render_developer_message( } out.push_str("\n\n# The user's message is: "); write_chat_content(out, content)?; - write_user_like_suffix(out, opens_thinking); + write_user_like_suffix(out, opens_thinking, add_generation_prompt); Ok(()) } @@ -219,17 +226,20 @@ fn render_user_message( out: &mut String, content: &ChatContent, opens_thinking: bool, + add_generation_prompt: bool, ) -> Result<()> { out.push_str("<｜User｜>"); write_chat_content(out, content)?; - write_user_like_suffix(out, opens_thinking); + write_user_like_suffix(out, opens_thinking, add_generation_prompt); Ok(()) } /// Shared trailing wrapper used by both real user turns and native developer /// turns after their content has already been written. -// TODO: respect `add_generation_prompt` option -fn write_user_like_suffix(out: &mut String, opens_thinking: bool) { +fn write_user_like_suffix(out: &mut String, opens_thinking: bool, add_generation_prompt: bool) { + if !add_generation_prompt { + return; + } out.push_str("<｜Assistant｜>"); if opens_thinking { out.push_str(THINKING_START_TOKEN); diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 3f93e3653efd..56c03a92cab4 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -281,6 +281,44 @@ fn assistant_after_last_user_requires_reasoning_or_tool_calls() { .assert_eq(&error.to_report_string()); } +#[test] +fn last_user_turn_appends_generation_prompt_by_default() { + let request = thinking_request(vec![ChatMessage::user("Hello")]); + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜>")); +} + +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::user("Hello")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello")); + assert!(!rendered.contains("<｜Assistant｜>")); + assert!(!rendered.contains("")); +} + +#[test] +fn user_then_assistant_still_writes_suffix_without_generation_prompt() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant_text("Hi there."), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>Hi there.<｜end▁of▁sentence｜>"] + .assert_eq(&rendered); +} + #[test] fn continue_final_assistant_omits_final_eos() { let mut request = ChatRequest { diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 172b67463c42..9404b54296a7 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -633,3 +633,94 @@ fn continue_final_assistant_omits_final_eos() { expect!["<｜begin▁of▁sentence｜><｜User｜>write<｜Assistant｜>partial answer"] .assert_eq(&rendered); } + +#[test] +fn last_user_turn_appends_generation_prompt_by_default() { + let request = thinking_request(vec![ChatMessage::user("Hello")]); + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello<｜Assistant｜>")); +} + +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::user("Hello")]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + assert!(rendered.ends_with("<｜User｜>Hello")); + assert!(!rendered.contains("<｜Assistant｜>")); + assert!(!rendered.contains("")); +} + +#[test] +fn user_then_assistant_still_writes_suffix_without_generation_prompt() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant_text("Hi there."), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("thinking".to_string(), Value::Bool(false)); + + let rendered = render_request(&request); + + expect![ + "<｜begin▁of▁sentence｜><｜User｜>Hello<｜Assistant｜>Hi there.<｜end▁of▁sentence｜>" + ] + .assert_eq(&rendered); +} + +#[test] +fn trailing_system_omits_generation_prompt_when_disabled() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::system("rules"), + ChatMessage::user("question"), + ChatMessage::system("final rules"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("enable_thinking".to_string(), Value::Bool(false)); + request.chat_options.reasoning_effort = Some(ReasoningEffort::Low); + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜>rules<｜User｜>questionfinal rules"].assert_eq(&rendered); +} + +#[test] +fn last_merged_tool_result_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ + ChatMessage::user("question"), + ChatMessage::tool_response("result", "call-1"), + ]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>question\n\nresult"] + .assert_eq(&rendered); + assert!(!rendered.contains("<｜Assistant｜>")); +} + +#[test] +fn last_developer_omits_generation_prompt_when_disabled() { + let mut request = thinking_request(vec![ChatMessage::developer("latest instruction", None)]); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + let rendered = render_request(&request); + + expect!["<｜begin▁of▁sentence｜><｜User｜>latest instruction"].assert_eq(&rendered); +} diff --git a/rust/src/chat/src/renderer/deepseek_v41/tests.rs b/rust/src/chat/src/renderer/deepseek_v41/tests.rs index 3babb5c269f8..3d621c2246fd 100644 --- a/rust/src/chat/src/renderer/deepseek_v41/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v41/tests.rs @@ -10,7 +10,9 @@ use super::DeepSeekV41ChatRenderer; use crate::ChatRenderer; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; -use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ReasoningEffort}; +use crate::request::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort, +}; fn render(request: &ChatRequest) -> String { DeepSeekV41ChatRenderer::new() @@ -281,6 +283,49 @@ fn mid_system_advances_last_user_boundary_and_drops_old_reasoning() { .assert_eq(&render(&request)); } +#[test] +fn last_user_turn_omits_generation_prompt_when_disabled() { + let mut request = request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + + expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 50 (range 1-100, the higher the value, the more thorough the reasoning)\n\n<｜User｜>question"] + .assert_eq(&render(&request)); +} + +#[test] +fn last_mid_system_omits_generation_prompt_when_disabled() { + let mut request = ChatRequest { + messages: vec![ + ChatMessage::system("sys"), + ChatMessage::user("q1"), + ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::Reasoning { text: "r1".into() }, + AssistantContentBlock::Text { text: "a1".into() }, + ]), + ChatMessage::developer("old policy", None), + ChatMessage::system("mid sys"), + ], + ..ChatRequest::for_test() + }; + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request + .chat_options + .template_kwargs + .insert("reasoning_effort".into(), json!(88)); + + expect!["<｜begin▁of▁sentence｜><｜System｜>Reasoning Effort: 88 (range 1-100, the higher the value, the more thorough the reasoning)\n\nsys<｜User｜>q1<｜Assistant｜>a1<｜end▁of▁sentence｜><｜System｜>mid sys"] + .assert_eq(&render(&request)); +} + +#[test] +fn chat_mode_omits_generation_prompt_when_disabled() { + let mut request = request(); + request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt; + request.chat_options.template_kwargs.insert("thinking".into(), json!(false)); + + expect!["<｜begin▁of▁sentence｜><｜User｜>question"].assert_eq(&render(&request)); +} + #[test] fn joins_text_parts_with_reference_separator_and_preserves_literal_tags() { let mut request = ChatRequest { diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index bf92e05ee186..e61c48ee3c83 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -33,6 +33,7 @@ vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true vllm-server.workspace = true +vllm-text.workspace = true vllm-tracing.workspace = true [dev-dependencies] diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 70b28018a3f1..cda2e6c1a8b9 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -33,6 +33,7 @@ use vllm_server::{ DEFAULT_KEEP_ALIVE_TIMEOUT, HttpListenerMode, LoraModulePath, ParserSelection, RenderConfig, RendererSelection, }; +use vllm_text::backend::hf::HfOverrides; use crate::cli::ssl::SslArgs; use crate::cli::unsupported::UnsupportedArgs; @@ -113,6 +114,9 @@ pub struct RenderArgs { /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). #[arg(long)] revision: Option, + /// JSON Merge Patch (RFC 7396) for config.json; null removes a field. + #[arg(long, value_parser = parse_json::, default_value = "{}", value_name = "JSON")] + hf_overrides: HfOverrides, /// HTTP bind host. #[arg(long, default_value = "127.0.0.1")] host: String, @@ -161,6 +165,7 @@ impl RenderArgs { RenderConfig { model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, served_model_name: self.served_model_name, host: self.host, port: self.port, @@ -204,6 +209,12 @@ pub struct SharedRuntimeArgs { #[serde(default)] pub revision: Option, + /// JSON Merge Patch (RFC 7396) for config.json; null removes a field. + /// Objects merge recursively and arrays/scalars replace existing values. + #[arg(long, value_parser = parse_json::, default_value = "{}", value_name = "JSON")] + #[serde(default)] + pub hf_overrides: HfOverrides, + /// The source of generation-config sampling defaults. `"auto"` loads the /// model's defaults, while `"vllm"` uses vLLM's neutral defaults. #[arg(long, default_value_t)] @@ -500,6 +511,7 @@ impl SharedRuntimeArgs { }, model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode: HttpListenerMode::InheritedFd { fd: listen_fd }, @@ -556,6 +568,7 @@ impl SharedRuntimeArgs { coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, revision: self.revision, + hf_overrides: self.hf_overrides, generation_config: self.generation_config, served_model_name: self.served_model_name, listener_mode, @@ -761,6 +774,9 @@ impl ServeArgs { let reasoning_parser = effective_engine_reasoning_parser(&self.runtime.reasoning_parser, &self.runtime.model); let profiler_config = self.runtime.profiler_config_json(); + let hf_overrides = (!self.runtime.hf_overrides.is_empty()).then(|| { + serde_json::to_string(&self.runtime.hf_overrides).expect("JSON object serializes") + }); self.managed_engine.clone().into_config( self.runtime.model.clone(), @@ -773,6 +789,7 @@ impl ServeArgs { self.runtime.shutdown_timeout, handshake_port, self.runtime.limit_mm_per_prompt_json(), + hf_overrides, ) } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 411d2562959c..5cc6b1921450 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -177,6 +177,9 @@ fn serve_args_forward_python_flags_with_separator() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1013,6 +1016,9 @@ fn frontend_args_accept_json() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: None, @@ -1120,6 +1126,73 @@ fn frontend_args_json_ignores_engine_owned_max_model_len() { assert_eq!(args.runtime.model, "Qwen/Qwen3-0.6B"); } +#[test] +fn hf_overrides_preserve_merge_patch_in_cli_and_python_bootstrap() { + let patch = r#"{"text_config":{"rope_parameters":{"factor":4}},"sliding_window":null}"#; + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--hf-overrides", + patch, + ]) + .unwrap(); + let Command::Serve(args) = cli.command else { + panic!("expected serve args") + }; + let expected: serde_json::Value = serde_json::from_str(patch).unwrap(); + assert_eq!( + serde_json::to_value(&args.runtime.hf_overrides).unwrap(), + expected + ); + let engine = args.to_managed_engine_config(1234); + let index = engine.python_args.iter().position(|arg| arg == "--hf-overrides").unwrap(); + assert_eq!( + serde_json::from_str::(&engine.python_args[index + 1]).unwrap(), + expected + ); + assert!(!engine.python_args.iter().any(|arg| arg == "--hf-config-path")); + let config = args.to_frontend_config("tcp://127.0.0.1:1234".to_string()); + assert_eq!(serde_json::to_value(config.hf_overrides).unwrap(), expected); + + let payload = serde_json::json!({"model_tag":"Qwen/Qwen3-0.6B", "hf_overrides":expected}); + let cli = Cli::try_parse_from([ + "vllm-rs", + "frontend", + "--listen-fd", + "3", + "--input-address", + "ipc:///tmp/input.sock", + "--output-address", + "ipc:///tmp/output.sock", + "--args-json", + &payload.to_string(), + ]) + .unwrap(); + let Command::Frontend(args) = cli.command else { + panic!("expected frontend args") + }; + assert_eq!( + serde_json::to_value(args.into_config().hf_overrides).unwrap(), + expected + ); +} + +#[test] +fn hf_overrides_reject_non_object_inputs() { + for patch in ["null", "[]", "1", r#""callable""#] { + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--hf-overrides", + patch, + ]) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); + } +} + #[test] fn frontend_args_json_accepts_supported_non_default_fields() { let cli = Cli::try_parse_from([ @@ -1613,6 +1686,9 @@ fn serve_args_accept_handshake_aliases() { runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, engine_ready_timeout_secs: 600, tool_call_parser: Auto, @@ -1763,6 +1839,9 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1852,6 +1931,9 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: BindTcp { @@ -1963,6 +2045,9 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present }, model: "Qwen/Qwen3-0.6B", revision: None, + hf_overrides: HfOverrides( + {}, + ), generation_config: Auto, served_model_name: [], listener_mode: InheritedFd { diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index 29ea5d06b0e7..8e514c95349c 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -235,11 +235,6 @@ pub struct EngineUnsupportedArgs { #[arg(long, default_missing_value = "true", num_args = 0..=1)] pub hf_token: Option, - /// If a dictionary, contains arguments to be forwarded to the Hugging Face - /// config. If a callable, it is called to update the HuggingFace config. - #[arg(long)] - pub hf_overrides: Option, - /// Overrides or sets generation config. e.g. `{"temperature": 0.5}`. If /// used with `--generation-config auto`, the override parameters will be /// merged with the default config from the model. If used with diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 63689ea3db83..4cd7420e561e 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -510,6 +510,15 @@ impl EngineCoreClient { self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum() } + /// Return the effective attention block size if all engines report the same value. + pub fn effective_attention_block_size(&self) -> Option { + let size = self.ready_response().effective_attention_block_size?; + self.engines + .iter() + .all(|engine| engine.ready_response.effective_attention_block_size == Some(size)) + .then_some(size) + } + /// Return the minimum engine-reported `max_model_len` across all engines. /// /// This is the auto-fitted value after KV cache profiling and may differ diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 117473ce47cc..7318fa01b8ce 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -19,7 +19,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; -use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder}; +use crate::metrics::{IterationMetricHandles, LoraInfoExporter, SchedulerStatsRecorder}; use crate::protocol::encode_msgpack; use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs}; use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; @@ -47,6 +47,7 @@ pub(crate) struct ClientInner { /// `VLLM_MSGPACK_ZERO_COPY_THRESHOLD` when this inner client is created. msgpack_zero_copy_threshold: usize, scheduler_stats_recorder: SchedulerStatsRecorder, + iteration_metrics: BTreeMap, request_reg: Mutex, utility_reg: Mutex, health_error: ArcSwapOption, @@ -64,12 +65,21 @@ impl ClientInner { ) -> Self { let scheduler_stats_recorder = SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines); + let iteration_metrics = engines + .iter() + .filter_map(|engine| { + let engine = engine.engine_id.engine_index()?; + let handles = IterationMetricHandles::new(&METRICS.request, &model_name, engine); + Some((engine, handles)) + }) + .collect(); Self { input_send, handle, model_name, msgpack_zero_copy_threshold: msgpack_zero_copy_threshold(), scheduler_stats_recorder, + iteration_metrics, request_reg: Mutex::new(RequestRegistry::new(engines)), utility_reg: Mutex::new(UtilityRegistry::default()), health_error: ArcSwapOption::empty(), @@ -415,6 +425,8 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs { EngineCoreOutputs::RequestBatch(batch) => { + let has_outputs = !batch.outputs.is_empty(); + let mut iteration_tokens = 0_u64; let senders = inner.take_senders_for_outputs(&batch.outputs); for (output, sender) in batch.outputs.into_iter().zip(senders) { let request_id = output.request_id.clone(); @@ -423,6 +435,12 @@ pub(crate) async fn run_output_dispatcher_loop( continue; }; + iteration_tokens += output.new_token_ids.len() as u64; + if let Some(prefill) = &output.prefill_stats { + // The engine emits prefill_stats once, on the first output. + iteration_tokens += u64::from(prefill.num_computed_tokens); + } + let wrapped_output = EngineCoreStreamOutput { engine_index: batch.engine_index, timestamp: batch.timestamp, @@ -433,6 +451,12 @@ pub(crate) async fn run_output_dispatcher_loop( } } + if has_outputs + && let Some(handles) = inner.iteration_metrics.get(&batch.engine_index) + { + handles.iteration_tokens_total.observe(iteration_tokens as f64); + } + // The sender for normally-finished requests should have already been removed // from the registry when their final output was dispatched // above. This serves as a safety net to capture any diff --git a/rust/src/engine-core-client/src/metrics.rs b/rust/src/engine-core-client/src/metrics.rs index e15258d6081d..b392d5260185 100644 --- a/rust/src/engine-core-client/src/metrics.rs +++ b/rust/src/engine-core-client/src/metrics.rs @@ -8,8 +8,8 @@ use std::time::{SystemTime, UNIX_EPOCH}; use vllm_metrics::{ EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames, LoraInfoLabels, MooncakeOperationCounterFamily, MooncakeOperationHistogramFamily, - MooncakeOperationLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge, - WaitingReasonLabels, + MooncakeOperationLabels, RequestMetrics, SchedulerLogStatsAccumulator, SchedulerMetrics, + U64Counter, U64Gauge, WaitingReasonLabels, }; use crate::protocol::stats::{ @@ -20,6 +20,23 @@ use crate::transport::ConnectedEngine; const WAITING_REASON_CAPACITY: &str = "capacity"; const WAITING_REASON_DEFERRED: &str = "deferred"; +/// Cached output-batch metric handles for one model and engine index. +pub(crate) struct IterationMetricHandles { + pub iteration_tokens_total: HistogramMetric, +} + +impl IterationMetricHandles { + pub(crate) fn new(metrics: &RequestMetrics, model_name: &str, engine: u32) -> Self { + let labels = EngineLabels { + model_name: model_name.to_string(), + engine, + }; + Self { + iteration_tokens_total: metrics.iteration_tokens_total.get_or_create_owned(&labels), + } + } +} + /// Cached scheduler-stats metric handles for all engines connected to one /// frontend client. pub(crate) struct SchedulerStatsRecorder { diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 77d72b525fbd..958fcf8aa1b3 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -72,6 +72,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { weight_transfer_backend: None, enable_sleep_mode: false, supports_draft_weight_updates: false, + effective_attention_block_size: None, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index c14c19a50373..52d897be52dc 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -104,6 +104,9 @@ pub struct EngineCoreReadyResponse { /// Whether the engine has a speculative draft model that can be updated. #[serde(default)] pub supports_draft_weight_updates: bool, + /// Full-attention block size in tokens after initialization, or unavailable. + #[serde(default)] + pub effective_attention_block_size: Option, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index d5bc40e06159..646c106a04ed 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray}; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct TokenLogprob { pub token_id: u32, + /// Preserves the engine's value, including NaN and infinities. pub logprob: f32, /// The sampled/selected token uses its actual vocab rank. Remaining entries /// use 1-based top-k ranks matching the engine's returned candidate /// order. + /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's + /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values. pub rank: u32, } @@ -53,10 +56,6 @@ impl PositionLogprobs { logprobs.len() ); } - if sampled_rank == 0 { - bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs"); - } - let mut entries = Vec::with_capacity(token_ids.len()); for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() { let rank = if index == 0 { diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index 2470460b52a6..f69ff22f1ae0 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() { "new_logprobs: zero-column logprobs payload with 2 rows" ); } + +fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value { + let rows = ranks.len(); + let cols = ids.len() / rows; + Value::Array(vec![ + ndarray_value( + " { + self.num_preemptions += 1; self.handles.num_preemptions.inc(); } } @@ -276,6 +281,7 @@ fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetri .get_or_create_owned(&labels), request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels), request_params_n: metrics.request_params_n.get_or_create_owned(&labels), + request_num_preemptions: metrics.request_num_preemptions.get_or_create_owned(&labels), request_prefill_kv_computed_tokens: metrics .request_prefill_kv_computed_tokens .get_or_create_owned(&labels), diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 2ff3faa8d459..2d09aff889a8 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -790,6 +790,12 @@ async fn generate_records_request_metrics_in_prometheus_output() { assert!(rendered.contains(&format!( "vllm:num_preemptions_total{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"4\"}} 1.0" + ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" + ))); assert!(rendered.contains(&format!( "vllm:time_to_first_token_seconds_count{{model_name=\"{model_name}\",engine=\"4\"}} 1" ))); @@ -883,6 +889,12 @@ async fn dropping_stream_records_abort_terminal_request_metrics() { assert!(rendered.contains(&format!( "vllm:e2e_request_latency_seconds_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_sum{{model_name=\"{model_name}\",engine=\"5\"}} 0.0" + ))); + assert!(rendered.contains(&format!( + "vllm:request_num_preemptions_count{{model_name=\"{model_name}\",engine=\"5\"}} 1" + ))); llm.shutdown().await.unwrap(); } diff --git a/rust/src/managed-engine/src/cli.rs b/rust/src/managed-engine/src/cli.rs index f0cd7bf1715d..0a81ae598e9c 100644 --- a/rust/src/managed-engine/src/cli.rs +++ b/rust/src/managed-engine/src/cli.rs @@ -93,6 +93,7 @@ impl ManagedEngineArgs { shutdown_timeout: u64, handshake_port: u16, limit_mm_per_prompt: Option, + hf_overrides: Option, ) -> ManagedEngineConfig { let mut python_args = self.python_args; // Manually forward some args to the Python engine. @@ -136,6 +137,10 @@ impl ManagedEngineArgs { python_args.push("--limit-mm-per-prompt".to_string()); python_args.push(limit_mm_per_prompt); } + if let Some(hf_overrides) = hf_overrides { + python_args.push("--hf-overrides".to_string()); + python_args.push(hf_overrides); + } ManagedEngineConfig { python: self.python, diff --git a/rust/src/metrics/src/request.rs b/rust/src/metrics/src/request.rs index c1135781b23c..5bfc32e89534 100644 --- a/rust/src/metrics/src/request.rs +++ b/rust/src/metrics/src/request.rs @@ -21,6 +21,10 @@ const REQUEST_LATENCY_BUCKETS: [f64; 21] = [ 480.0, 960.0, 1920.0, 7680.0, ]; const REQUEST_PARAMS_N_BUCKETS: [f64; 5] = [1.0, 2.0, 5.0, 10.0, 20.0]; +const REQUEST_NUM_PREEMPTIONS_BUCKETS: [f64; 7] = [1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0]; +const ITERATION_TOKENS_BUCKETS: [f64; 13] = [ + 1.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0, 16384.0, +]; fn build_1_2_5_buckets(max_value: u32) -> Vec { let mut buckets = Vec::new(); @@ -66,6 +70,14 @@ fn request_params_n_histogram() -> Histogram { Histogram::new(REQUEST_PARAMS_N_BUCKETS.iter().copied()) } +fn request_num_preemptions_histogram() -> Histogram { + Histogram::new(REQUEST_NUM_PREEMPTIONS_BUCKETS.iter().copied()) +} + +fn iteration_tokens_histogram() -> Histogram { + Histogram::new(ITERATION_TOKENS_BUCKETS.iter().copied()) +} + #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] pub struct FinishedReasonLabels { pub model_name: String, @@ -83,8 +95,12 @@ pub struct PromptTokenSourceLabels { pub(crate) type FinishedReasonCounterFamily = Family; pub(crate) type PromptTokenSourceCounterFamily = Family; -/// Request-lifecycle Prometheus families exported from the `llm` layer. +/// Request and output-batch Prometheus families. pub struct RequestMetrics { + /// Computed prompt and generated tokens per nonempty engine output batch, + /// recorded by the engine-core client for requests tracked by this frontend. + pub iteration_tokens_total: HistogramFamily, + // Request-derived counters. pub num_preemptions: Family, pub prompt_tokens: Family, @@ -92,11 +108,6 @@ pub struct RequestMetrics { pub prompt_tokens_cached: Family, pub generation_tokens: Family, - // We intentionally don't support iteration-level histograms for now, since it seems to make - // more sense if the engine maintains these metrics and frontend simply forwards. - // - // pub iteration_tokens_total: HistogramFamily, - // Request lifecycle counters and histograms. pub request_success: FinishedReasonCounterFamily, pub request_prompt_tokens: HistogramFamily, @@ -104,6 +115,7 @@ pub struct RequestMetrics { pub request_max_num_generation_tokens: HistogramFamily, pub request_params_max_tokens: HistogramFamily, pub request_params_n: HistogramFamily, + pub request_num_preemptions: HistogramFamily, pub request_prefill_kv_computed_tokens: HistogramFamily, pub time_to_first_token_seconds: HistogramFamily, pub inter_token_latency_seconds: HistogramFamily, @@ -118,6 +130,14 @@ pub struct RequestMetrics { impl RequestMetrics { /// Register the request-oriented metric families into the shared registry. pub(crate) fn register(registry: &mut Registry) -> Self { + let iteration_tokens_total = + Family::new_with_constructor(iteration_tokens_histogram as fn() -> Histogram); + registry.register( + "vllm:iteration_tokens_total", + "Histogram of computed prompt and generated tokens per nonempty engine output batch.", + iteration_tokens_total.clone(), + ); + // Request-derived counters. let num_preemptions = Family::default(); registry.register( @@ -202,6 +222,14 @@ impl RequestMetrics { request_params_n.clone(), ); + let request_num_preemptions = + Family::new_with_constructor(request_num_preemptions_histogram as fn() -> Histogram); + registry.register( + "vllm:request_num_preemptions", + "Histogram of the number of times a request was preempted.", + request_num_preemptions.clone(), + ); + let request_prefill_kv_computed_tokens = Family::new_with_constructor(request_token_count_histogram as fn() -> Histogram); registry.register( @@ -276,6 +304,7 @@ impl RequestMetrics { ); Self { + iteration_tokens_total, num_preemptions, prompt_tokens, prompt_tokens_by_source, @@ -287,6 +316,7 @@ impl RequestMetrics { request_max_num_generation_tokens, request_params_max_tokens, request_params_n, + request_num_preemptions, request_prefill_kv_computed_tokens, time_to_first_token_seconds, inter_token_latency_seconds, diff --git a/rust/src/mock-engine/src/engine.rs b/rust/src/mock-engine/src/engine.rs index c21682915910..38f9b21be1d7 100644 --- a/rust/src/mock-engine/src/engine.rs +++ b/rust/src/mock-engine/src/engine.rs @@ -19,6 +19,7 @@ use vllm_engine_core_client::protocol::output::{ UtilityCallOutput, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::protocol::utility::{ EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope, }; @@ -187,6 +188,11 @@ impl ActiveRequest { /// Advance this request by one mock engine step. fn step(&mut self, opt: &Opt) -> EngineCoreOutput { + let prefill_stats = (self.generated == 0).then(|| PrefillStats { + num_prompt_tokens: self.prompt_len as u32, + num_computed_tokens: self.prompt_len as u32, + ..Default::default() + }); let remaining = self.max_tokens - self.generated; let chunk_len = remaining.min(opt.output_token_chunk_size); let mut new_token_ids = Vec::with_capacity(chunk_len); @@ -196,11 +202,14 @@ impl ActiveRequest { self.generated += chunk_len; let finished = self.generated >= self.max_tokens; - request_output( - self.request_id.clone(), - new_token_ids, - finished.then_some(EngineCoreFinishReason::Length), - ) + EngineCoreOutput { + prefill_stats, + ..request_output( + self.request_id.clone(), + new_token_ids, + finished.then_some(EngineCoreFinishReason::Length), + ) + } } } diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 9bfee620ec74..bee710fd0c15 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -10,6 +10,7 @@ use tokio_util::sync::CancellationToken; use vllm_engine_core_client::protocol::output::EngineCoreFinishReason; use vllm_engine_core_client::protocol::request::EngineCoreRequest; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::stats::PrefillStats; use vllm_engine_core_client::test_utils::IpcNamespace; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode}; @@ -119,12 +120,22 @@ async fn chunk_size_one_outputs_one_token_per_update() { let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 1); assert_eq!(first.finish_reason, None); + assert_eq!( + first.prefill_stats, + Some(PrefillStats { + num_prompt_tokens: 3, + num_computed_tokens: 3, + ..Default::default() + }) + ); let second = stream.next().await.expect("second").expect("second ok"); assert_eq!(second.new_token_ids.len(), 1); assert_eq!(second.finish_reason, None); + assert!(second.prefill_stats.is_none()); let third = stream.next().await.expect("third").expect("third ok"); assert_eq!(third.new_token_ids.len(), 1); assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length)); + assert!(third.prefill_stats.is_none()); assert!(stream.next().await.is_none()); shutdown_mock(client, shutdown, task).await; diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index fc9914e3e2f2..ad3872843a1f 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -61,6 +61,7 @@ async fn main() -> Result<()> { coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, revision: None, + hf_overrides: Default::default(), generation_config: Default::default(), served_model_name: vec![], listener_mode: HttpListenerMode::BindTcp { diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 0f0deac86d2c..4c337adbef8f 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -16,6 +16,7 @@ use vllm_chat::{ ChatTemplateContentFormatOption, GenerationConfigMode, ParserSelection, RendererSelection, }; use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode}; +use vllm_text::backend::hf::HfOverrides; /// Default keep-alive idle timeout (seconds); also the head-read bound /// when keep-alive is disabled (`0`). @@ -214,6 +215,8 @@ pub struct Config { pub model: String, /// Model revision on the Hugging Face Hub (branch, tag, or commit SHA). pub revision: Option, + /// JSON Merge Patch applied to the model config before backend initialization. + pub hf_overrides: HfOverrides, /// Which generation-config sampling defaults to inherit. pub generation_config: GenerationConfigMode, /// Model name(s) exposed to clients via the OpenAI API. When non-empty, diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index ddb74bf40f1d..de9c5ae65fb9 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -195,6 +195,7 @@ impl pb::control_server::Control for ControlServiceImpl { max_batched_tokens: ready.max_num_batched_tokens, max_loras: ready.max_loras, rl_capabilities: Some(self.rl_capabilities()), + effective_attention_block_size: self.client().effective_attention_block_size(), })) } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index f5f74d557e69..a21e013323d4 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1911,8 +1911,16 @@ async fn control_abort_resolves_external_id_and_empty_is_noop() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_reports_server_and_model_info() { + let mut ready = default_ready_response(); + ready.effective_attention_block_size = Some(64); let (generate_service, control_service, engine_health, _engine_task) = - setup_grpc_service(b"engine-grpc-info", default_stream_output_specs()).await; + setup_grpc_service_with_engine_script( + b"engine-grpc-info".to_vec(), + ready, + Arc::new(FakeTextBackend), + |_, _| boxed_test_future(async {}), + ) + .await; let (channel, server_task) = start_grpc_test_server( generate_service, control_service, @@ -1928,6 +1936,7 @@ async fn control_reports_server_and_model_info() { .expect("get server info") .into_inner(); assert_eq!(server.engine_version, "test-vllm-version"); + assert_eq!(server.effective_attention_block_size, Some(64)); assert_eq!(server.api_version, "vllm"); assert_eq!(server.instance_id, "test-instance"); assert_eq!(server.max_model_len, DEFAULT_MOCK_MAX_MODEL_LEN as u32); @@ -2195,6 +2204,7 @@ async fn control_aggregates_multi_engine_capacity() { ready_0.weight_transfer_backend = Some("nccl".to_string()); ready_0.enable_sleep_mode = true; ready_0.supports_draft_weight_updates = true; + ready_0.effective_attention_block_size = Some(64); let mut ready_1 = default_ready_response(); ready_1.max_model_len = 4_096; @@ -2248,6 +2258,7 @@ async fn control_aggregates_multi_engine_capacity() { .into_inner(); assert_eq!(server.max_model_len, 4_096); assert_eq!(server.total_kv_blocks, 30); + assert!(server.effective_attention_block_size.is_none()); let rl = server.rl_capabilities.expect("RL capabilities"); assert!(!rl.weight_transfer_enabled); assert!(rl.weight_transfer_backend.is_empty()); diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 95ae359f35b3..49d3d57f3264 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -99,6 +99,7 @@ async fn build_state(config: &Config) -> Result> { &config.model, LoadModelBackendsOptions { revision: config.revision.clone(), + hf_overrides: config.hf_overrides.clone(), generation_config: config.generation_config, renderer: config.renderer, language_model_only: config.language_model_only, diff --git a/rust/src/server/src/render.rs b/rust/src/server/src/render.rs index fe7033fa8cad..4fc55fe60be9 100644 --- a/rust/src/server/src/render.rs +++ b/rust/src/server/src/render.rs @@ -13,6 +13,7 @@ use vllm_chat::{ ParserSelection, RendererSelection, load_model_backends, }; use vllm_text::TextRequestProcessor; +use vllm_text::backend::hf::HfOverrides; use crate::{ HttpListenerMode, TlsConfig, @@ -25,6 +26,7 @@ use crate::{ pub struct RenderConfig { pub model: String, pub revision: Option, + pub hf_overrides: HfOverrides, pub served_model_name: Vec, pub host: String, pub port: u16, @@ -69,6 +71,7 @@ async fn build_state(config: &RenderConfig) -> Result> { &config.model, LoadModelBackendsOptions { revision: config.revision.clone(), + hf_overrides: config.hf_overrides.clone(), generation_config: Default::default(), renderer: config.renderer, language_model_only: true, @@ -153,6 +156,7 @@ mod tests { RenderConfig { model: "test-model".to_string(), revision: None, + hf_overrides: Default::default(), served_model_name: Vec::new(), host: "127.0.0.1".to_string(), port: 8000, diff --git a/rust/src/server/src/routes/openai/utils/logprobs.rs b/rust/src/server/src/routes/openai/utils/logprobs.rs index 66c7a8fad9b8..1de2cdc10015 100644 --- a/rust/src/server/src/routes/openai/utils/logprobs.rs +++ b/rust/src/server/src/routes/openai/utils/logprobs.rs @@ -369,4 +369,51 @@ mod tests { assert_eq!(chat_top_logprobs_len(1), 1); assert_eq!(chat_top_logprobs_len(-1), 3); } + + #[test] + fn chat_logprobs_clamps_nan_and_negative_infinity_for_json() { + let mut logprobs = sample_logprobs(); + logprobs.positions[0].entries[0].logprob = f32::NAN; + logprobs.positions[0].entries[0].rank = 0; + logprobs.positions[0].entries[1].logprob = f32::NEG_INFINITY; + + let response = decoded_logprobs_to_openai_chat(&logprobs, -1, false).unwrap(); + let json = serde_json::to_string_pretty(&response).unwrap(); + expect_test::expect![[r#" + { + "content": [ + { + "token": "A", + "logprob": -9999.0, + "bytes": [ + 65 + ], + "top_logprobs": [ + { + "token": "A", + "logprob": -9999.0, + "bytes": [ + 65 + ] + }, + { + "token": "B", + "logprob": -9999.0, + "bytes": [ + 66 + ] + }, + { + "token": "C", + "logprob": -2.0, + "bytes": [ + 67 + ] + } + ] + } + ] + }"#]] + .assert_eq(&json); + } } diff --git a/rust/src/text/Cargo.toml b/rust/src/text/Cargo.toml index 7bda7f976e1c..36152d7cc4a1 100644 --- a/rust/src/text/Cargo.toml +++ b/rust/src/text/Cargo.toml @@ -16,6 +16,7 @@ reqwest.workspace = true serde.workspace = true serde_json.workspace = true serde_with.workspace = true +tempfile.workspace = true thiserror.workspace = true thiserror-ext.workspace = true tracing.workspace = true @@ -28,7 +29,6 @@ vllm-tokenizer.workspace = true expect-test.workspace = true futures.workspace = true serial_test.workspace = true -tempfile.workspace = true tokio.workspace = true vllm-llm = { workspace = true, features = ["test-util"] } vllm-tokenizer = { workspace = true, features = ["test-utils"] } diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index ae768767d99a..2f32f2cb483d 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -3,6 +3,7 @@ mod config; mod model_files; +mod overrides; use std::collections::BTreeSet; use std::sync::Arc; @@ -18,6 +19,7 @@ pub use self::config::{ load_tokenizer_config, }; pub use self::model_files::{ResolvedModelFiles, TokenizerSource}; +pub use self::overrides::HfOverrides; use crate::backend::{GenerationConfigMode, SamplingHints, TextBackend}; use crate::error::Result; diff --git a/rust/src/text/src/backend/hf/model_files.rs b/rust/src/text/src/backend/hf/model_files.rs index de802fb4390c..3ffd7b2bb79c 100644 --- a/rust/src/text/src/backend/hf/model_files.rs +++ b/rust/src/text/src/backend/hf/model_files.rs @@ -2,11 +2,13 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::path::{Path, PathBuf}; +use std::sync::Arc; use hf_hub::api::tokio::{Api, ApiBuilder, ApiRepo}; use hf_hub::{Cache, Repo, RepoType}; use thiserror_ext::AsReport as _; +use super::HfOverrides; use super::config::{HfTokenizerConfig, load_tokenizer_config}; use crate::error::{Error, Result}; @@ -50,9 +52,27 @@ pub struct ResolvedModelFiles { pub processor_config_path: Option, pub chat_template_path: Option, pub config_path: Option, + /// Keep a patched config alive across clones and backend ownership. + config_temp_path: Option>, } impl ResolvedModelFiles { + /// Materialize a merge-patched config owned by these resolved model files. + pub fn apply_overrides(&mut self, overrides: &HfOverrides) -> Result<()> { + if overrides.is_empty() { + return Ok(()); + } + let path = materialize_config(self.config_path.as_deref(), overrides).map_err(|error| { + Error::Tokenizer(format!( + "failed to materialize HF overrides: {}", + error.as_report() + )) + })?; + self.config_path = Some(path.to_path_buf()); + self.config_temp_path = Some(Arc::new(path)); + Ok(()) + } + /// Resolve tokenizer/config files from a local model directory first when /// `model_id` points to one, otherwise consult the local HF cache and /// finally the Hub. @@ -84,6 +104,24 @@ impl ResolvedModelFiles { } } +fn materialize_config( + path: Option<&Path>, + overrides: &HfOverrides, +) -> anyhow::Result { + let mut config: serde_json::Value = match path { + Some(path) => serde_json::from_slice(&std::fs::read(path)?)?, + None => serde_json::json!({}), + }; + anyhow::ensure!(config.is_object(), "HF model config must be a JSON object"); + overrides.apply(&mut config); + let file = tempfile::Builder::new() + .prefix("vllm-rs-hf-config-") + .suffix(".json") + .tempfile()?; + serde_json::to_writer(file.as_file(), &config)?; + Ok(file.into_temp_path()) +} + fn resolve_local_model_files(model_dir: &Path) -> Result { let tokenizer_config_path = local_file_if_exists(model_dir, "tokenizer_config.json"); let tokenizer_config = load_tokenizer_config(tokenizer_config_path.as_deref())?; @@ -101,6 +139,7 @@ fn resolve_local_model_files(model_dir: &Path) -> Result { processor_config_path: local_file_if_exists(model_dir, "processor_config.json"), chat_template_path: discover_chat_template_in_dir(model_dir), config_path: local_file_if_exists(model_dir, "config.json"), + config_temp_path: None, }) } @@ -164,6 +203,7 @@ async fn resolve_remote_model_files( processor_config_path, chat_template_path, config_path, + config_temp_path: None, }) } @@ -204,6 +244,7 @@ fn resolve_cached_model_files(cache: &Cache, repo: &Repo) -> Result); + +impl HfOverrides { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// Apply the patch to an in-memory config, preserving the source file. + pub fn apply(&self, config: &mut Value) { + merge_object(config, &self.0); + } +} + +/// Apply RFC 7396 object-member updates at every nesting level. +fn merge_object(target: &mut Value, patch: &Map) { + if !target.is_object() { + *target = Value::Object(Map::new()); + } + let target = target.as_object_mut().unwrap(); + for (key, value) in patch { + match value { + Value::Null => { + target.remove(key); + } + Value::Object(object) => { + merge_object(target.entry(key.clone()).or_insert(Value::Null), object) + } + _ => { + target.insert(key.clone(), value.clone()); + } + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::HfOverrides; + + #[test] + fn merge_patch_preserves_siblings_replaces_arrays_and_deletes_null_members() { + let mut config = json!({ + "rope_parameters": {"rope_type": "yarn", "factor": 2}, + "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 2}}, + "architectures": ["Old", "Other"], + "sliding_window": 4096 + }); + let patch: HfOverrides = serde_json::from_value(json!({ + "rope_parameters": {"factor": 4}, + "text_config": {"rope_parameters": {"factor": 4}}, + "architectures": ["New"], + "sliding_window": null, + "missing": null + })) + .unwrap(); + patch.apply(&mut config); + assert_eq!( + config, + json!({ + "rope_parameters": {"rope_type": "yarn", "factor": 4}, + "text_config": {"rope_parameters": {"rope_type": "yarn", "factor": 4}}, + "architectures": ["New"] + }) + ); + } + + #[test] + fn merge_patch_creates_objects_and_preserves_nulls_inside_replacement_arrays() { + let mut config = json!({"a": 1, "b": [1, 2]}); + let patch: HfOverrides = serde_json::from_value(json!({ + "a": {"nested": {"removed": null}}, + "b": {}, + "c": [null, {"literal": null}] + })) + .unwrap(); + patch.apply(&mut config); + assert_eq!( + config, + json!({"a": {"nested": {}}, "b": {}, "c": [null, {"literal": null}]}) + ); + } +} diff --git a/setup.py b/setup.py index 13692f3d20ca..4012000e634e 100644 --- a/setup.py +++ b/setup.py @@ -1257,35 +1257,39 @@ def get_vllm_version() -> str: version = get_version(write_to="vllm/_version.py") sep = "+" if "+" not in version else "." # dev versions might contain + - if _no_device(): - if envs.VLLM_TARGET_DEVICE == "empty": - version += f"{sep}empty" - elif _is_cuda(): - if USE_PRECOMPILED_EXTENSIONS and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX: - version += f"{sep}precompiled" + if not envs.VLLM_SKIP_VERSION_SUFFIX: + if _no_device(): + if envs.VLLM_TARGET_DEVICE == "empty": + version += f"{sep}empty" + elif _is_cuda(): + if ( + USE_PRECOMPILED_EXTENSIONS + and not envs.VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX + ): + version += f"{sep}precompiled" + else: + cuda_version = str(get_nvcc_cuda_version()) + if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: + cuda_version_str = cuda_version.replace(".", "")[:3] + # skip this for source tarball, required for pypi + if "sdist" not in sys.argv: + version += f"{sep}cu{cuda_version_str}" + elif _is_hip(): + # Get the Rocm Version + rocm_version = get_rocm_version() or torch.version.hip + if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: + version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" + elif _is_tpu(): + version += f"{sep}tpu" + elif _is_cpu(): + # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), + # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. + if VLLM_TARGET_DEVICE == "cpu": + version += f"{sep}cpu" + elif _is_xpu(): + version += f"{sep}xpu" else: - cuda_version = str(get_nvcc_cuda_version()) - if cuda_version != envs.VLLM_MAIN_CUDA_VERSION: - cuda_version_str = cuda_version.replace(".", "")[:3] - # skip this for source tarball, required for pypi - if "sdist" not in sys.argv: - version += f"{sep}cu{cuda_version_str}" - elif _is_hip(): - # Get the Rocm Version - rocm_version = get_rocm_version() or torch.version.hip - if rocm_version and rocm_version != envs.VLLM_MAIN_CUDA_VERSION: - version += f"{sep}rocm{rocm_version.replace('.', '')[:3]}" - elif _is_tpu(): - version += f"{sep}tpu" - elif _is_cpu(): - # Check the local VLLM_TARGET_DEVICE (may be set by auto-detect above), - # not envs.VLLM_TARGET_DEVICE, so CPU-only hosts still get `+cpu`. - if VLLM_TARGET_DEVICE == "cpu": - version += f"{sep}cpu" - elif _is_xpu(): - version += f"{sep}xpu" - else: - raise RuntimeError("Unknown runtime environment") + raise RuntimeError("Unknown runtime environment") return version diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 94060d1724a8..590e97c1eb54 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -215,6 +215,9 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): base_hash = CacheConfig().compute_hash() assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash + config = CacheConfig() + config.effective_attention_block_size = 64 + assert config.compute_hash() == base_hash def test_scheduler_config_hash_includes_max_num_seqs(): diff --git a/tests/distributed/test_sharded_rdt_producer.py b/tests/distributed/test_sharded_rdt_producer.py index 62fc63876262..5cfab1957d65 100644 --- a/tests/distributed/test_sharded_rdt_producer.py +++ b/tests/distributed/test_sharded_rdt_producer.py @@ -20,17 +20,23 @@ import contextlib import gc +import os +import sys import threading import time from concurrent.futures import ThreadPoolExecutor from concurrent.futures import TimeoutError as FuturesTimeout +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock import pytest import torch import vllm.distributed.weight_transfer.sharded_rdt_trainer as trainer_mod +from vllm.distributed.nixl_utils import alias_nixl_for_ray from vllm.distributed.weight_transfer.sharded_rdt_common import ( buffer_alloc_bytes, + register_nixl_memory, ) from vllm.distributed.weight_transfer.sharded_rdt_trainer import ( DEFAULT_GATHER_LOOKAHEAD, @@ -45,6 +51,114 @@ GROUP_B = ("model.layers.1.w",) +@pytest.fixture +def ray_nixl_transport(monkeypatch): + import ray + from ray.experimental.rdt import util + from ray.experimental.rdt.nixl_tensor_transport import NixlTensorTransport + + from vllm.distributed import nixl_utils + from vllm.platforms import current_platform + + transport = NixlTensorTransport() + package = ModuleType("nixl_rocm") + api = ModuleType("nixl_rocm._api") + api.__dict__.update(nixl_agent=Mock(), nixl_agent_config=Mock()) + package.__dict__["_api"] = api + modules = {package.__name__: package, api.__name__: api} + imported_rcache_limits = [] + + def import_nixl(name): + imported_rcache_limits.append(os.environ.get("UCX_RCACHE_MAX_UNRELEASED")) + monkeypatch.setitem(sys.modules, name, modules[name]) + return modules[name] + + nixl_names = ("nixl", "nixl._api", *modules) + for name in nixl_names: + monkeypatch.delitem(sys.modules, name, raising=False) + monkeypatch.delenv("UCX_RCACHE_MAX_UNRELEASED", raising=False) + importer = Mock(side_effect=import_nixl) + monkeypatch.setattr( + nixl_utils, "importlib", SimpleNamespace(import_module=importer) + ) + monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: True) + monkeypatch.setattr(util, "transport_managers", {"NIXL": transport}) + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + monkeypatch.setattr( + ray, + "get_runtime_context", + lambda: SimpleNamespace(get_actor_id=lambda: "rdt-actor"), + ) + try: + yield transport, api, importer, imported_rcache_limits + finally: + # The helper inserts aliases directly; remove them before monkeypatch + # restores any modules and environment values present before this test. + for name in nixl_names: + sys.modules.pop(name, None) + os.environ.pop("UCX_RCACHE_MAX_UNRELEASED", None) + + +def test_ray_registers_memory_through_rocm_alias(ray_nixl_transport): + """Ray's registration must reach the ROCm agent through the alias.""" + _, api, importer, _ = ray_nixl_transport + register_nixl_memory(torch.empty(4, dtype=torch.uint8)) + api.nixl_agent.return_value.register_memory.assert_called() + assert sys.modules["nixl"] is sys.modules["nixl_rocm"] + assert sys.modules["nixl._api"] is api + + importer.reset_mock() + alias_nixl_for_ray() + importer.assert_not_called() + + +@pytest.mark.parametrize("rcache_limit", [None, "4096"]) +def test_ray_nixl_sets_ucx_limit_before_import( + monkeypatch, ray_nixl_transport, rcache_limit +): + """UCX must see the default or the user's override during the first import.""" + _, _, _, imported_rcache_limits = ray_nixl_transport + if rcache_limit is not None: + monkeypatch.setenv("UCX_RCACHE_MAX_UNRELEASED", rcache_limit) + alias_nixl_for_ray() + assert imported_rcache_limits == [rcache_limit or "1024"] * 2 + + +@pytest.mark.parametrize("existing_api", [False, True]) +def test_ray_nixl_preserves_non_rocm_or_existing_imports( + monkeypatch, ray_nixl_transport, existing_api +): + from vllm.platforms import current_platform + + _, _, importer, _ = ray_nixl_transport + if existing_api: + api = ModuleType("nixl._api") + monkeypatch.setitem(sys.modules, "nixl._api", api) + else: + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + alias_nixl_for_ray() + importer.assert_not_called() + assert "UCX_RCACHE_MAX_UNRELEASED" not in os.environ + if existing_api: + assert sys.modules["nixl._api"] is api + else: + assert "nixl._api" not in sys.modules + + +def test_ray_nixl_alias_reports_missing_rocm_package(monkeypatch, ray_nixl_transport): + from vllm.distributed import nixl_utils + + transport, api, importer, _ = ray_nixl_transport + monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: False) + with pytest.raises(ImportError, match="nixl_rocm"): + register_nixl_memory(torch.empty(1)) + assert transport._nixl_agent is None + api.nixl_agent.assert_not_called() + importer.assert_not_called() + assert "nixl" not in sys.modules + assert "nixl._api" not in sys.modules + + @pytest.fixture def server_factory(monkeypatch): """Build real `_RDTProducerServer` instances without a GPU. diff --git a/tests/entrypoints/cli/test_serve.py b/tests/entrypoints/cli/test_serve.py new file mode 100644 index 000000000000..9ba9ac7455c3 --- /dev/null +++ b/tests/entrypoints/cli/test_serve.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the ``vllm serve`` CLI subcommand.""" + +import argparse +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.entrypoints.cli.serve import run_headless + + +class _StopAfterEngineConfig(Exception): + pass + + +def test_headless_imports_reasoning_parser_plugin_before_engine_config(): + plugin_path = "/tmp/custom_reasoning_parser.py" + args = argparse.Namespace( + api_server_count=0, + reasoning_parser_plugin=plugin_path, + ) + engine_args = MagicMock() + + with ( + patch( + "vllm.entrypoints.cli.serve.vllm.AsyncEngineArgs.from_cli_args", + return_value=engine_args, + ), + patch( + "vllm.reasoning.ReasoningParserManager.import_reasoning_parser" + ) as import_plugin, + ): + + def create_engine_config(**kwargs): + import_plugin.assert_called_once_with(plugin_path) + raise _StopAfterEngineConfig + + engine_args.create_engine_config.side_effect = create_engine_config + + with pytest.raises(_StopAfterEngineConfig): + run_headless(args) diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 40fefc237404..668f11544858 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -5,6 +5,7 @@ import socket import threading import time +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -14,6 +15,7 @@ exit_before_report_worker, ) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, get_engine_client_zmq_addr, @@ -32,6 +34,18 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None): print("Mock worker completed successfully") +def update_admission_stats_worker(listen_address, sock, args, client_config): + stats = SharedAdmissionStats( + client_config, + client_config["client_count"], + client_config["client_index"], + ) + value = client_config["client_index"] + 1 + stats.set_num_requests(value) + args.barrier.wait() + args.results.put(stats.get_num_requests()) + + # Module-level stub for the gather_actual_addresses test. Must be # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): @@ -131,6 +145,31 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +@pytest.mark.skip_global_cleanup +def test_api_server_processes_share_admission_stats(api_server_args): + args = api_server_args.copy() + spawn_context = multiprocessing.get_context("spawn") + worker_args = SimpleNamespace( + barrier=spawn_context.Barrier(args["num_servers"]), + results=spawn_context.Queue(), + max_num_queued_reqs=10, + ) + args["args"] = worker_args + args["target_server_fn"] = update_admission_stats_worker + manager = APIServerProcessManager(**args) + + try: + for process in manager.processes: + process.join(timeout=30) + assert process.exitcode == 0 + + assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3 + finally: + manager.shutdown() + worker_args.results.close() + worker_args.results.join_thread() + + @patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker) def test_wait_for_completion_or_failure(api_server_args): """Test that wait_for_completion_or_failure works with failures.""" diff --git a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh index b99b1c3ca5b0..462ada22d712 100755 --- a/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh +++ b/tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh @@ -181,8 +181,9 @@ run_scale_out_ec() { # Encode-only EC producer instance. It runs the vision encoder and # publishes embeddings through the ECExampleConnector shared storage. echo "Starting encode instance on GPU $GPU_E, port $ENCODE_PORT" - env CUDA_VISIBLE_DEVICES="$GPU_E" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + env CUDA_VISIBLE_DEVICES="$GPU_E" \ vllm serve "$MODEL" \ + --enable-scale-out \ --port "$ENCODE_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ @@ -204,8 +205,9 @@ run_scale_out_ec() { # features plus ec_transfer_params and loads the embeddings that the # encode instance published. echo "Starting prefill instance on GPU $GPU_PD, port $PREFILL_PORT" - env CUDA_VISIBLE_DEVICES="$GPU_PD" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \ + env CUDA_VISIBLE_DEVICES="$GPU_PD" \ vllm serve "$MODEL" \ + --enable-scale-out \ --port "$PREFILL_PORT" \ --max-model-len "$MAX_MODEL_LEN" \ --enforce-eager \ diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml new file mode 100644 index 000000000000..009871925925 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml @@ -0,0 +1,20 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +tolerance: 0.0 +num_questions: 1319 +num_fewshot: 5 +max_tokens: 256 +max_concurrency: 32 +request_timeout_seconds: 1800 +startup_max_wait_seconds: 1200 +server_args: >- + --tensor-parallel-size 4 + --max-model-len 8192 + --max-num-seqs 32 + --max-num-batched-tokens 4096 + --enforce-eager + --kv-cache-dtype bfloat16 + --attention-config '{"hisparse_config":{}}' + --kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":32}}' +env: + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml index f57467ac7418..8f3e9ce918ef 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -3,14 +3,18 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# Preserve KV-cache headroom after CUDA graph memory profiling and bound the +# DCP-aware FlashInfer workspace to 2.39 GiB at the 4,096-token scheduler cap. server_args: >- --compilation-config '{"cudagraph_mode":"PIECEWISE"}' + --gpu-memory-utilization 0.96 --max-model-len 4096 - --max-num-batched-tokens 32768 + --max-num-batched-tokens 4096 --safetensors-load-strategy prefetch --moe-backend flashinfer_cutlass --prefill-context-parallel-size 4 --decode-context-parallel-size 4 + --dcp-comm-backend ag_rs --enable-expert-parallel --kv-cache-dtype fp8 env: diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml index 72f7fdf1ae91..bc5c11fa53e2 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml @@ -3,8 +3,10 @@ accuracy_threshold: 0.90 num_questions: 1319 num_fewshot: 5 max_concurrency: 100 +# Leave headroom for FlashInfer CUTLASS autotuning after KV-cache allocation. server_args: >- --compilation-config '{"cudagraph_mode":"NONE"}' + --gpu-memory-utilization 0.85 --max-model-len 4096 --max-num-batched-tokens 32768 --safetensors-load-strategy prefetch diff --git a/tests/evals/gsm8k/configs/models-hisparse.txt b/tests/evals/gsm8k/configs/models-hisparse.txt new file mode 100644 index 000000000000..53ff181f4f66 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-hisparse.txt @@ -0,0 +1 @@ +GLM-5.2-NVFP4-HiSparse.yaml diff --git a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py index d67df8ee2b4d..d5c37724aacd 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py @@ -30,7 +30,10 @@ def reset_aiter_mla_support_cache(monkeypatch: pytest.MonkeyPatch) -> None: def _install_fake_aiter_modules( - monkeypatch: pytest.MonkeyPatch, *, supports_fp8: bool + monkeypatch: pytest.MonkeyPatch, + *, + supports_fp8: bool, + supports_causal: bool = False, ) -> None: aiter_mod: Any = types.ModuleType("aiter") mla_mod: Any = types.ModuleType("aiter.mla") @@ -65,6 +68,14 @@ def mla_decode_fwd_without_fp8( mla_decode_fwd = mla_decode_fwd_without_fp8 + if supports_causal: + inner = mla_decode_fwd + + def mla_decode_fwd_with_causal(*args, causal=True, **kwargs): + return inner(*args, **kwargs) + + mla_decode_fwd = mla_decode_fwd_with_causal + mla_mod.mla_decode_fwd = mla_decode_fwd aiter_mod.mla = mla_mod @@ -135,3 +146,107 @@ def test_aiter_mla_fp8_support_result_is_cached(monkeypatch): assert _check_aiter_mla_fp8_support() is True assert _check_aiter_mla_fp8_support() is True assert signature_mock.call_count == 1 + + +def test_missing_causal_arg_fails_closed(monkeypatch): + """Old aiter wheels have no causal=; do not report that as a valid draft path.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=False) + with pytest.raises(RuntimeError, match="causal-only"): + rocm_aiter_ops.mla_decode_supports_non_causal() + with pytest.raises(RuntimeError, match="causal-only"): + AiterMLABackend.supports_non_causal() + + +@pytest.mark.parametrize("has_kernels", [False, True]) +def test_supports_non_causal_requires_arch_kernels(monkeypatch, has_kernels): + """The Python argument alone does not guarantee an arch has the kernel.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.v1.attention.backends.mla import rocm_aiter_mla + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True) + monkeypatch.setattr( + rocm_aiter_mla, + "_aiter_mla_non_causal_asm_kernels", + lambda: has_kernels, + ) + assert rocm_aiter_ops.mla_decode_supports_non_causal() is True + assert AiterMLABackend.supports_non_causal() is has_kernels + + +@pytest.mark.parametrize("causal", [True, False]) +def test_decode_fwd_forwards_causal_to_aiter(monkeypatch, causal): + """The decode op must pass the block mask through; omitting it is causal-only.""" + seen: list = [] + + def mla_decode_fwd(*args, causal=None, **kwargs): + seen.append(causal) + + mla_mod: Any = types.ModuleType("aiter.mla") + mla_mod.mla_decode_fwd = mla_decode_fwd + aiter_mod: Any = types.ModuleType("aiter") + aiter_mod.mla = mla_mod + monkeypatch.setitem(sys.modules, "aiter", aiter_mod) + monkeypatch.setitem(sys.modules, "aiter.mla", mla_mod) + + import torch + + import vllm._aiter_ops as aiter_ops + + aiter_ops._rocm_aiter_mla_decode_fwd_impl( + torch.zeros(2, 1, 8), + torch.zeros(2, 1, 1, 8), + torch.zeros(2, 1, 8), + torch.zeros(3, dtype=torch.int32), + 1, + causal=causal, + ) + assert seen == [causal] + + +def _aiter_mla_validate( + monkeypatch, + *, + dtype, + kv_cache_dtype, + use_non_causal, +): + from vllm.platforms.interface import DeviceCapability + from vllm.v1.attention.backend import AttentionType + from vllm.v1.attention.backends.mla import rocm_aiter_mla + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLABackend + + _install_fake_aiter_modules(monkeypatch, supports_fp8=True, supports_causal=True) + monkeypatch.setattr( + rocm_aiter_mla, "_aiter_mla_non_causal_asm_kernels", lambda: True + ) + return AiterMLABackend.validate_configuration( + head_size=576, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + block_size=1, + use_mla=True, + has_sink=False, + use_sparse=False, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(9, 5), + attn_type=AttentionType.DECODER, + use_non_causal=use_non_causal, + ) + + +def test_non_causal_fp16_auto_cache_is_rejected(monkeypatch): + """Pinned AITER aborts fp16 Q; auto-select must fall through to Triton.""" + import torch + + reasons = _aiter_mla_validate( + monkeypatch, + dtype=torch.float16, + kv_cache_dtype="auto", + use_non_causal=True, + ) + assert any("non-causal fp16" in reason for reason in reasons) diff --git a/tests/kernels/moe/test_deepep_v2_async_finalize.py b/tests/kernels/moe/test_deepep_v2_async_finalize.py new file mode 100644 index 000000000000..193d8906f37c --- /dev/null +++ b/tests/kernels/moe/test_deepep_v2_async_finalize.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Contract tests for DeepEPV2PrepareAndFinalize.finalize_async. + +Hermetic: fake buffer/event, single process, no GPU collectives — only +requires deep_ep to be importable. + + T1 finalize_async -> combine is issued with async_with_compute_stream=True + and the join is deferred to the receiver, which waits + on the combine event (device-side) before the copy. + T2 DBO active -> combine falls back to async_with_compute_stream=False; + the receiver still produces the correct output. + T3 finalize (sync) -> combine is issued synchronously and the output is + copied immediately. +""" + +import pytest +import torch + +from vllm.utils.import_utils import has_deep_ep_v2 + +requires_deep_ep_v2 = pytest.mark.skipif( + not has_deep_ep_v2(), + reason="Requires DeepEP v2 (ElasticBuffer)", +) + +if has_deep_ep_v2(): + import vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 as _dv2 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + DeepEPV2PrepareAndFinalize, + ) + from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + ) + + +class _FakeEvent: + def __init__(self, has_event: bool): + self.event = object() if has_event else None + self.waited = False + + def current_stream_wait(self): + assert self.event is not None + self.waited = True + + +class _FakeBuffer: + def __init__(self, out: torch.Tensor): + self.out = out + self.calls: list[dict] = [] + self.last_event: _FakeEvent | None = None + + def combine( + self, + x, + handle, + topk_weights, + async_with_compute_stream, + allocate_on_comm_stream, + ): + self.calls.append({"async_with_compute_stream": async_with_compute_stream}) + # DeepEP returns a completion event only for an async combine. + self.last_event = _FakeEvent(has_event=async_with_compute_stream) + return self.out, None, self.last_event + + +def _make_pf(out: torch.Tensor): + pf = DeepEPV2PrepareAndFinalize( + buffer=_FakeBuffer(out), + num_dispatchers=1, + dp_size=1, + rank_expert_offset=0, + num_experts=8, + num_topk=2, + ) + pf.handles[0] = object() + return pf + + +def _run(pf, output: torch.Tensor, do_async: bool): + empty = torch.empty(0, dtype=torch.bfloat16) + weights = torch.empty(0, 2) + ids = torch.empty(0, 2, dtype=torch.int64) + fn = pf.finalize_async if do_async else pf.finalize + return fn(output, empty, weights, ids, False, TopKWeightAndReduceContiguous()) + + +@requires_deep_ep_v2 +def test_finalize_async_defers_join_to_receiver(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is True + assert not torch.equal(dst, out_src), "join must not happen before receiver" + recv() + assert pf.buffer.last_event.waited, "receiver must join via event wait" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_async_under_dbo_uses_sync_combine(monkeypatch): + monkeypatch.setattr(_dv2, "dbo_enabled", lambda: True) + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + recv() + assert not pf.buffer.last_event.waited, "no event to wait on for sync combine" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_sync(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + ret = _run(pf, dst, do_async=False) + assert ret is None + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + assert torch.equal(dst, out_src) diff --git a/tests/kernels/moe/test_rocm_aiter_moe.py b/tests/kernels/moe/test_rocm_aiter_moe.py index 89244bed1b0f..0967f9499e0f 100644 --- a/tests/kernels/moe/test_rocm_aiter_moe.py +++ b/tests/kernels/moe/test_rocm_aiter_moe.py @@ -565,6 +565,74 @@ def test_aiter_moe_shared_experts_enablement_follows_env( assert rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() is expected +def test_aiter_moe_situv2_syncs_aiter_a4w4_env( + monkeypatch: pytest.MonkeyPatch, +): + """SiTUv2 MoE should route to AITER a4w4 and clear legacy a8w4 overrides.""" + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.delenv("AITER_SITUV2_A8W4", raising=False) + mp.delenv("AITER_SITUV2_A4W4", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert os.environ.get("AITER_SITUV2_A4W4") == "1" + assert "AITER_SITUV2_A8W4" not in os.environ + + +def test_aiter_moe_situv2_legacy_a8w4_alias_enables_a4w4( + monkeypatch: pytest.MonkeyPatch, +): + """Deprecated VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 still enables SiTUv2 a4w4.""" + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.delenv("AITER_SITUV2_A8W4", raising=False) + mp.delenv("AITER_SITUV2_A4W4", raising=False) + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "1") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert os.environ.get("AITER_SITUV2_A4W4") == "1" + assert "AITER_SITUV2_A8W4" not in os.environ + + +def test_aiter_moe_situv2_clears_aiter_env_when_disabled( + monkeypatch: pytest.MonkeyPatch, +): + from vllm._aiter_ops import rocm_aiter_ops + + _assert_aiter_supported() + + with monkeypatch.context() as mp: + mp.setenv("AITER_SITUV2_A4W4", "1") + mp.setenv("VLLM_ROCM_USE_AITER", "1") + mp.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + mp.delenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", raising=False) + mp.setenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0") + _reload_envs() + rocm_aiter_ops.refresh_env_variables() + + import os + + assert "AITER_SITUV2_A4W4" not in os.environ + + @pytest.mark.parametrize("moe_padding", [True, False]) def test_aiter_moe_padding_env_var( moe_padding: bool, diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index fd530e7ffb53..d02527aba048 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -27,6 +27,7 @@ from vllm.utils.deep_gemm import ( fp8_gemm_nt, get_tma_aligned_size, + is_deep_gemm_e8m0_used, per_block_cast_to_fp8, should_use_deepgemm_for_fp8_linear, ) @@ -291,7 +292,11 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed): A_fp8, As_fp8 = per_token_group_quant_fp8( A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True ) - B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size) + B_fp8, Bs_fp8 = per_block_cast_to_fp8( + B_fp32, + block_size=block_size, + use_ue8m0=is_deep_gemm_e8m0_used(), + ) As = As_fp8.to(torch.float32) Bs = Bs_fp8.to(torch.float32) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 0f1976cf683e..c1de94be92fc 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -42,6 +42,122 @@ from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize( + "cache_dtype,kv_mxfp8", + [ + (torch.uint8, False), + (torch.uint8, True), + (torch.bfloat16, False), + (torch.float8_e4m3fn, False), + ], +) +@pytest.mark.parametrize("num_tokens", [1, 17, 1023, 1024]) +@pytest.mark.parametrize("use_graph", [False, True]) +def test_dspark_context_kv_matches_query_insert( + cache_dtype, kv_mxfp8, num_tokens, use_graph +): + """KV-only insertion must preserve every cache byte, including graph replay.""" + from vllm.models.deepseek_v41.nvidia.dspark import _insert_context_kv + + torch.manual_seed(42) + block_size = 256 + num_blocks = math.ceil((num_tokens + 7) / block_size) + row_size = (528 if kv_mxfp8 else 584) if cache_dtype == torch.uint8 else 512 + page_stride = math.ceil(block_size * row_size / 576) * 576 + 576 + backing = torch.full((num_blocks, page_stride), 3, device="cuda").to(cache_dtype) + cache = backing.as_strided( + (num_blocks, block_size, row_size), (page_stride, row_size, 1) + ) + reference_backing = backing.clone() + reference = reference_backing.as_strided(cache.shape, cache.stride()) + kv = torch.randn(num_tokens + 3, 512, device="cuda", dtype=torch.bfloat16) + positions = torch.arange(num_tokens + 3, device="cuda") + 7 + angles = torch.randn(num_tokens + 16, 32, device="cuda") + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + slots = torch.randperm(num_blocks * block_size, device="cuda")[:num_tokens] + slots[::5] = -1 + scale = torch.tensor([0.7], device="cuda") + # Only what _insert_context_kv actually reads: the record width comes off + # the cache tensor, so no per-record width attribute is stubbed here. + attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(kv_cache=cache, block_size=block_size), + kv_mxfp8=kv_mxfp8, + rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin), + _flashinfer_fp8_kv_scale=scale, + ) + + def legacy_insert(): + q = torch.zeros(kv.shape[0], 8, 512, dtype=kv.dtype, device="cuda") + if cache_dtype == torch.uint8: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + reference.view(num_blocks, -1), + slots, + positions, + cos_sin, + 8, + 1e-20, + block_size, + True, + kv_mxfp8, + ) + elif cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + reference, + slots, + positions, + cos_sin, + 1e-20, + block_size, + ) + else: + q_fp8 = torch.empty_like(q, dtype=cache_dtype) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + reference, + slots, + positions, + cos_sin, + scale, + scale, + 1e-20, + block_size, + ) + + def insert(): + _insert_context_kv(attn, kv, positions, slots) + + insert() + graph = None + if use_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + insert() + for iteration in range(3): + kv.normal_() + positions.add_(1) + if iteration == 2: + slots.fill_(-1) + backing.view(torch.uint8).fill_(165) + reference_backing.copy_(backing) + before = kv.clone() + legacy_insert() + insert() if graph is None else graph.replay() + torch.testing.assert_close( + backing.view(torch.uint8), + reference_backing.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close(kv, before, rtol=0, atol=0) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="graph capture coverage" ) @@ -215,6 +331,136 @@ def run(): torch.testing.assert_close(cache_backing, expected_cache, rtol=0, atol=0) +def _mxfp8_record_reference( + latent_row: torch.Tensor, cos_sin_row: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference V4.1 record for one token: 512 fp8 bytes + 16 UE8M0 scales. + + Mirrors FlashMLA's ``KVCacheLayout.V41_FP8Sparse`` quantizer: rotate the + RoPE tail first, then scale every 32-dim tile, RoPE tiles included. + """ + row = latent_row.float() + c, s = cos_sin_row.chunk(2) + rotated = row.clone() + even, odd = row[448::2], row[449::2] + rotated[448::2] = even * c - odd * s + rotated[449::2] = odd * c + even * s + quantized, scales = _ue8m0_reference(rotated, 32, 448.0) + return quantized.view(torch.uint8), (scales.log2() + 127).to(torch.uint8) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("compress_ratio", [1, 2]) +def test_v41_rope_insert_mxfp8_record(compress_ratio: int): + """The 528-byte V4.1 record quantizes the RoPE dims instead of keeping them + in bf16, so a page is [block_size x 512 fp8][block_size x 16 UE8M0]. + + Byte-exact against the reference quantizer, including the rows and page + padding the kernel must leave alone. + """ + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + rope_quant_insert, + ) + + torch.manual_seed(7) + device = "cuda" + cache_block = 64 + cache_stride = math.ceil(cache_block * 528 / 512) * 512 + num_tokens = 12 + + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + latent = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + angles = torch.randn(64, 32, device=device) + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + + cache_backing = torch.full((2, cache_stride), 165, dtype=torch.uint8, device=device) + cache = cache_backing.as_strided((2, cache_block, 528), (cache_stride, 528, 1)) + # Leave one token unmapped so the kernel's negative-slot skip is covered. + slots = torch.arange(num_tokens, dtype=torch.int64, device=device) + slots[3] = -1 + + rope_quant_insert(latent, positions, cos_sin, cache, slots, compress_ratio) + + expected = torch.full_like(cache_backing, 165) + for t in range(num_tokens): + slot = slots[t].item() + pos = positions[t].item() + if slot < 0 or (pos + 1) % compress_ratio: + continue + page, row = divmod(slot, cache_block) + values, scales = _mxfp8_record_reference( + latent[t], cos_sin[pos // compress_ratio * compress_ratio] + ) + expected[page, row * 512 : (row + 1) * 512] = values + scale_offset = cache_block * 512 + row * 16 + expected[page, scale_offset : scale_offset + 16] = scales + + torch.testing.assert_close(cache_backing, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_v41_mxfp8_cache_round_trip(): + """quantize_and_insert -> dequantize_and_gather recovers the V4.1 record. + + The RoPE dims now go through fp8, so they carry quantization error too -- + the check is that every dim survives within one MXFP8 tile step, and that + the bytes on the way match the reference quantizer. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache, + quantize_and_insert_k_cache, + ) + + torch.manual_seed(11) + device = "cuda" + block_size = 64 + num_tokens = 70 + num_blocks = 4 + page_bytes = math.ceil(block_size * 528 / 512) * 512 + + k = torch.randn(num_tokens, 512, dtype=torch.bfloat16, device=device) + cache = torch.zeros(num_blocks, page_bytes, dtype=torch.uint8, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + quantize_and_insert_k_cache( + k, cache, slot_mapping, block_size=block_size, bytes_per_token=528 + ) + + for t in (0, 1, block_size, num_tokens - 1): + values, scales = _ue8m0_reference(k[t].float(), 32, 448.0) + page, row = divmod(t, block_size) + torch.testing.assert_close( + cache[page, row * 512 : (row + 1) * 512], + values.view(torch.uint8), + rtol=0, + atol=0, + ) + scale_offset = block_size * 512 + row * 16 + torch.testing.assert_close( + cache[page, scale_offset : scale_offset + 16], + (scales.log2() + 127).to(torch.uint8), + rtol=0, + atol=0, + ) + + out = torch.zeros(1, num_tokens, 512, dtype=torch.bfloat16, device=device) + dequantize_and_gather_k_cache( + out, + cache.view(num_blocks, block_size, 528), + seq_lens=torch.tensor([num_tokens], dtype=torch.int32, device=device), + gather_lens=None, + block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view( + 1, -1 + ), + block_size=block_size, + offset=0, + ) + # Half an e4m3 ULP at the top of the range (16 units) bounds the error. + tile_amax = k.float().abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) + tolerance = 16.0 * (tile_amax / 448.0).log2().ceil().exp2() + error = (out[0].float() - k.float()).abs() + assert (error <= tolerance.repeat_interleave(32, dim=-1)).all() + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") @pytest.mark.parametrize("compress_ratio", [1, 2]) @pytest.mark.parametrize("store_fp8", [False, True]) diff --git a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py index 0e25884368c1..68eabfb94e84 100644 --- a/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py +++ b/tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py @@ -424,6 +424,104 @@ def _dequant(k_cache_2d): ) +# ── Test 2c: DeepSeek V4.1 MXFP8 record (kv_mxfp8=True) ────────────────────── + +V41_HEAD_BYTES = HEAD_DIM + HEAD_DIM // 32 # 512 fp8 + 16 UE8M0 scales = 528 + + +@pytest.mark.parametrize("num_tokens", [1, 17, 64]) +@pytest.mark.parametrize("block_size", [32, 64]) +def test_kv_path_mxfp8_record_matches_reference(num_tokens: int, block_size: int): + """``kv_mxfp8=True`` writes DeepSeek's V4.1 record: all 512 dims as fp8 e4m3 + with one UE8M0 scale per 32 dims, no bf16 RoPE tail. + + Both sides quantize the same bf16-rounded rotated row, so the bytes must + match the Triton encoder exactly apart from the documented 1-ULP bf16 + rounding split on the rotation itself. + """ + from vllm.models.deepseek_v41.common.ops import ( + dequantize_and_gather_k_cache as v41_dequantize_and_gather_k_cache, + ) + from vllm.models.deepseek_v41.common.ops import ( + quantize_and_insert_k_cache as v41_quantize_and_insert_k_cache, + ) + + torch.manual_seed(3) + device = "cuda" + dtype = torch.bfloat16 + max_pos = 4096 + + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype, device=device) + positions = torch.arange(num_tokens, dtype=torch.int64, device=device) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + + k_cache_ref = torch.zeros( + num_blocks, block_size * V41_HEAD_BYTES, dtype=torch.uint8, device=device + ) + v41_quantize_and_insert_k_cache( + apply_rope_gptj_last_k(kv, positions, cos_sin_cache), + k_cache_ref, + slot_mapping, + block_size=block_size, + use_fnuz=USE_FNUZ, + bytes_per_token=V41_HEAD_BYTES, + ) + + k_cache_fused = torch.zeros_like(k_cache_ref) + q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q_dummy, + kv, + k_cache_fused, + slot_mapping, + positions, + cos_sin_cache, + 64, + 1e-6, + block_size, + True, + True, # kv_mxfp8 + ) + + # A pre-quantization bf16 tie can land the two encoders on adjacent e4m3 + # codes, so compare after dequantizing and allow one step at the tile scale. + def _dequant(cache_2d): + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=dtype, device=device) + v41_dequantize_and_gather_k_cache( + out, + cache_2d.view(num_blocks, block_size, V41_HEAD_BYTES), + torch.tensor([num_tokens], dtype=torch.int32, device=device), + None, + torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(0), + block_size, + offset=0, + use_fnuz=USE_FNUZ, + ) + return out[0] + + rec_ref = _dequant(k_cache_ref).float() + rec_fused = _dequant(k_cache_fused).float() + tile_scale = ( + (rec_ref.abs().view(num_tokens, 16, 32).amax(-1).clamp(min=1e-4) / FP8_MAX) + .log2() + .ceil() + .exp2() + ) + # One e4m3 ULP at the top of the range is 32 quantized units. + step = 32.0 * tile_scale.repeat_interleave(32, dim=-1) + assert ((rec_fused - rec_ref).abs() <= step).all() + # The NoPE dims see no rotation at all, so they must be byte-identical. + nope_bytes = slice(0, block_size * HEAD_DIM) + ref_rows = k_cache_ref[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + fused_rows = k_cache_fused[:, nope_bytes].view(num_blocks, block_size, HEAD_DIM) + torch.testing.assert_close( + fused_rows[..., :NOPE_DIM], ref_rows[..., :NOPE_DIM], rtol=0, atol=0 + ) + + # ── Test 2b: DP padding (slot_mapping shorter than q/kv) ───────────────────── diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 8770882ddf7c..7464c9076973 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -560,3 +560,55 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat torch.testing.assert_close( k_scale[emitted], expected_scale[emitted], rtol=0, atol=0 ) + + +@pytest.mark.parametrize("use_cutedsl", [False, True]) +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", +) +@torch.inference_mode() +def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): + """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's + sparse MQA-logits kernels take) instead of fp32; the values are the fp32 + result rounded once, so the scoring kernel needs no cast.""" + if use_cutedsl and not has_cutedsl(): + pytest.skip("cutedsl (cutlass) not installed") + + device = "cuda" + torch.manual_seed(0) + num_tokens, n_head = 257, 32 + q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) + positions = torch.randint( + 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device + ) + cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) + weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) + softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125 # pow2 head_scale, see above + + _, weights_ref = _reference( + positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True + ) + cutedsl_patch = ( + mock.patch( + "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl", + return_value=False, + ) + if not use_cutedsl + else contextlib.nullcontext() + ) + with cutedsl_patch: + _, weights_fused = fused_indexer_q_rope_quant( + positions, + q.clone(), + cos_sin_cache, + weights, + softmax_scale, + head_scale, + use_fp4=True, + weights_out_dtype=torch.bfloat16, + ) + assert weights_fused.dtype == torch.bfloat16 + assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16)) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index f805fe39a57b..5fab150d9fce 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -11,6 +11,7 @@ import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( _torch_hc_prenorm_gemm, + mhc_post_tilelang, mhc_pre_delayed_tilelang, ) from vllm.model_executor.kernels.mhc.tilelang_kernels import ( @@ -39,6 +40,10 @@ from vllm.models.deepseek_v41.nvidia.model import ( DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, ) +from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import ( + is_mega_mhc_supported, + mhc_shifted_post_pre_deep_gemm, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -147,6 +152,27 @@ def forward(hidden, positions, ids, pre, post, res, residual, **kwargs): ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("num_tokens", [1, 7, 32]) +@pytest.mark.parametrize("tp_size", [4, 8]) +def test_hc_collapse_before_gather_is_bitwise_equal(num_tokens, tp_size): + """Simulate SP on one GPU, including padding and ranks without valid tokens.""" + set_random_seed(0) + padded_tokens = num_tokens + (-num_tokens) % tp_size + streams = torch.randn(padded_tokens, 4, 5120, dtype=torch.bfloat16, device=DEVICE) + pre_mix = torch.rand(padded_tokens, 4, device=DEVICE) + stream_shards = streams.chunk(tp_size) + mix_shards = pre_mix.chunk(tp_size) + before = hc_collapse_triton( + torch.cat(stream_shards)[:num_tokens], torch.cat(mix_shards)[:num_tokens] + ) + after = torch.cat( + [hc_collapse_triton(x, pre) for x, pre in zip(stream_shards, mix_shards)] + )[:num_tokens] + assert after.shape == (num_tokens, 5120) + torch.testing.assert_close(after, before, atol=0, rtol=0) + + @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.parametrize( "num_tokens,sinkhorn_repeat,norm_eps", @@ -474,7 +500,7 @@ def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, ** reference, ) monkeypatch.setattr( - "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang", + "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre", fused_reference, ) expected = decoder(x, positions, None, **kwargs) @@ -559,6 +585,104 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried): ) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="DeepGEMM Mega mHC requires SM100-family CUDA", +) +@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024]) +@pytest.mark.parametrize("hidden_size", [5120, 7168]) +def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size): + set_random_seed(0) + hc_mult = 4 + if not is_mega_mhc_supported(hidden_size, 4): + pytest.skip("DeepGEMM Mega mHC is not available") + + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) + post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) + res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) + fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01 + scale = torch.randn(3, device=DEVICE) * 0.1 + base = torch.randn(24, device=DEVICE) * 0.1 + norm_weight = torch.empty( + hidden_size, dtype=torch.bfloat16, device=DEVICE + ).uniform_(0.9, 1.1) + + expected_residual = ( + mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual + ) + ( + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) = mhc_pre_delayed_tilelang( + expected_residual, + fn, + scale, + base, + 2e-5, + 3e-4, + 2e-6, + 1.25, + 10, + pre_mix=previous_mix, + norm_weight=norm_weight, + norm_eps=7e-6, + ) + args = ( + x, + residual, + previous_mix, + post_mix, + res_mix, + fn, + scale, + base, + 2e-5, + 3e-4, + 1.25, + 2e-6, + 10, + norm_weight, + 7e-6, + ) + actual = mhc_shifted_post_pre_deep_gemm(*args) + expected = ( + expected_residual, + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) + for i, (result, ref) in enumerate(zip(actual, expected, strict=True)): + # Match the existing mHC BF16 tolerance: post rounding differences + # can be amplified by cancellation in the carried collapse. + atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3) + torch.testing.assert_close(result, ref, atol=atol, rtol=rtol) + + if num_tokens: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # DeepGEMM initializes barriers per stream before capture. + mhc_shifted_post_pre_deep_gemm(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = mhc_shifted_post_pre_deep_gemm(*args) + for multiplier in (0.5, -1.0): + x.mul_(multiplier) + graph.replay() + eager = mhc_shifted_post_pre_deep_gemm(*args) + for result, ref in zip(captured, eager, strict=True): + torch.testing.assert_close(result, ref, atol=0, rtol=0) + torch.cuda.current_stream().wait_stream(stream) + + def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: x = x.softmax(-1) + eps x = x / (x.sum(-2, keepdim=True) + eps) diff --git a/tests/kernels/test_minimax_m3_amd_ops.py b/tests/kernels/test_minimax_m3_amd_ops.py index 182672597a9b..997c61980913 100644 --- a/tests/kernels/test_minimax_m3_amd_ops.py +++ b/tests/kernels/test_minimax_m3_amd_ops.py @@ -42,6 +42,7 @@ from vllm.models.minimax_m3.amd.ops import ( # noqa: E402 gemma_fused_add_rmsnorm, gemma_rmsnorm, + swiglu_oai_quantize_mxfp8, swiglu_oai_split, ) from vllm.models.minimax_m3.amd.ops.gemma_rmsnorm import _num_warps # noqa: E402 @@ -167,6 +168,25 @@ def test_swiglu_oai_split(m, inter, limit, dtype): assert _relerr(got, ref) < 5e-3 +@torch.inference_mode() +def test_swiglu_oai_quantize_mxfp8_uses_e4m3_range_for_scale(): + # Keep a small value in two MX blocks next to their maxima. The scale must + # use E4M3's finite range (448), otherwise that value underflows to zero. + # Keep the third block empty to cover the reference's finite tiny clamp. + m, inter = 3, 96 + gate = torch.full((m, inter), 2.0, device=DEVICE, dtype=torch.float16) + up = torch.full((m, inter), 2**-10, device=DEVICE, dtype=torch.float16) + up[:, :64:32] = 1.0 + up[:, 64:] = 0.0 + gate_up = torch.cat((gate, up), dim=-1) + + got_q, got_s = swiglu_oai_quantize_mxfp8(gate_up, alpha=0.0, beta=0.0, limit=None) + ref_q, ref_s = _mxfp8_e4m3_quantize_torch(up, is_sf_swizzled_layout=False) + + assert torch.equal(got_s, ref_s) + assert torch.equal(got_q, ref_q) + + # --------------------------------------------------------------------------- # # Fused MXFP8 activation quant (Triton vs torch reference) # --------------------------------------------------------------------------- # @@ -179,8 +199,8 @@ def test_mxfp8_quant_triton_matches_torch(shape, dtype): xq_t, s_t = _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout=False) xq_k, s_k = _mxfp8_e4m3_quantize_triton(x) assert s_k.shape == s_t.shape == (shape[0], shape[1] // 32) - # E8M0 block exponents share the floor(log2(amax))+127 algorithm; allow at - # most a 1-step difference at exact powers of two. + # Both paths use the E4M3-aware scale calculation; allow a 1-step + # difference at exact powers of two due to floating-point rounding. assert (s_k.int() - s_t.int()).abs().max().item() <= 1 # Dequantized values agree to fp8 granularity. deq_t = dequant_mxfp8_to_bf16(xq_t, s_t) @@ -424,6 +444,69 @@ def test_mxfp8_linear_emulation_bf16_at_load( assert _relerr(out.float(), out_ref.float()) < 2e-2 +# --------------------------------------------------------------------------- # +# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load +# --------------------------------------------------------------------------- # +@requires_gfx950 +# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8. +@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)]) +@torch.inference_mode() +def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape): + """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it. + + ``process_weights_after_loading`` converts those to BF16 once; a weight that + never went through it has to fall back per step rather than take the + invalid path, which is what keying the dispatch on the weight dtype alone + would do. + """ + from vllm.model_executor.kernels.linear.mxfp8 import rocm_native + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, + ) + + N, K = shape + aligned = K % 128 == 0 + torch.manual_seed(0) + w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1 + w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False) + x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5 + out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale)) + # dot_scaled re-quantizes the activation, the dequant paths do not. + tol = 5e-2 if aligned else 2e-3 + + def make_layer(): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False) + layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False) + return layer + + kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig()) + layer = make_layer() + kernel.process_weights_after_loading(layer) + assert (layer.weight.element_size() >= 2) is not aligned + assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol + + # An unaligned weight that is still FP8 must not reach dot_scaled. + unprocessed = make_layer() + unprocessed.weight_scale = torch.nn.Parameter( + w_scale[:N, : K // 32].contiguous(), requires_grad=False + ) + calls = [] + real = rocm_native._mxfp8_dot_scaled_linear + + def _spy(*a, **k): + calls.append(1) + return real(*a, **k) + + rocm_native._mxfp8_dot_scaled_linear = _spy + try: + out = kernel.apply_weights(unprocessed, x) + finally: + rocm_native._mxfp8_dot_scaled_linear = real + assert bool(calls) is aligned + assert _relerr(out.float(), out_ref.float()) < tol + + # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ──────────────── # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index c629ae15dea9..289539470631 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -22,11 +22,22 @@ def _has_device_capability(major: int) -> bool: return current_platform.is_cuda() and current_platform.has_device_capability(major) +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + # DeepSelect is compiled for sm_100a/sm_103a only. requires_sm100 = pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="DeepSelect requires SM100a/SM103a", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), reason="This test exercises the gfx950 launch configuration" +) COOPERATIVE_TOPK_BACKEND = pytest.param( @@ -487,6 +498,150 @@ def test_top_k_per_row_decode_gfx950_long_c4a_1d_seq_lens() -> None: ) +def _assert_exact_topk( + logits: torch.Tensor, indices: torch.Tensor, row_ends: torch.Tensor +) -> None: + ends = row_ends.reshape(-1, 1).clamp_min(0) + valid = (indices >= 0) & (indices < ends) + assert torch.equal(valid.sum(1), ends.clamp_max(indices.shape[1]).flatten()) + assert torch.all(valid | (indices == -1)) + ordered = indices.sort(1).values + assert not ((ordered[:, 1:] == ordered[:, :-1]) & (ordered[:, 1:] >= 0)).any() + + width = int(ends.max()) + visible = logits[:, :width].clone() + columns = torch.arange(width, device=logits.device) + visible.masked_fill_(columns[None, :] >= ends, float("-inf")) + expected = visible.topk(min(indices.shape[1], width), dim=1).values + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, float("-inf")) + selected = selected.sort(1, descending=True).values + assert torch.equal(selected[:, : expected.shape[1]], expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_1d_seq_lens() -> None: + """Honor causal offsets and exact compressed lengths, including short rows.""" + next_n, width, top_k = 6, 16_384, 512 + lengths = torch.tensor( + [0, 1, 511, 512, 513, 1023, 1024, 16_000], + dtype=torch.int32, + device="cuda", + ) + offsets = torch.arange(next_n, dtype=torch.int32, device="cuda") + row_ends = (lengths[:, None] - next_n + offsets + 1).clamp_min(0).flatten() + rows = row_ends.numel() + storage = torch.full((rows, width + 8), 1e20, dtype=torch.float32, device="cuda") + logits = storage[:, 1 : width + 1] + logits[:, :16_000] = torch.arange(16_000, dtype=torch.float32, device="cuda") + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + _assert_exact_topk(logits, indices, row_ends) + + +@pytest.mark.parametrize( + ("rows", "row_length"), + [ + pytest.param(128, 65_536, id="single-before-64k"), + pytest.param(128, 65_537, id="four-splits-after-64k"), + pytest.param(129, 262_144, id="single-before-three-split-length"), + pytest.param(129, 262_145, id="three-splits-after-256k"), + pytest.param(160, 262_145, id="three-splits-upper-row-bound"), + pytest.param(161, 262_144, id="single-before-two-split-length"), + pytest.param(161, 262_145, id="two-splits-after-256k"), + pytest.param(256, 262_145, id="two-splits-upper-row-bound"), + pytest.param(257, 262_145, id="large-row-single-block"), + pytest.param(384, 262_145, id="optimized-row-limit"), + pytest.param(385, 262_145, id="native-fallback"), + ], +) +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_split_policy_boundaries( + rows: int, row_length: int +) -> None: + width, top_k = 262_145, 512 + row = torch.arange(width, dtype=torch.float32, device="cuda") + logits = row.expand(rows, -1) + lengths = torch.full((rows,), row_length, dtype=torch.int32, device="cuda") + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + + _run_topk_backend("top_k_per_row_decode", logits, lengths, indices, top_k, width) + + expected = torch.arange( + row_length - top_k, row_length, dtype=torch.int32, device="cuda" + ).expand_as(indices) + assert torch.equal(indices.sort(1).values, expected) + + +@requires_gfx950 +@torch.inference_mode() +def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: + """Preserve visible top-k scores when masks and lengths change on replay.""" + rows, next_n, width, top_k = 24, 6, 262_145, 512 + logits = torch.full((rows, width), 1e20, dtype=torch.float32, device="cuda") + lengths = torch.full( + (rows // next_n, next_n), 70_000, dtype=torch.int32, device="cuda" + ) + indices = torch.full((rows, top_k), -777, dtype=torch.int32, device="cuda") + logits[:, :70_000] = float("-inf") + bin_sizes = [128, 129, 1024, 1025] + for row in range(rows): + pattern = row % 9 + if pattern < 4: + count = [0, 1, 511, 513][pattern] + logits[row, :count] = torch.arange( + count, dtype=torch.float32, device="cuda" + ) + elif pattern == 4: + logits[row, :511] = 0 + logits[row, 511] = torch.finfo(torch.float32).min + else: + count = bin_sizes[pattern - 5] + higher = max(0, top_k - count // 2) + logits[row, :higher] = 2 + logits[row, higher : higher + count] = 1 + logits[row, higher] = 1 + 2**-23 + + def run() -> None: + _run_topk_backend( + "top_k_per_row_decode", logits, lengths, indices, top_k, width, next_n + ) + + run() + _assert_exact_topk(logits, indices, lengths) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + bounds = [ + 0, + 511, + 512, + 513, + 16_384, + 16_385, + 131_072, + 131_073, + 262_144, + 262_145, + ] + row_bounds = torch.tensor(bounds, dtype=torch.int32, device="cuda") + repeats = (rows + len(bounds) - 1) // len(bounds) + lengths.flatten().copy_(row_bounds.repeat(repeats)[:rows]) + logits[:] = torch.arange(width, dtype=torch.float32, device="cuda") + indices.fill_(-777) + graph.replay() + _assert_exact_topk(logits, indices, lengths) + + @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @torch.inference_mode() def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: @@ -620,6 +775,31 @@ def decode_kernel(*args, **kwargs) -> None: compress_ratio=4, num_rows=257, max_valid_seq_len=250_000, + num_columns=524_288, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=385, + max_valid_seq_len=250_000, + num_columns=524_288, + topk_tokens=512, + on_gfx950=True, + ), + decode_kernel, + ), + ( + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=50_000, + num_columns=1_048_577, + topk_tokens=512, on_gfx950=True, ), decode_kernel, @@ -638,6 +818,15 @@ def fail_if_imported() -> None: ) ineligible_cases = [ dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True), + dict( + is_prefill=False, + compress_ratio=2, + num_rows=384, + max_valid_seq_len=1_048_576, + num_columns=1_048_576, + topk_tokens=512, + on_gfx950=True, + ), dict( is_prefill=False, compress_ratio=4, @@ -743,6 +932,106 @@ def test_deepseek_workspace_topk( ) +@pytest.mark.skipif(not _has_device_capability(80), reason="Requires SM80+") +@pytest.mark.parametrize("rows", [65, 128]) +@pytest.mark.parametrize("top_k", [512, 1024, 2048]) +@pytest.mark.parametrize( + "distribution", ["random", "10LSBits", "ascending", "constant", "sampled_peaks"] +) +@torch.inference_mode() +def test_persistent_topk_sampled_graph( + rows: int, top_k: int, distribution: str +) -> None: + """Sampling and both fallbacks preserve exact values across graph replays.""" + if torch.cuda.get_device_properties(0).shared_memory_per_block_optin < 144 * 1024: + pytest.skip("Sampled top-k requires at least 144 KiB of shared memory") + set_random_seed(42) + min_sampled_length = 98304 if top_k == 512 else 65536 + width = min_sampled_length + 3 + lengths = torch.full((rows,), width, dtype=torch.int32, device="cuda") + logits = create_random_logits( + torch.zeros_like(lengths), + lengths, + torch.float32, + 42, + False, + "10LSBits" if distribution == "10LSBits" else "random", + ) + if distribution == "ascending": + logits.copy_(torch.arange(width, device="cuda", dtype=torch.float32)) + elif distribution == "constant": + logits.fill_(1.0) + elif distribution == "sampled_peaks": + # A biased sample leaves fewer than k survivors and must fall back. + logits.zero_() + sample = torch.arange(4096, device="cuda") + positions = (sample // 32) * (width // 128) + sample % 32 + logits[:, positions] = 1 + sample.float() / 4096 + original = logits.clone() + indices = torch.empty((rows, top_k), dtype=torch.int32, device="cuda") + workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda") + + def run() -> None: + torch.ops._C.persistent_topk( + logits, + lengths.view(-1, 4 if rows == 128 else 1), + indices, + workspace, + top_k, + width, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + positions = torch.arange(width, device="cuda") + slots = torch.arange(top_k, device="cuda") + for step in range(3): + bounds = torch.tensor( + [ + -1, + 0, + 1, + top_k - 1, + top_k, + 32769, + min_sampled_length - 1, + min_sampled_length, + width, + width + 17, + ], + dtype=torch.int32, + device="cuda", + ) + lengths.copy_( + bounds[(torch.arange(rows, device="cuda") + step) % bounds.numel()] + ) + logits.copy_(original if step != 1 else original.flip(1)) + logits.masked_fill_(positions[None] >= lengths[:, None], float("nan")) + indices.fill_(-2) + graph.replay() + valid = slots[None] < lengths.clamp(max=top_k)[:, None] + assert torch.all(indices[~valid] == -1) + assert torch.all(((indices >= 0) & (indices < lengths[:, None])) == valid) + ordered_indices = indices.sort(dim=1).values + assert torch.all( + (ordered_indices[:, 1:] != ordered_indices[:, :-1]) + | (ordered_indices[:, 1:] == -1) + ) + selected = logits.gather(1, indices.clamp_min(0).long()) + selected.masked_fill_(~valid, -float("inf")) + expected = logits.masked_fill( + positions[None] >= lengths[:, None], -float("inf") + ) + torch.testing.assert_close( + selected.sort(dim=1, descending=True).values, + expected.topk(top_k, dim=1).values, + atol=0, + rtol=0, + ) + + def run_large_context_topk_test( batch_size: int, seq_lens: list[int], diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index f8e7a89ccad2..41a06d0546f0 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -450,3 +450,547 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block) reference = torch.where(keep, 3.0, -torch.inf) torch.testing.assert_close(logits, reference, rtol=0, atol=0) + + +# --- DeepGEMM sparse-MQA indexer path (V4.1 two-level candidate filtering) --- + + +def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): + """Production-faithful candidates: top min(avail, K) blocks per row — + full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" + avail = ((lens.cpu() + cbk - 1) // cbk).tolist() + rows = [] + for a in avail: + n = min(a, num_candidates) + picks = torch.randperm(a)[:n].sort().values[:num_candidates].tolist() + rows.append(picks + [-1] * (num_candidates - len(picks))) + return torch.tensor(rows, dtype=torch.int32) + + +def _candidate_token_mask( + candidates: torch.Tensor, cbk: int, total: int, ks: torch.Tensor | None = None +): + """[rows, total] bool mask of the candidate-covered token positions. + + ``ks`` converts request-local candidate blocks to absolute positions + (prefill's packed workspace); pass None when already request-local. + """ + rows = candidates.shape[0] + device = candidates.device + # Sentinel block keeps -1 padding from scattering onto real columns. + width = max(total, (candidates.shape[1] + 1) * cbk) + if ks is not None: + width += int(ks.max()) + cbk + in_cand = torch.zeros(rows, width, dtype=torch.bool, device=device) + block_id = torch.where(candidates >= 0, candidates.long(), candidates.shape[1]) + tok = block_id.unsqueeze(2) * cbk + torch.arange(cbk, device=device) + if ks is not None: + tok = tok + ks.unsqueeze(1).unsqueeze(2) + keep = (candidates >= 0).unsqueeze(2).expand(-1, -1, cbk) + in_cand.scatter_(1, tok.flatten(1), keep.flatten(1)) + return in_cand[:, :total] + + +def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: + """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid + count are -1 (the indexer's "no token" sentinel).""" + values, indices = masked.topk(topk, dim=1) + return torch.where(values > float("-inf"), indices.int(), -1) + + +def _assert_same_selection( + logits: torch.Tensor, + got: torch.Tensor, + expected: torch.Tensor, + row_starts: torch.Tensor | None = None, +) -> None: + """Both top-k index sets must pick the same multiset of logit values per + row (-1 padding must match too). Indices themselves may differ: the + top-k kernels break ties among equal values in nondeterministic order.""" + for r in range(logits.shape[0]): + base = int(row_starts[r]) if row_starts is not None else 0 + g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] + assert g.numel() == e.numel(), f"row {r}: {g.numel()} vs {e.numel()} selected" + vg = logits[r][(base + g).long()].sort().values + ve = logits[r][(base + e).long()].sort().values + assert torch.equal(vg, ve), f"row {r}: selected values differ" + + +def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: + """Caller-owned buffers the sparse path writes into (the metadata builder + owns these in production).""" + i32 = dict(dtype=torch.int32, device="cuda") + return dict( + sparse_indices=torch.zeros(rows, num_sparse, **i32), + end=torch.zeros(rows, **i32), + col_indices=torch.zeros(rows, topk, **i32), + ) + + +def _expand_candidates_reference( + candidates: torch.Tensor, ks: torch.Tensor, ke: torch.Tensor, cbk: int, sbk: int +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand + each candidate block into ``cbk // sbk`` sparse blocks, keep those inside + ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by + repeating the last valid block, and count the valid sparse columns.""" + ratio = cbk // sbk + rows, num_candidates = candidates.shape + num_sparse = num_candidates * ratio + big = torch.iinfo(torch.int32).max + offsets = torch.arange(ratio, device=candidates.device, dtype=torch.int32) + blocks = (candidates.unsqueeze(-1) * ratio + offsets).flatten(1) + max_block = ((ke - ks).clamp(min=0) + sbk - 1) // sbk + valid = (blocks >= 0) & (blocks < max_block.unsqueeze(1)) + sorted_blocks, _ = torch.where(valid, blocks, big).sort(dim=1) + n_valid = (sorted_blocks != big).sum(dim=1, dtype=torch.int32) + base = (ks // sbk).unsqueeze(1) + kernel_blocks = torch.where(sorted_blocks != big, sorted_blocks, 0) + base + last_slot = (n_valid - 1).clamp(min=0).unsqueeze(1).long() + last_block = kernel_blocks.gather(1, last_slot) + pad = torch.arange(num_sparse, device=candidates.device) >= n_valid.unsqueeze(1) + sparse_indices = torch.where(pad, last_block, kernel_blocks) + last_start = last_block.squeeze(1) * sbk + ks % sbk + last_fill = (ke - last_start).clamp(min=1, max=sbk) + end = torch.where(n_valid > 0, (n_valid - 1) * sbk + last_fill, 0) + return sparse_indices.int(), end.int() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +def test_candidate_blocks_to_sparse_indices_math(): + """Candidate->sparse-block expansion: ks anchoring, range filtering, + repeat-last padding, and the per-row valid column count (``end``).""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + # Rows cover: exactly one block, an unaligned ks with a partial tail + # block, an empty row, duplicates / -1 padding / out-of-range garbage + # (warmup runs touch uninitialized buffers), and int32-overflow junk. + ks = torch.tensor([64, 13, 100, 8, 0], dtype=torch.int32, device="cuda") + ke = torch.tensor([72, 30, 100, 40, 1000], dtype=torch.int32, device="cuda") + candidates = torch.tensor( + [ + [0, -1, -1, -1], + [2, 0, -1, -1], + [3, -1, -1, -1], + [1, 1, 7, 0], + [2**30, -(2**30), 0, -1], + ], + dtype=torch.int32, + device="cuda", + ) + indices, end = candidate_blocks_to_sparse_indices( + candidates, ks, ke, candidate_block_size=8, sparse_block_kv=8 + ) + assert indices.tolist() == [ + [8, 8, 8, 8], + [1, 3, 3, 3], + [12, 12, 12, 12], + [1, 2, 2, 2], + [0, 0, 0, 0], + ] + # end: row 1's last block covers [29, 37) but ke = 30 -> one valid token; + # row 3 keeps the duplicated block (production candidates are unique). + assert end.tolist() == [8, 9, 0, 24, 8] + + # Wider candidate blocks expand into multiple sparse blocks. + idx2, end2 = candidate_blocks_to_sparse_indices( + torch.tensor([[1, 0]], dtype=torch.int32, device="cuda"), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([48], dtype=torch.int32, device="cuda"), + candidate_block_size=16, + sparse_block_kv=8, + ) + assert idx2.tolist() == [[0, 1, 2, 3]] + assert end2.tolist() == [32] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)]) +def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): + """The Triton expansion equals the PyTorch reference on production-shaped + random inputs, including strided (column-sliced) candidate rows and + strided output buffers.""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + torch.manual_seed(0) + rows, num_candidates = 41, 2048 + ratio = cbk // sbk + lens = torch.randint(1, 40000, (rows,), dtype=torch.int32) + lens[0], lens[1] = 1, 3 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke = ks.cuda(), (ks + lens).cuda() + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + # Row 2 gets out-of-range garbage past the padding, as a CUDA-graph + # warm-up on uninitialized buffers would. + candidates[2, -3:] = torch.tensor([2**30, -(2**29), 10**6], dtype=torch.int32) + + ref_indices, ref_end = _expand_candidates_reference(candidates, ks, ke, cbk, sbk) + + cand_buf = torch.full((rows, num_candidates + 64), -7, dtype=torch.int32).cuda() + cand_buf[:, :num_candidates] = candidates + si_buf = torch.zeros(rows, num_candidates * ratio + 32, dtype=torch.int32).cuda() + end_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf[::2] = ks + got_indices, got_end = candidate_blocks_to_sparse_indices( + cand_buf[:, :num_candidates], + ks_buf[::2], + ke, + cbk, + sbk, + out=(si_buf[:, : num_candidates * ratio], end_buf[::2]), + ) + assert torch.equal(got_indices, ref_indices) + assert torch.equal(got_end, ref_end) + assert torch.equal( + cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7) + ) + assert not si_buf[:, num_candidates * ratio :].any() + assert not end_buf[1::2].any() + + +def _skip_unless_sm100_sparse_kernels(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 required") + deep_gemm = pytest.importorskip("deep_gemm") + if not hasattr(deep_gemm, "fp8_fp4_sparse_mqa_logits"): + pytest.skip("DeepGEMM >= 2.8 required for sparse MQA logits") + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + ) + + if not has_deep_select(): + pytest.skip("DeepSelect extension (vllm._deepselect_C) required") + return deep_gemm + + +def _quant_fp4(x: torch.Tensor): + from deep_gemm.utils import per_token_cast_to_fp4 + + return per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): + """The sparse-MQA candidate path selects exactly the same logit values as + ``torch.topk`` over dense bf16 logits masked to the candidates (the two + DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA + graph capture, and a later layer reusing the step's expansion/schedule + selects the same values.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_prefill_chunk, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 37, 32, 128 + cbk, topk, num_candidates = 8, 512, 2048 # V4.1-Flash production values + lens = torch.randint(1, 30000, (rows,), dtype=torch.int32) + lens[0], lens[1], lens[2] = 1, 5, 2048 # single token, sub-block, long + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + + # Reference: dense bf16 logits masked to the candidates, then torch.topk. + dense = deep_gemm.fp8_fp4_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + ks, + ke, + clean_logits=False, + logits_dtype=torch.bfloat16, + ) + cols = torch.arange(total, device="cuda") + in_range = (cols - ks.unsqueeze(1) >= 0) & ( + cols - ks.unsqueeze(1) < lens.unsqueeze(1) + ) + masked = dense.masked_fill( + ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), + float("-inf"), + ).float() + ref_abs = _reference_topk(masked, topk) # packed-workspace positions + ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(kernel_metadata=None): + return sparse_mqa_logits_prefill_chunk( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + k_fp.view(torch.int8), + k_sf.view(total), + weights, + ks, + ke, + candidates, + cbk, + cbk, + topk, + out, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse() # warm up JIT + the DeepGEMM workspace pre-capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse() + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse() + _assert_same_selection(dense, out, ref, row_starts=ks) + out.fill_(-1) + run_sparse(kernel_metadata) + + _assert_same_selection(dense, out, ref, row_starts=ks) + + +def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): + """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' + fp4 data first, then their packed UE8M0 scales.""" + device = req_lens.device + pages_per = (req_lens + page_kv - 1) // page_kv + num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) + kv_cache = torch.zeros( + num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device + ) + block_table = torch.zeros( + len(req_lens), max_pages, dtype=torch.int32, device=device + ) + page_pool = torch.randperm(num_pages, device=device) + offset = 0 + for b in range(len(req_lens)): + n = int(pages_per[b]) + block_table[b, :n] = page_pool[offset : offset + n].to(torch.int32) + offset += n + k_fp, k_sf = _quant_fp4( + torch.randn(int(req_lens[b]), head_dim, device=device, dtype=torch.bfloat16) + ) + rows68 = torch.zeros(n * page_kv, 68, dtype=torch.uint8, device=device) + rows68[: int(req_lens[b]), : head_dim // 2] = k_fp.view(torch.uint8) + rows68[: int(req_lens[b]), head_dim // 2 :] = k_sf.view(torch.uint8).view(-1, 4) + by_page = rows68.view(n, page_kv, 68) + split = torch.cat( + [by_page[..., :64].reshape(n, -1), by_page[..., 64:].reshape(n, -1)], dim=1 + ).view(n, page_kv, 1, head_dim // 2 + 4) + kv_cache[block_table[b, :n].long()] = split + return kv_cache, block_table + + +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool): + """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks + from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to + one query row each), CUDA graph capture/replay and schedule reuse by a + later layer.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_paged_decode, + ) + + torch.manual_seed(0) + batch, num_heads, head_dim, page_kv = 9, 32, 128, 128 + cbk, topk, num_candidates = 8, 512, 2048 + rows = batch * next_n + req_lens = torch.randint(1, 30000, (batch,), dtype=torch.int32, device="cuda") + req_lens[0], req_lens[1] = 1, 30000 # tiny row + a filtered (subset) row + # Row (b, j) sees L_b - next_n + j + 1 tokens (native MTP layout). + row_lens = ( + (req_lens.unsqueeze(1) - next_n + torch.arange(1, next_n + 1, device="cuda")) + .clamp(min=0) + .to(torch.int32) + ) + kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim) + assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it + + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + q = q_fp.view(torch.int8).view(batch, next_n, num_heads, head_dim // 2) + q_scale = q_sf.view(batch, next_n, num_heads) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() + + def reference_topk(): + context_lens_2d = row_lens.view(batch, next_n) + sched = deep_gemm.get_paged_mqa_logits_metadata( + context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + (q, q_scale), + kv_cache, + weights, + context_lens_2d, + block_table, + sched, + int(req_lens.max()), + False, + torch.bfloat16, + None, + ) + lens_flat = row_lens.flatten() + cols = torch.arange(logits.shape[1], device="cuda") + masked = logits.masked_fill( + ~( + (cols.unsqueeze(0) < lens_flat.unsqueeze(1)) + & _candidate_token_mask(candidates, cbk, logits.shape[1]) + ), + float("-inf"), + ).float() + return _reference_topk(masked, topk), logits + + expected, dense_logits = reference_topk() + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + + # The metadata builder flattens every query to one row: per-row context + # lengths, an expanded block table and a row -> request map. + context_lens = row_lens.flatten().contiguous() + row_block_table = block_table.repeat_interleave(next_n, dim=0) + row_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + row_indices = row_indices.repeat_interleave(next_n) + row_ks = torch.zeros(rows, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(o: torch.Tensor, kernel_metadata=None): + return sparse_mqa_logits_paged_decode( + q.view(rows, 1, num_heads, head_dim // 2), + q_scale.view(rows, 1, num_heads), + kv_cache, + weights, + context_lens, + row_block_table, + row_indices, + candidates, + cbk, + cbk, + topk, + o, + row_ks=row_ks, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse(out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse(out) + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse(out) + _assert_same_selection(dense_logits, out, expected) + out.fill_(-1) + run_sparse(out, kernel_metadata) + + _assert_same_selection(dense_logits, out, expected) + + +def test_sparse_mqa_logits_accuracy_vs_torch(): + """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized + MXFP4 inputs.""" + _skip_unless_sm100_sparse_kernels() + from deep_gemm.utils import cast_back_from_fp4 + + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + from vllm.utils.deep_gemm import ( + fp8_fp4_sparse_mqa_logits, + get_sparse_mqa_logits_metadata, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 17, 32, 128 + sbk, num_candidates = 8, 2048 + lens = torch.randint(1, 1500, (rows,), dtype=torch.int32) + lens[0] = 1 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, sbk).cuda() + + sparse_idx, end = candidate_blocks_to_sparse_indices(candidates, ks, ke, sbk, sbk) + metadata = get_sparse_mqa_logits_metadata( + ks, ke, total, sparse_idx, torch.int8, sbk + ) + logits = fp8_fp4_sparse_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + metadata, + sparse_idx.shape[1], + sbk, + ) + + # PyTorch fp32 reference on the dequantized inputs. + q_deq = cast_back_from_fp4(q_fp, q_sf, gran_k=32, use_packed_ue8m0=True) + k_deq = cast_back_from_fp4(k_fp, k_sf, gran_k=32, use_packed_ue8m0=True) + dots = torch.einsum( + "rhd,td->rht", + q_deq.view(rows, num_heads, head_dim).float(), + k_deq.float(), + ) + ref = (dots.relu() * weights.float().unsqueeze(-1)).sum(1) # [rows, total] + + # Compare on valid sparse columns (absolute packed-workspace positions). + width = sparse_idx.shape[1] * sbk + cols = torch.arange(width, device="cuda") + valid = cols.unsqueeze(0) < end.unsqueeze(1).long() + pos = ( + sparse_idx.repeat_interleave(sbk, dim=1).long() * sbk + + (ks % sbk).unsqueeze(1) + + (cols % sbk).unsqueeze(0) + ).clamp(min=0, max=total - 1) + abs_err = (logits.float() - ref.gather(1, pos)).abs()[valid] + row_scale = ref.abs().amax(dim=1) + norm_err = ((logits.float() - ref.gather(1, pos)).abs() / row_scale.unsqueeze(1))[ + valid + ] + print( + f"sparse vs torch fp32: max_abs={abs_err.max().item():.4f}; " + f"rowmax-normalized p50={norm_err.median().item():.5f} " + f"p99={norm_err.quantile(0.99).item():.5f} max={norm_err.max().item():.5f}" + ) + # bf16-level numerics: far below the value gaps that change top-k selection. + assert norm_err.quantile(0.99).item() < 1e-2 + assert norm_err.max().item() < 2e-2 diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index f19ec35560f1..b4a544d2fc0f 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc +import importlib.machinery import inspect +import sys +import types from unittest.mock import Mock from weakref import WeakKeyDictionary, ref @@ -817,6 +820,110 @@ def materialize_with_sentinel(meta_tensor): assert "0.scale" not in model.state_dict() +@pytest.fixture +def hpc_rope_norm(monkeypatch, default_vllm_config): + """Import HpcRopeNorm with the external ``hpc`` package stubbed out.""" + if "hpc" not in sys.modules: + stub = types.ModuleType("hpc") + stub.__spec__ = importlib.machinery.ModuleSpec("hpc", loader=None) + stub.QuantType = types.SimpleNamespace( # type: ignore[attr-defined] + QPERTOKEN_PERHEAD_KPERTENSOR_VPERTENSOR=types.SimpleNamespace(value=0) + ) + monkeypatch.setitem(sys.modules, "hpc", stub) + from vllm.model_executor.layers.hpc import rope_norm + + monkeypatch.setattr(rope_norm, "_hpc_rope_norm_instances", {}) + return rope_norm + + +def test_hpc_rope_norm_kernel_sees_refit_norm_weights(monkeypatch, hpc_rope_norm): + """The fused HPC kernel is handed the live QK-norm weights after a refit. + + Drives the production ``_forward_impl`` with a recording ``hpc`` stub. The + Q/K norm weights it receives must be the model's own float32 parameters, + so a layerwise reload that rewrites them in place (same storage) is what + the kernel sees. Previously the kernel read separate mirrors that no + reload path refreshed. + """ + from vllm.model_executor.layers.layernorm import RMSNorm + + head_dim, num_heads, num_kv_heads, block_size = 128, 8, 1, 4 + layer = torch.nn.Module() + layer.q_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) + layer.k_norm = RMSNorm(head_dim, 1e-6, dtype=torch.float32) + layer.hpc_rope_norm = hpc_rope_norm.HpcRopeNorm( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + cos_sin_cache=torch.ones(16, head_dim), + use_qk_norm=True, + fallback_qnorm=layer.q_norm, + fallback_knorm=layer.k_norm, + kv_cache_dtype="auto", + layer_name="hpc_test_layer", + ) + model = torch.nn.Sequential(layer) + rnorm = layer.hpc_rope_norm + + calls: list[dict] = [] + + def record(*args, **kwargs): + calls.append(kwargs) + + monkeypatch.setattr(sys.modules["hpc"], "rope_norm_store_kv", record, raising=False) + + def kernel_norm_weights(): + q_size, kv_size = num_heads * head_dim, num_kv_heads * head_dim + qkv = torch.zeros(1, q_size + 2 * kv_size, dtype=torch.bfloat16) + kv_cache = torch.zeros( + 2, num_kv_heads, block_size, 2 * head_dim, dtype=torch.bfloat16 + ) + attn_layer = types.SimpleNamespace( + _k_scale=torch.ones(1), _v_scale=torch.ones(1) + ) + attn_metadata = types.SimpleNamespace( + num_actual_tokens=1, + num_decodes=1, + num_decode_tokens=1, + num_prefills=0, + num_prefill_tokens=0, + max_query_len=1, + decode_query_len=1, + qo_indptr=None, + qo_indptr_decode=None, + slot_mapping=torch.tensor([4]), + seq_lens=torch.tensor([1]), + block_table_tensor=torch.tensor([[1]]), + hpc_kv_written=False, + ) + output = torch.zeros(1, q_size, dtype=torch.bfloat16) + rnorm._forward_impl(qkv, kv_cache, attn_metadata, attn_layer, output) + return calls[-1]["q_norm_weight"], calls[-1]["k_norm_weight"] + + def loaded(value): + return torch.full((head_dim,), value, dtype=torch.bfloat16) + + default_weight_loader(layer.q_norm.weight, loaded(0.5)) + default_weight_loader(layer.k_norm.weight, loaded(0.25)) + q, k = kernel_norm_weights() + assert torch.equal(q, loaded(0.5).float()) + assert torch.equal(k, loaded(0.25).float()) + q_ptr, k_ptr = q.data_ptr(), k.data_ptr() + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + layer.q_norm.weight.weight_loader(layer.q_norm.weight, loaded(2.0)) + layer.k_norm.weight.weight_loader(layer.k_norm.weight, loaded(3.0)) + finalize_layerwise_reload(model, model_config=None) + + q, k = kernel_norm_weights() + assert q.dtype == k.dtype == torch.float32 + assert torch.equal(q, loaded(2.0).float()) + assert torch.equal(k, loaded(3.0).float()) + assert (q.data_ptr(), k.data_ptr()) == (q_ptr, k_ptr) + assert not hasattr(rnorm, "qnorm_weight") + + @pytest.mark.parametrize( "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])] ) diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 42accbef85c5..30d9909a5a2e 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -10,6 +10,7 @@ import torch.nn as nn from transformers import PretrainedConfig +from vllm.config import ParallelConfig from vllm.multimodal.processing import InputProcessingContext @@ -371,7 +372,7 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary(): attention_backend=None, kv_cache_dtype=None, ), - parallel_config=SimpleNamespace(tensor_parallel_size=1), + parallel_config=ParallelConfig(), attention_config=SimpleNamespace(backend=None), cache_config=SimpleNamespace(), load_config=SimpleNamespace(), diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 1bd9f6230353..2c13e07ff367 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -21,7 +21,7 @@ from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.kv_cache_interface import SlidingWindowSpec -from ....utils import ROCM_ENGINE_KWARGS +from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn" @@ -192,6 +192,7 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa "cudagraph_mode", [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], ) +@create_new_process_for_each_test(method="spawn") def test_voxtral_realtime_cudagraph( audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode ): diff --git a/tests/models/multimodal/processing/test_gemma4_unified.py b/tests/models/multimodal/processing/test_gemma4_unified.py index 67a81ddb7b41..fe544e233884 100644 --- a/tests/models/multimodal/processing/test_gemma4_unified.py +++ b/tests/models/multimodal/processing/test_gemma4_unified.py @@ -6,9 +6,13 @@ import pytest import torch from PIL import Image as PILImage +from transformers.models.gemma4 import Gemma4AudioFeatureExtractor from vllm.exceptions import VLLMValidationError -from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs +from vllm.model_executor.models.gemma4_mm import ( + Gemma4ImagePixelInputs, + _dummy_audio_num_samples, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import MultiModalFieldConfig @@ -204,3 +208,41 @@ def test_limit_mm_per_prompt( mm_items=processor.info.parse_mm_data(mm_data), hf_processor_mm_kwargs={}, ) + + +def test_dummy_audio_length_falls_back_for_the_unified_extractor(): + """The unified audio feature extractor does not define ``fft_length``. + + ``Gemma4DummyInputsBuilder.get_dummy_mm_data`` read that attribute + unconditionally, so memory profiling raised ``AttributeError`` for every + unified Gemma 4 deployment in which audio is the modality with the largest + per-item token count -- which is the case as soon as ``audio_seq_length`` + is raised above the video budget. + """ + + class _UnifiedFeatureExtractor: + sampling_rate = 16000 + + class _UnifiedProcessor: + feature_extractor = _UnifiedFeatureExtractor() + audio_seq_length = 3000 + audio_ms_per_token = 40 + + # 3000 tokens * 40 ms = 120 s at 16 kHz + assert _dummy_audio_num_samples(_UnifiedProcessor()) == 120 * 16000 + + +def test_dummy_audio_length_still_uses_fft_length_for_the_tower_extractor(): + """The mel/tower extractor keeps its existing behaviour.""" + + class _TowerFeatureExtractor(Gemma4AudioFeatureExtractor): + def __init__(self): # noqa: D107 - bypass the real feature-extractor init + self.fft_length = 512 + self.sampling_rate = 16000 + + class _TowerProcessor: + feature_extractor = _TowerFeatureExtractor() + audio_seq_length = 750 + audio_ms_per_token = 40 + + assert _dummy_audio_num_samples(_TowerProcessor()) == 512 diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index 1c1029b51a01..7cbf9791a7b0 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -574,7 +574,7 @@ def test_reasoning_end_token_ids_drop_multi_token_tool_opener( # the text lexer. assert thinking_parser.reasoning_end_token_ids == {_THINK_END_ID} assert thinking_parser.find_reasoning_end_offset([7, _THINK_END_ID, 8]) == 1 - assert thinking_parser.find_reasoning_end_offset([7, 8]) is None + assert thinking_parser.find_reasoning_end_offset([7, 8]) == 2 def test_streaming_reasoning_implicit_end(self, thinking_parser): chunks = [ diff --git a/tests/parser/engine/test_gemma4_streaming_reasoning.py b/tests/parser/engine/test_gemma4_streaming_reasoning.py index 74224c15b749..ae21d69fc303 100644 --- a/tests/parser/engine/test_gemma4_streaming_reasoning.py +++ b/tests/parser/engine/test_gemma4_streaming_reasoning.py @@ -1616,3 +1616,36 @@ def test_non_streaming_multiple_commas(self, multi_comma_parser, request_obj): assert result.tools_called is True args = json.loads(result.tool_calls[0].function.arguments) assert args["destination"] == "456 Oakwood Avenue, Rivermist, 83214" + + +class TestGemma4IsReasoningEnd: + @staticmethod + def _parser(thinking: bool): + return Gemma4Parser( + _make_tokenizer(_PLAIN_ANSWER_TOKENS), + chat_template_kwargs={"enable_thinking": thinking}, + ) + + @pytest.mark.parametrize( + ("ids", "thinking", "ended"), + [ + ([CHANNEL_START_ID, 3000, CHANNEL_END_ID], True, True), + ([CHANNEL_START_ID, 3000, TOOL_CALL_START_ID], True, True), + ([CHANNEL_START_ID, 3000], True, False), + ([NEW_TURN_ID, 9100], True, False), + ([CHANNEL_END_ID, NEW_TURN_ID, 9100], True, False), + ([NEW_TURN_ID, 9100], False, True), + # The bundled template's thinking-off prompt: closed channel. + ([NEW_TURN_ID, 9100, CHANNEL_START_ID, 3000, CHANNEL_END_ID], False, True), + # A marker still wins with thinking off: the grammar honours it. + ([CHANNEL_START_ID, 3000], False, False), + # A tool-call opener after the channel closed: still ended. + ( + [CHANNEL_START_ID, 3000, CHANNEL_END_ID, TOOL_CALL_START_ID, 9200], + True, + True, + ), + ], + ) + def test_is_reasoning_end(self, ids, thinking, ended): + assert self._parser(thinking).is_reasoning_end(ids) is ended diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index f27ce52cb37b..8fd7c4d2623f 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -185,8 +185,8 @@ def test_config_without_reasoning_has_empty_set(self): def test_find_reasoning_end_offset_returns_first_match(self): engine = _make_engine() assert engine.find_reasoning_end_offset([5, 201, 6, 201]) == 1 - assert engine.find_reasoning_end_offset([5, 6]) is None - assert engine.find_reasoning_end_offset([]) is None + assert engine.find_reasoning_end_offset([5, 6]) == 2 + assert engine.find_reasoning_end_offset([]) == 0 # Rejected-draft placeholders never match. assert engine.find_reasoning_end_offset([-1, 201]) == 1 diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py index 37e6b5d13090..3f285e4d0315 100644 --- a/tests/parser/engine/test_qwen3_reasoning.py +++ b/tests/parser/engine/test_qwen3_reasoning.py @@ -180,9 +180,8 @@ def test_prompt_tool_example_before_generation_think_not_end(self, parser): """Tool examples before the generation must not end reasoning.""" assert not parser.is_reasoning_end([_TOOL_CALL_ID, _TEXT_ID, _THINK_START_ID]) - def test_paired_tool_call_not_end(self, parser): - """Paired ... (from template) is NOT end.""" - assert not parser.is_reasoning_end( + def test_completed_tool_call_after_open_think_is_end(self, parser): + assert parser.is_reasoning_end( [_THINK_START_ID, 1, _TOOL_CALL_ID, 2, _TOOL_CALL_END_ID] ) @@ -196,6 +195,43 @@ def test_empty_ids(self, parser): assert not parser.is_reasoning_end([]) +class TestIsReasoningEndDerivedFromGrammar: + @pytest.mark.parametrize( + ("ids", "thinking", "ended"), + [ + # Explicit close, and the implicit close on . + ([_THINK_START_ID, 1, _THINK_END_ID], True, True), + ([_THINK_START_ID, 1, _TOOL_CALL_ID], True, True), + # Still inside the block. + ([_THINK_START_ID, 1], True, False), + ([_THINK_END_ID, _THINK_START_ID, 1], True, False), + # Uncommitted: a bare turn, an empty sequence. Thinking on starts + # in REASONING, so reasoning is expected. + ([_IM_END_ID, _IM_START_ID, 1], True, False), + ([], True, False), + # Thinking off starts in CONTENT with no way into REASONING, so + # a the template left behind is inert, and a silent tail + # is ended because no block is expected. + ([_THINK_START_ID, 1], False, True), + ([_THINK_END_ID, _THINK_START_ID, 1], False, True), + ([_IM_END_ID, _IM_START_ID, 1], False, True), + ([], False, True), + ], + ) + def test_is_reasoning_end(self, mock_tokenizer, ids, thinking, ended): + parser = Qwen3Parser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking} + ) + assert parser.is_reasoning_end(ids) is ended + + @pytest.mark.parametrize("thinking", [True, False]) + def test_wait_for_reasoning_follows_thinking(self, mock_tokenizer, thinking): + parser = Qwen3Parser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": thinking} + ) + assert parser.parser_engine_config.wait_for_reasoning is thinking + + class TestIsReasoningEndTurnBoundaries: """Reasoning markers from earlier turns must not mark the new turn's reasoning as finished (think-marker injection / interleaved-thinking diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index ae236d71f805..9f00adf78314 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -1668,3 +1668,148 @@ def test_suffix_match_at_module_boundary(prefix, ignored_layer, expected): ) is expected ) + + +_GLM5_MXFP4_WEIGHT = { + "dtype": "fp4", + "qscheme": "per_group", + "group_size": 32, + "scale_format": "e8m0", + "is_dynamic": False, +} +_GLM5_BLOCK_FP8_WEIGHT = { + "dtype": "fp8_e4m3", + "qscheme": "per_block", + "is_dynamic": False, + "block_size": [128, 128], + "symmetric": True, +} +_GLM5_BLOCK_FP8_INPUT = { + "dtype": "fp8_e4m3", + "qscheme": "per_group", + "is_dynamic": True, + "group_size": 128, + "symmetric": True, +} +_GLM5_GATE_UP = "language_model.model.layers.0.mlp.gate_up_proj" + + +def _glm5_mixed_precision_config() -> QuarkConfig: + """Global MXFP4 with the dense-MLP gate/up shards marked block-FP8.""" + fp8 = {"weight": _GLM5_BLOCK_FP8_WEIGHT, "input_tensors": _GLM5_BLOCK_FP8_INPUT} + return QuarkConfig( + { + "global_quant_config": { + "weight": _GLM5_MXFP4_WEIGHT, + "input_tensors": None, + }, + "layer_type_quant_config": {}, + "layer_quant_config": {"*mlp.gate_proj": fp8, "*mlp.up_proj": fp8}, + "exclude": [], + } + ) + + +class _GLM5RecordingParam: + """Stand-in for a BF16 target param that records what gets loaded.""" + + def __init__(self): + self.loaded_weight: torch.Tensor | None = None + self.loaded_shard: object = "unset" + + def weight_loader(self, param, loaded_weight, shard_id=None): + self.loaded_weight = loaded_weight + self.loaded_shard = shard_id + + +def _glm5_block_fp8(out_dim: int, in_dim: int, block: int = 128): + """Return a (fp8 weight, f32 per-block scale) pair.""" + weight = (torch.randn(out_dim, in_dim) * 0.1).to(torch.float8_e4m3fn) + scale = torch.rand(out_dim // block, in_dim // block, dtype=torch.float32) + 0.5 + return weight, scale + + +def test_glm5next_gate_up_proj_mapping_is_genuine_fusion(): + from vllm.models.glm5next.common.model import Glm5NextForConditionalGeneration + + mapping = Glm5NextForConditionalGeneration.packed_modules_mapping + assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"] + + +def test_glm5next_genuine_fusion_resolves_gate_up_to_block_fp8(): + # The fused module must expand to its real shard names so each resolves to + # the per-layer block-FP8 entry (rather than the global MXFP4 scheme). + config = _glm5_mixed_precision_config() + config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]} + _, _, scheme_cls = config.get_scheme_cls(LinearBase, _GLM5_GATE_UP) + assert scheme_cls is QuarkW8A8Fp8PerBlock + + +def test_glm5next_attn_loader_accepts_quark_weight_scale(): + # Regression for KeyError on '...kv_a_proj_with_mqa.weight_scale': the fused + # q_a/kv_a projection is kept BF16 and dequantized on load; the Quark scale + # name must be recognized and routed to the correct fused shard. + from vllm.models.glm5next.common.model import ( + _dequant_fp8_block, + _try_load_fp8_attn_proj, + ) + + prefix = "layers.3.self_attn" + target = _GLM5RecordingParam() + params_dict = {f"{prefix}.fused_qkv_a_proj.weight": target} + buf: dict = {} + loaded: set = set() + + weight, scale = _glm5_block_fp8(256, 256) + # fp8 weight arrives first -> buffered, nothing loaded yet. + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.kv_a_proj_with_mqa.weight", weight, buf, params_dict, loaded, 0 + ) + is True + ) + assert target.loaded_weight is None + # Quark-named scale completes the pair -> dequantize + load. + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.kv_a_proj_with_mqa.weight_scale", + scale, + buf, + params_dict, + loaded, + 0, + ) + is True + ) + assert target.loaded_weight is not None + assert target.loaded_weight.dtype == torch.bfloat16 + assert target.loaded_shard == 1 # kv_a is shard 1 of fused_qkv_a_proj + assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) + assert f"{prefix}.fused_qkv_a_proj.weight" in loaded + + +def test_glm5next_attn_loader_accepts_deepseek_weight_scale_inv(): + # DeepSeek scale name keeps working unchanged. + from vllm.models.glm5next.common.model import ( + _dequant_fp8_block, + _try_load_fp8_attn_proj, + ) + + prefix = "layers.7.self_attn" + target = _GLM5RecordingParam() + params_dict = {f"{prefix}.o_proj.weight": target} + buf: dict = {} + loaded: set = set() + + weight, scale = _glm5_block_fp8(128, 256) + _try_load_fp8_attn_proj( + f"{prefix}.o_proj.weight", weight, buf, params_dict, loaded, 0 + ) + assert ( + _try_load_fp8_attn_proj( + f"{prefix}.o_proj.weight_scale_inv", scale, buf, params_dict, loaded, 0 + ) + is True + ) + assert target.loaded_shard is None # o_proj is a direct (non-fused) proj + assert torch.equal(target.loaded_weight, _dequant_fp8_block(weight, scale, 128)) diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index 91dee5634b52..8bcf4c48ab59 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -6,13 +6,17 @@ """ import json +from types import SimpleNamespace +from unittest.mock import Mock import jsonschema import pytest from transformers import AutoModelForSeq2SeqLM +from vllm import CompletionOutput, RequestOutput from vllm.assets.audio import AudioAsset from vllm.entrypoints.llm import LLM +from vllm.logprobs import Logprob, SampleLogprobs from vllm.platforms import current_platform from vllm.sampling_params import BeamSearchParams, StructuredOutputsParams @@ -45,6 +49,87 @@ MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"] +@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) +@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) +def test_beam_search_abort_returns_partial_outputs_and_continues_other_prompts( + monkeypatch, + abort_after: int, + prompt_token: int, + terminal_logprobs: SampleLogprobs | None, +) -> None: + """Return tokens and scores from the last completed step for aborted prompts. + + Other prompts in the batch continue generating. + """ + + def run_requests(prompts, **kwargs): + results = [] + for prompt in prompts: + tokens = prompt["prompt_token_ids"] + if tokens[0] == prompt_token: + assert len(tokens) <= abort_after + 1, "Continued after abort" + aborted = ( + tokens[0] == prompt_token + and len(tokens) == abort_after + 1 + and tokens[-1] != 12 + ) + results.append( + RequestOutput( + request_id="inner", + prompt=None, + prompt_token_ids=tokens, + prompt_logprobs=None, + finished=True, + outputs=[ + CompletionOutput( + index=0, + text="", + token_ids=[] if aborted else [11], + cumulative_logprob=None, + logprobs=terminal_logprobs + if aborted + else [{11: Logprob(-0.1), 12: Logprob(-0.2)}], + finish_reason="abort" if aborted else "length", + ) + ], + ) + ) + return results + + llm = LLM.__new__(LLM) + llm.llm_engine = Mock() + tokenizer = SimpleNamespace( + eos_token_id=0, + decode=lambda tokens, skip_special_tokens=False: " ".join( + str(token) for token in tokens if not skip_special_tokens or token != 0 + ), + ) + llm.renderer = Mock(get_tokenizer=Mock(return_value=tokenizer)) + monkeypatch.setattr(llm, "_preprocess_cmpl", lambda prompts: prompts) + monkeypatch.setattr(llm, "_render_and_run_requests", run_requests) + prompts = [ + {"type": "token", "prompt_token_ids": [token]} for token in [prompt_token, 2] + ] + params = BeamSearchParams(beam_width=2, max_tokens=3) + outputs = llm.beam_search(prompts, params) + + aborted, normal = outputs + expected_tokens = ( + [[prompt_token, 11], [prompt_token, 12]] if abort_after else [[prompt_token]] + ) + assert [beam.tokens for beam in aborted.sequences] == expected_tokens + for beam in aborted.sequences: + assert beam.finish_reason == "abort" + assert beam.text == tokenizer.decode( + beam.tokens, skip_special_tokens=params.skip_special_tokens + ) + assert len(beam.logprobs) == abort_after + assert aborted.sequences[0].cum_logprob == pytest.approx(-0.1 * abort_after) + assert len(normal.sequences) == 2 + assert all(len(beam.tokens) == 4 for beam in normal.sequences) + assert all(beam.finish_reason != "abort" for beam in normal.sequences) + + @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", MAX_TOKENS) diff --git a/tests/samplers/test_beam_search_online.py b/tests/samplers/test_beam_search_online.py index 102e66dcbe79..ccb40cc65906 100644 --- a/tests/samplers/test_beam_search_online.py +++ b/tests/samplers/test_beam_search_online.py @@ -5,7 +5,7 @@ from vllm import CompletionOutput, RequestOutput from vllm.entrypoints.generate.beam_search.online import BeamSearchOnlineMixin -from vllm.logprobs import Logprob +from vllm.logprobs import Logprob, SampleLogprobs from vllm.sampling_params import BeamSearchParams @@ -108,3 +108,55 @@ async def test_beam_search_respects_skip_special_tokens( assert outputs[0].outputs[0].text == expected_text assert outputs[0].outputs[0].token_ids == [_Tokenizer.special_token_id] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("abort_after", "prompt_token"), [(0, 1), (1, 1), (0, 0)]) +@pytest.mark.parametrize("terminal_logprobs", [None, [], [{11: Logprob(-0.1)}]]) +async def test_beam_search_abort_returns_partial_outputs( + monkeypatch, + abort_after: int, + prompt_token: int, + terminal_logprobs: SampleLogprobs | None, +) -> None: + """Abort returns each beam's tokens and scores from the last completed step.""" + calls = 0 + + async def generate(prompt, *args, **kwargs): + nonlocal calls + calls += 1 + depth = len(prompt["prompt_token_ids"]) - 1 + assert depth <= abort_after, "Beam search continued after abort" + result = await anext(_EngineClient().generate(prompt, *args, **kwargs)) + output = result.outputs[0] + output.token_ids = [11] + output.logprobs = [{11: Logprob(-0.1), 12: Logprob(-0.2)}] + output.finish_reason = "length" + if depth == abort_after and prompt["prompt_token_ids"][-1] != 12: + output.token_ids = [] + output.logprobs = terminal_logprobs + output.finish_reason = "abort" + yield result + + serving = _Serving() + monkeypatch.setattr(serving.engine_client, "generate", generate) + prompt = {"type": "token", "prompt": "prompt", "prompt_token_ids": [prompt_token]} + outputs = [ + output + async for output in serving.beam_search( + prompt, "request", BeamSearchParams(beam_width=2, max_tokens=4) + ) + ] + + assert calls == 1 + 2 * abort_after + assert len(outputs) == 1 + assert outputs[0].finished + assert outputs[0].request_id == "request" + expected_tokens = [[11], [12]] if abort_after else [[]] + assert [output.token_ids for output in outputs[0].outputs] == expected_tokens + for output in outputs[0].outputs: + assert output.finish_reason == "abort" + assert output.text == _Tokenizer().decode(list(output.token_ids)) + assert output.logprobs is not None + assert len(output.logprobs) == abort_after + assert outputs[0].outputs[0].cumulative_logprob == pytest.approx(-0.1 * abort_after) diff --git a/tests/test_config.py b/tests/test_config.py index 4d39b56d265f..f75592f089e2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -337,6 +337,108 @@ def test_hisparse_rejects_non_cuda(monkeypatch): VllmConfig(attention_config=AttentionConfig(hisparse_config=HiSparseConfig())) +@pytest.mark.parametrize( + "kv_transfer_config", + [ + KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_both", + kv_connector_extra_config={ + "connectors": [ + { + "kv_connector": "OffloadingConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, + }, + { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 128}, + }, + ] + }, + ), + KVTransferConfig( + kv_connector="MultiConnector", + kv_role="kv_consumer", + kv_connector_extra_config={ + "connectors": [ + {"kv_connector": "NixlConnector", "kv_role": "kv_consumer"}, + { + "kv_connector": "HiSparseConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"host_pool_gib": 128}, + }, + ] + }, + ), + ], + ids=["standalone", "multi-connector", "pd-decode"], +) +def test_hisparse_connector_implies_attention_config(monkeypatch, kv_transfer_config): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + config = VllmConfig( + kv_transfer_config=kv_transfer_config, + # Skip the HMA auto-detect block: it imports the connector class, + # which pulls the platform's compiled attention extensions. + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=False, + ), + ) + assert isinstance(config.attention_config.hisparse_config, HiSparseConfig) + + +def test_hisparse_connector_preserves_explicit_attention_config(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr("vllm.config.vllm.HAS_TRITON", True) + config = VllmConfig( + attention_config=AttentionConfig( + hisparse_config=HiSparseConfig(device_buffer_size=512) + ), + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ), + scheduler_config=SchedulerConfig( + max_model_len=2048, + is_encoder_decoder=False, + disable_hybrid_kv_cache_manager=False, + ), + ) + assert config.attention_config.hisparse_config.device_buffer_size == 512 + + +def test_hisparse_connector_without_cuda_still_rejected(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: False) + with pytest.raises(ValueError, match="requires NVIDIA CUDA"): + VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 128}, + ) + ) + + +def test_no_hisparse_connector_keeps_attention_config_unset(monkeypatch): + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + config = VllmConfig( + kv_transfer_config=KVTransferConfig( + kv_connector="OffloadingConnector", kv_role="kv_both" + ) + ) + assert config.attention_config.hisparse_config is None + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" diff --git a/tests/test_request_input_bounds.py b/tests/test_request_input_bounds.py index 44f536e7fbe7..990f10cfebe3 100644 --- a/tests/test_request_input_bounds.py +++ b/tests/test_request_input_bounds.py @@ -12,6 +12,7 @@ import sys from collections.abc import Callable from typing import Protocol +from unittest.mock import Mock import pytest from pydantic import ValidationError @@ -27,10 +28,39 @@ from vllm.exceptions import VLLMValidationError from vllm.sampling_params import BeamSearchParams from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding +from vllm.v1.engine.input_processor import InputProcessor pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] +@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127]) +def test_routed_experts_offset_validated_before_engine_submission(offset): + """Reject invalid offsets even when routed-expert output is disabled.""" + processor = Mock(spec=InputProcessor) + processor.tokenizer = None + processor.generation_config_fields = {} + processor.renderer = Mock() + processor.renderer.get_eos_token_id.return_value = None + processor.vllm_config = Mock() + params = SamplingParams(routed_experts_prompt_start=offset) + prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]} + + if not 0 <= offset <= 3: + with pytest.raises(VLLMValidationError, match="routed_experts_prompt_start"): + InputProcessor.process_inputs( + processor, "invalid", prompt, params, ("generate",) + ) + params.routed_experts_prompt_start = 0 + + request = InputProcessor.process_inputs( + processor, "valid", prompt, params, ("generate",) + ) + assert ( + request.sampling_params.routed_experts_prompt_start + == params.routed_experts_prompt_start + ) + + # --- Stop strings: public requests cap the number of stop strings --------- diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 8d74f0431934..97ed0cd69d4c 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -897,3 +897,111 @@ def test_streaming_trailing_bare_bool_not_duplicated(self, parser, mock_request) } assert args_text.count("replace_all") == 1 + + +class TestBareAndMalformedOpeners: + """Tests for bare and malformed tool call openers. + + See https://github.com/vllm-project/vllm/issues/53431. + """ + + def test_bare_opener_non_streaming(self, parser, mock_request): + model_output = '<|tool_call>:get_weather{location:<|"|>London<|"|>}' + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "London"} + + def test_documented_form_still_works(self, parser, mock_request): + model_output = ( + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + assert result.tool_calls[0].function.name == "get_weather" + + def test_bare_opener_multiple_args(self, parser, mock_request): + model_output = ( + "<|tool_call>:get_weather{" + 'location:<|"|>Tokyo<|"|>,unit:<|"|>celsius<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "get_weather" + args = json.loads(result.tool_calls[0].function.arguments) + assert args == {"location": "Tokyo", "unit": "celsius"} + + def test_malformed_no_brace_bounded(self, parser, mock_request): + model_output = "<|tool_call>call:bad_func no brace" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 1 + name = result.tool_calls[0].function.name + assert "" not in name + + def test_malformed_then_well_formed(self, parser, mock_request): + model_output = ( + "<|tool_call>call:bad_func no brace" + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "bad_func no brace" + assert result.tool_calls[1].function.name == "get_weather" + assert "London" in result.tool_calls[1].function.arguments + + def test_bare_malformed_then_well_formed(self, parser, mock_request): + model_output = ( + "<|tool_call>:bad_func no brace" + '<|tool_call>call:get_weather{location:<|"|>London<|"|>}' + "" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 2 + assert result.tool_calls[0].function.name == "bad_func no brace" + assert result.tool_calls[1].function.name == "get_weather" + + def test_three_mixed_calls(self, parser, mock_request): + model_output = ( + "<|tool_call>:bad no brace" + '<|tool_call>call:search{input:<|"|>hello<|"|>}' + "<|tool_call>:set{flag:true,count:5}" + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert len(result.tool_calls) == 3 + assert result.tool_calls[0].function.name == "bad no brace" + assert result.tool_calls[1].function.name == "search" + assert result.tool_calls[2].function.name == "set" + + def test_colon_in_content_not_broken(self, parser, mock_request): + model_output = ( + "Hello: world" + '<|tool_call>:get_weather{location:<|"|>London<|"|>}' + ) + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.content is not None + assert "Hello" in result.content + + def test_bare_opener_empty_args(self, parser, mock_request): + model_output = "<|tool_call>:set_status{}" + result = parser.extract_tool_calls(model_output, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "set_status" diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index b77b4efaffe6..88e1f3f3832c 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -5,7 +5,7 @@ from unittest.mock import MagicMock import pytest -from xgrammar import Grammar, StructuralTag +from xgrammar import Grammar, StructuralTag, builtin_structural_tag from xgrammar.testing import _is_grammar_accept_string from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -81,10 +81,14 @@ def test_supported_structural_tag_models_include_vllm_builtins(): @pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) -def test_deepseek_v41_format_constraints_ignore_parameter_schema( +def test_deepseek_v41_older_xgrammar_keeps_format_constraints( choice, sample_tools_strict, + monkeypatch, ): + monkeypatch.setattr( + builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None, raising=False + ) tool_choice = ( ChatCompletionNamedToolChoiceParam( function=ChatCompletionNamedFunction(name=choice) @@ -128,7 +132,11 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): reasoning=False, ) grammar = Grammar.from_structural_tag(tag) - call = '<｜DSML｜ invoke name="get_weather">\n\n' + call = ( + '<｜DSML｜ invoke name="get_weather">\n' + '<｜DSML｜ parameter name="city" string="true">Paris\n' + "\n" + ) assert _is_grammar_accept_string( grammar, "\n\n<｜DSML｜ calls>\n" + call + "" ) @@ -141,6 +149,72 @@ def test_deepseek_v41_named_choice_emits_one_call(sample_tools): ) +@pytest.mark.skipif( + not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), + reason="Requires XGrammar's DeepSeek V4.1 builtin", +) +@pytest.mark.parametrize("choice", ["auto", "required", "get_weather"]) +@pytest.mark.parametrize("reasoning", [False, True]) +def test_deepseek_v41_constrains_parameters_after_reasoning( + sample_tools_strict, choice, reasoning +): + """Strict calls must enforce the schema without requiring another think close.""" + tool_choice = ( + ChatCompletionNamedToolChoiceParam( + function=ChatCompletionNamedFunction(name=choice) + ) + if choice == "get_weather" + else choice + ) + sample_tools_strict[0].function.parameters["additionalProperties"] = False + tag = get_model_structural_tag( + "deepseek_v41", sample_tools_strict, tool_choice, reasoning=reasoning + ) + grammar = Grammar.from_structural_tag(tag) + begin = '\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name="get_weather">\n' + end = "\n" + parameter = ( + '<｜DSML｜ parameter name="city" string="true">Paris\n' + ) + assert _is_grammar_accept_string(grammar, begin + parameter + end) + for invalid in ( + "", # missing required city + parameter * 2, + parameter.replace('name="city"', 'name="unknown"'), + parameter.replace('string="true">Paris', 'string="false">42'), + parameter.replace("｜ parameter", "｜parameter"), + '{"city":"Paris"}', # the encoder emits DSML parameters, not a JSON body + ): + assert not _is_grammar_accept_string(grammar, begin + invalid + end) + assert not _is_grammar_accept_string( + grammar, "reason" + begin + parameter + end + ) + assert _is_grammar_accept_string(grammar, "Hello") == (choice == "auto") + + +@pytest.mark.skipif( + not hasattr(builtin_structural_tag, "get_deepseek_v4_1_structural_tag"), + reason="Requires XGrammar's DeepSeek V4.1 builtin", +) +def test_deepseek_v41_non_strict_parallel_calls_keep_typed_dsml(sample_tools): + sample_tools[0].function.strict = False + tag = get_model_structural_tag( + "deepseek_v41", sample_tools, "required", reasoning=False + ) + grammar = Grammar.from_structural_tag(tag) + call = ( + '<｜DSML｜ invoke name="get_weather">\n' + '<｜DSML｜ parameter name="extra" string="false">' + '{"nested":[true,null,1.5]}\n' + "\n" + ) + output = "\n\n<｜DSML｜ calls>\n" + call * 2 + "" + assert _is_grammar_accept_string(grammar, output) + assert not _is_grammar_accept_string( + grammar, output.replace('{"nested":[true,null,1.5]}', "invalid") + ) + + @pytest.mark.parametrize("model", sorted(XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS)) def test_get_model_structural_tag_supports_all_xgrammar_builtins( model: str, diff --git a/tests/utils_/test_gpu_sync_debug.py b/tests/utils_/test_gpu_sync_debug.py index dc8b12280010..10b550c0a174 100644 --- a/tests/utils_/test_gpu_sync_debug.py +++ b/tests/utils_/test_gpu_sync_debug.py @@ -123,6 +123,24 @@ def test_suppressing_works_while_compiling(monkeypatch): with_gpu_sync_check(suppressed)() +@create_new_process_for_each_test() +def test_copy_checker_is_traceable_while_compiling(monkeypatch): + """The copy wrapper must bypass ContextVars while Dynamo is tracing.""" + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error") + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + gsd._install_copy_checkers() + + def convert_like(x, like): + return x.to(device=like.device, dtype=like.dtype) + + compiled = torch.compile(convert_like, fullgraph=True) + result = compiled( + torch.ones(4, device="cuda", dtype=torch.float64), + torch.ones(4, device="cuda", dtype=torch.float32), + ) + assert result.dtype == torch.float32 + + @create_new_process_for_each_test() def test_sync_debug_mode_restored_after_checked_call(monkeypatch): """The mode is armed only for the duration of a checked call. Leaving it diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index 6901a162a407..79dd3e320ae6 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest import torch from vllm.utils.torch_utils import ( OMP_NUM_THREADS_SET_BY_VLLM, + async_tensor_h2d, available_cpu_count, common_broadcastable_dtype, current_stream, @@ -192,3 +194,26 @@ def test_runtime_threads_override_vllm_set_omp_num_threads( torch.set_num_threads(3) set_torch_threads_for_runtime() assert torch.get_num_threads() == 1 + + +@pytest.mark.parametrize( + "device", ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"] +) +def test_async_tensor_h2d_staging(device): + """Staging must preserve values while normalizing dtype and layout.""" + # Non-contiguous source with a dtype conversion: would previously sync. + src = torch.arange(12, dtype=torch.int64).reshape(3, 4).T + result = async_tensor_h2d(src, device=device, dtype=torch.int32) + assert result.dtype == torch.int32 + assert result.is_contiguous() + assert torch.equal(result.cpu(), src.to(torch.int32)) + + # numpy and list sources still work. + assert torch.equal( + async_tensor_h2d(np.arange(4, dtype=np.int64), device=device), + torch.arange(4, device=device), + ) + assert torch.equal( + async_tensor_h2d([1, 2, 3], device=device, dtype=torch.int32), + torch.tensor([1, 2, 3], dtype=torch.int32, device=device), + ) diff --git a/tests/v1/attention/test_dcp_a2a_pack_mask.py b/tests/v1/attention/test_dcp_a2a_pack_mask.py index 36032e7a69a0..cff1228baf03 100644 --- a/tests/v1/attention/test_dcp_a2a_pack_mask.py +++ b/tests/v1/attention/test_dcp_a2a_pack_mask.py @@ -13,7 +13,7 @@ mask_dcp_empty_shards_, ) -pytestmark = pytest.mark.skipif( +requires_accelerator = pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="needs a CUDA or ROCm device for the Triton pack kernel", ) @@ -44,6 +44,7 @@ def _pack(out, lse, world_size, h_per_rank, head_dim, seq_lens, query_start_loc) @pytest.mark.parametrize("world_size,h_per_rank", [(8, 2), (8, 16), (4, 4), (2, 1)]) @pytest.mark.parametrize("tokens_per_req", [1, 2, 3]) @pytest.mark.parametrize("num_pad_rows", [0, 2]) +@requires_accelerator @pytest.mark.parametrize( "seq_lens_list", [ @@ -97,6 +98,7 @@ def test_fused_mask_matches_eager( @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("seed", range(8)) +@requires_accelerator def test_fused_mask_matches_eager_ragged(dtype, seed): """Non-uniform per-request query lengths. @@ -137,6 +139,7 @@ def test_fused_mask_matches_eager_ragged(dtype, seed): ) +@requires_accelerator def test_mask_disabled_is_unmasked(): """Passing no seq_lens/query_start_loc must leave the LSE untouched.""" torch.manual_seed(0) @@ -149,3 +152,14 @@ def test_mask_disabled_is_unmasked(): packed = _pack(out, lse, world_size, h_per_rank, head_dim, None, None) assert torch.isfinite(packed.float()).all() + + +def test_eager_mask_handles_rank_with_no_local_sequences(): + """Padded graph rows are all empty when a DCP rank has no sequences.""" + lse = torch.randn(8, 4, dtype=torch.float32) + seq_lens = torch.empty(0, dtype=torch.int32) + query_start_loc = torch.tensor([0], dtype=torch.int32) + + mask_dcp_empty_shards_(lse, seq_lens, query_start_loc) + + assert torch.isneginf(lse).all() diff --git a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py index b20c16330659..8dc082b08bbc 100644 --- a/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py +++ b/tests/v1/attention/test_flashinfer_mla_sparse_sm90.py @@ -151,32 +151,44 @@ def test_forward_wiring(monkeypatch, qk_rope, kv_dtype): assert kwargs == {} -def test_builder_attaches_its_state(monkeypatch): +@pytest.mark.parametrize("use_mha", [False, True]) +@pytest.mark.parametrize("num_decodes", [0, 1]) +def test_builder_plans_only_rows_dispatched_to_mqa(monkeypatch, use_mha, num_decodes): + """MHA prefill rows must not make the MQA kernel read beyond its query.""" builder = object.__new__(FlashInferMLASparseSM90Builder) builder._index_topk = 2048 builder._index_kpool = 4 builder._async_scheduling = False builder.state = FakeState(TOPK) + builder._attention_layer = SimpleNamespace(_use_sparse_mha=lambda _: use_mha) metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata) metadata.state = None + metadata.num_prefills = 1 + metadata.num_decode_tokens = num_decodes monkeypatch.setattr( sm90_mod.FlashInferMLASparseMetadataBuilder, "build", lambda *_args, **_kwargs: metadata, ) cam = SimpleNamespace( - num_reqs=1, - query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), - seq_lens=torch.tensor([1], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32), + num_reqs=num_decodes + 1, + query_start_loc_cpu=torch.tensor( + [0, 1, 6] if num_decodes else [0, 5], dtype=torch.int32 + ), + seq_lens_cpu_upper_bound=torch.tensor( + [1402, 5] if num_decodes else [5], dtype=torch.int32 + ), positions=None, ) result = builder.build(0, cam) assert result.state is builder.state - assert builder.state.plan_calls[0][0] == 1 - assert builder.state.plan_calls[0][1].tolist() == [1] + expected_lens = [1402] if num_decodes else [] + if not use_mha: + expected_lens += [1, 2, 3, 4, 5] + assert builder.state.plan_calls[0][0] == len(expected_lens) + assert builder.state.plan_calls[0][1].tolist() == expected_lens def test_plan_uses_state_params(monkeypatch): diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py new file mode 100644 index 000000000000..5b63c2c609f9 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_sparse_mla_workspace.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only tests for FlashInfer sparse MLA backend constraints.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + FlashInferMLASparseTRTLLMBackend, + _required_workspace_bytes, + compute_trtllm_sparse_mla_workspace_bytes, +) + +# Reporter config from vllm#50781: GLM-5.2 (64 q heads), TP=8, DCP=8, +# max_num_batched_tokens=16384, FlashInfer 0.6.14. +_REPORTER_HEADS_PER_RANK = 64 // 8 +_REPORTER_DCP = 8 +_REPORTER_MAX_TOKENS = 16384 +# From the reported crash: trtllm-gen requested exactly this many bytes for +# trtllm_gen_softmax_workspace after an 8 MiB counter carve had left +# 404,750,336 of the 413,138,944-byte default buffer. +_REPORTER_OBSERVED_SOFTMAX_BYTES = 1_611_661_312 +_REPORTER_OBSERVED_PRECARVE_BYTES = 8_388_608 + + +def test_reporter_softmax_carve_is_reproduced_byte_exactly(): + # The crashing step scheduled 12288 tokens (a single 12K-token request); + # the softmax slab formula must reproduce the kernel's request exactly: + # 8 * (8 heads/rank * 8 dcp) * 12288 tokens * 256 + 1 MiB guard. + softmax_bytes = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=0, + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=12288, + ) + assert softmax_bytes == _REPORTER_OBSERVED_SOFTMAX_BYTES + + +def test_reporter_config_covers_observed_overflow(): + computed = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert computed >= ( + _REPORTER_OBSERVED_SOFTMAX_BYTES + _REPORTER_OBSERVED_PRECARVE_BYTES + ) + + +def test_non_dcp_size_is_unchanged(): + computed = compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes=_DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size=1, + num_heads_per_rank=128, + max_num_batched_tokens=65536, + ) + assert computed == _DEFAULT_WORKSPACE_BUFFER_SIZE + + +def test_combined_pcp_dcp_rejects_flashinfer_sparse(): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=4, + decode_context_parallel_size=4, + ), + model_config=None, + ) + with patch("vllm.config.get_current_vllm_config", return_value=config): + reason = FlashInferMLASparseTRTLLMBackend.supports_combination( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SimpleNamespace(major=10), + ) + + assert reason is not None + assert "use FLASHMLA_SPARSE" in reason + + +@pytest.mark.parametrize("pcp_size,dcp_size", [(4, 1), (1, 4)]) +def test_single_context_parallel_mode_allows_flashinfer_sparse( + pcp_size: int, dcp_size: int +): + config = SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=dcp_size, + ), + model_config=None, + ) + with patch("vllm.config.get_current_vllm_config", return_value=config): + reason = FlashInferMLASparseTRTLLMBackend.supports_combination( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=64, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + device_capability=SimpleNamespace(major=10), + ) + + assert reason is None + + +def test_default_constant_matches_envs_default(monkeypatch): + from vllm import envs + + monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False) + assert _DEFAULT_WORKSPACE_BUFFER_SIZE == envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + + +def test_env_unset_returns_computed(monkeypatch): + monkeypatch.delenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", raising=False) + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + _REPORTER_DCP, + _REPORTER_HEADS_PER_RANK, + _REPORTER_MAX_TOKENS, + ) + + +def test_env_override_below_computed_is_respected_with_warning(monkeypatch): + override = 100 * 1024 * 1024 + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override)) + with patch( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger" + ) as mock_logger: + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == override + mock_logger.warning_once.assert_called_once() + + +def test_env_override_above_computed_is_respected_without_warning(monkeypatch): + override = 8 * 1024 * 1024 * 1024 + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(override)) + with patch( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse.logger" + ) as mock_logger: + required = _required_workspace_bytes( + dcp_world_size=_REPORTER_DCP, + num_heads_per_rank=_REPORTER_HEADS_PER_RANK, + max_num_batched_tokens=_REPORTER_MAX_TOKENS, + ) + assert required == override + mock_logger.warning_once.assert_not_called() diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index eab4759640a6..8bc409adb704 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -307,6 +307,23 @@ def test_get_dcp_local_seq_lens_preserves_mtp_bounds_shape(): torch.testing.assert_close(actual, expected) +def test_get_dcp_local_seq_lens_reuses_device_rank_tensor(monkeypatch): + seq_lens = torch.tensor([0, 1, 7, 8, 17], dtype=torch.int32) + world = 4 + rank = 2 + rank_tensor = torch.tensor(rank, dtype=torch.int32) + expected = get_dcp_local_seq_lens(seq_lens, world, rank, 2) + + def fail_tensor_construction(*args, **kwargs): + raise AssertionError("rank tensor must be reused") + + with monkeypatch.context() as patch: + patch.setattr(torch, "tensor", fail_tensor_construction) + actual = get_dcp_local_seq_lens(seq_lens, world, rank_tensor, 2) + + torch.testing.assert_close(actual, expected) + + def test_get_dcp_local_seq_lens_must_run_after_decode_expansion(): world = 2 rank = 1 @@ -979,7 +996,7 @@ def test_pcp_plan_deinterleave_restores_global_order( monkeypatch.setattr("vllm.v1.attention.backends.mla.indexer.PIN_MEMORY", False) monkeypatch.setattr( - "vllm.v1.attention.backends.mla.indexer.async_copy_to_gpu", + "vllm.v1.attention.backends.mla.indexer.async_tensor_h2d", lambda x, device: x.to(device), ) scheduled = np.array(req_lens, dtype=np.int64) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 87df1e409cf3..f3b4b7f5ba24 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -179,6 +179,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, + prefill_context_parallel_size=1, cp_kv_cache_interleave_size=1, ), ) @@ -193,6 +194,37 @@ def test_indexer_warmup_normalizes_zero_compress_ratios(): } +def test_indexer_warmup_includes_pcp_normalized_dcp_key(monkeypatch): + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.indexer.get_dcp_group", + lambda: SimpleNamespace(rank_in_group=2), + ) + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(compress_ratios=[32], index_kpool=32) + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=4, + prefill_context_parallel_size=4, + cp_kv_cache_interleave_size=1, + ), + ) + + keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config) + + # Triton's compile key normalizes generic i32 values to 2 and divisible + # i32 values (including zero) to 16. + assert {(key.dcp_rank, key.dcp_world) for key in keys} == {(2, 2), (16, 1)} + assert { + ( + key.input_variant.is_aligned("uncompressed_seq_lens"), + key.input_variant.is_aligned("cu_compressed_seq_lens"), + ) + for key in keys + } == {(False, False), (False, True), (True, False), (True, True)} + + def test_compressed_slot_mapping_warmup_includes_index_kpool(): config = SimpleNamespace( cache_config=SimpleNamespace(block_size=256), diff --git a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py index 06684302cf39..a86c7e617be1 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py +++ b/tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py @@ -176,3 +176,8 @@ def test_pad_unpad_round_trip_preserves_head_order(num_heads): assert unpadded.shape == q.shape torch.testing.assert_close(unpadded, q) + + +def test_a_non_causal_block_never_routes_to_gluon(gluon_available): + """A small-head block that normally uses Gluon must use ASM when non-causal.""" + assert not AiterMLAHelper.use_gluon_verify(12, 8, "auto", causal=False) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 112cdff324df..efffcad927db 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -130,6 +130,7 @@ def _builder( _mla_q_dtype=torch.bfloat16, _mla_kv_dtype=torch.bfloat16, decode_attn_out_dtype=torch.bfloat16, + _decode_causal=True, ) # Bound to the stub rather than faked: the verify flatten's per-row view is # part of what _build_decode is being tested for. @@ -155,6 +156,25 @@ def test_backend_declares_uniform_batch_support(): ) +@pytest.mark.parametrize("causal", [True, False]) +def test_build_hands_common_causal_flag_to_aiter_builder(monkeypatch, causal): + seen = [] + + def fake_common_build(self, *args): + seen.append(self._decode_causal) + return SimpleNamespace(decode=None, prefill=None) + + monkeypatch.setattr( + rocm_aiter_mla.MLACommonMetadataBuilder, "build", fake_common_build + ) + builder = object.__new__(AiterMLAMetadataBuilder) + builder._fp8_prefill_enabled = False + + builder.build(0, SimpleNamespace(causal=causal)) + + assert seen == [causal] + + def test_dcp_verify_row_view_is_causal_per_row(): """Row t must cover this rank's share of positions [0, seq_len - qlen + t]. @@ -949,3 +969,81 @@ def test_persistent_metadata_gate_without_gluon_build( assert metadata.has_persistent_metadata is expect_persistent assert get_mla_metadata_v1.called is expect_persistent + + +def _build_non_causal(monkeypatch, *, num_heads, kv_cache_dtype, qlen, mtp_qlen): + """Drive _build_decode with causal=False and hand back the aiter mock.""" + get_mla_metadata_v1 = mock.MagicMock() + monkeypatch.setitem( + sys.modules, + "aiter", + SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1), + ) + monkeypatch.setattr( + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel() + ) + builder = _builder( + mtp_decode_qlen=mtp_qlen, + kv_cache_dtype=kv_cache_dtype, + num_heads=num_heads, + ) + builder._decode_causal = False + metadata = AiterMLAMetadataBuilder._build_decode( + builder, + block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8), + seq_lens_device=torch.tensor([7, 5], dtype=torch.int32), + max_seq_len=7, + query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32), + num_decode_tokens=2 * qlen, + dcp_tot_seq_lens_device=None, + ) + return metadata, get_mla_metadata_v1 + + +def test_non_causal_build_hands_the_mask_to_aiter(monkeypatch): + """The schedule carries the mask, so is_causal must follow the block.""" + metadata, get_mla_metadata_v1 = _build_non_causal( + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=8, mtp_qlen=8 + ) + assert metadata.has_persistent_metadata + # 6th positional arg of get_mla_metadata_v1 is is_causal. + assert get_mla_metadata_v1.call_args.args[5] is False + + +@pytest.mark.parametrize("num_heads", [12, 16, 48]) +def test_a_two_token_fp8_block_is_refused_without_a_fold(monkeypatch, num_heads): + """16-head (and padded-to-16) qlen-2 fp8 has no non-causal kernel. + + 48/80/112 (and 33 padded to 48) fold to H16 while keeping Q2, so they + must be refused the same way as native 16-head qlen 2. + """ + with pytest.raises(ValueError, match="2-token"): + _build_non_causal( + monkeypatch, + num_heads=num_heads, + kv_cache_dtype="fp8", + qlen=2, + mtp_qlen=8, + ) + + +def test_a_two_token_fp8_block_is_allowed_when_folded(monkeypatch): + """32 heads at qlen 2 folds onto the 16-head / 4-token kernel.""" + metadata, get_mla_metadata_v1 = _build_non_causal( + monkeypatch, + num_heads=32, + kv_cache_dtype="fp8", + qlen=2, + mtp_qlen=8, + ) + assert metadata.has_persistent_metadata + assert get_mla_metadata_v1.called + + +def test_a_two_token_bf16_block_is_allowed(monkeypatch): + """The bf16 fold has no such hole, so the guard must not fire there.""" + metadata, _ = _build_non_causal( + monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8 + ) + assert metadata.has_persistent_metadata diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 2ca839d59c50..6d2c4fbd5880 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -22,10 +22,17 @@ create_vllm_config, ) from vllm import _custom_ops as ops -from vllm.config import HiSparseConfig, SpeculativeConfig, set_current_vllm_config +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.config import ( + CUDAGraphMode, + HiSparseConfig, + SpeculativeConfig, + set_current_vllm_config, +) from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( HiSparseConnectorWorker, ) +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import mla_attention from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha from vllm.model_executor.layers.attention.sparse_mla_attention import ( @@ -897,6 +904,75 @@ def test_triton_convert_req_index_to_global_index_decode_only( torch.testing.assert_close(result, reference_result, rtol=0, atol=0) +def test_index_group_convert_during_piecewise_capture(): + """In piecewise cudagraph mode the indexer is captured while the convert + runs in the following eager break. The convert's side stream must not + wait on ``logical_topk_ready`` there: an event recorded inside a captured + segment is graph-local, and an eager wait on it raises + cudaErrorInvalidValue.""" + device = torch.device(DEVICE_TYPE) + num_tokens, num_topk, num_requests, blocks_per_req, block_size = 8, 128, 4, 4, 16 + + logical_topk_indices = torch.randint( + 0, + block_size * blocks_per_req, + (num_tokens, num_topk), + dtype=torch.int32, + device=device, + ) + builder = SparseMLAIndexGroupBuilder(logical_topk_indices) + group, layer_index = builder.register_layer(is_index_producing_layer=True) + assert layer_index == 0 and group.has_indexer + + req_id = torch.arange(num_tokens, dtype=torch.int32, device=device) % num_requests + block_table = torch.arange( + num_requests * blocks_per_req, dtype=torch.int32, device=device + ).view(num_requests, blocks_per_req) + attn_metadata = SimpleNamespace( + req_id_per_token=req_id, block_table=block_table, block_size=block_size + ) + expected = triton_convert_req_index_to_global_index( + req_id, + block_table, + logical_topk_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk, + ) + + vllm_config = create_vllm_config(model_name="Qwen/Qwen3.5-0.8B") + marker = torch.zeros(1, device=device) + capture_stream = torch.cuda.Stream(device=device) + # torch.cuda.stream routes through vLLM's patched torch.cuda.set_stream, + # so current_stream() tracks the capture stream as in production. + with ( + torch.cuda.stream(capture_stream), + set_forward_context( + None, vllm_config, cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE + ), + ): + capture = BreakableCUDAGraphCapture() + with capture: + marker.add_(1) + # Record inside a captured segment, convert in the eager + # break — mirrors mla.py -> unified_mla_attention_with_output. + group.set_logical_topk_ready(layer_index) + capture.add_eager( + lambda: group.convert_logical_to_physical_topk( + layer_index, + logical_topk_indices, + attn_metadata, + block_stride_rows=None, + return_valid_counts=False, + ) + ) + marker.add_(1) + capture.replay() + capture_stream.synchronize() + + result = group.physical_topk_indices[:num_tokens] + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("block_size", [16]) @pytest.mark.skipif( torch.cuda.get_device_capability() < (9, 0), diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index fd003fdce6ff..059b4005f0db 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -299,7 +299,14 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing(): running = request.num_computed_tokens assert not mamba.req_to_blocks["r"][running // block_size].is_null request.shared_prefix_boundary = running - assert mamba._cache_partial_tail_block(request, running) is None + assert ( + mamba._cache_partial_tail_block( + request, + running, + retention_interval=manager.coordinator.retention_interval, + ) + is None + ) @pytest.mark.parametrize("eagle_group", [None, 0]) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 8012c82f7841..393ec484651e 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -615,6 +615,77 @@ def test_internal_checkpoint_uses_partial_hash_lifecycle(): assert num_computed == 112 +@pytest.mark.parametrize( + "retention_interval,transient_published", + [(None, True), (0, False)], +) +def test_internal_checkpoint_publication_respects_retention( + retention_interval, transient_published +): + """With retention_interval=0, a mid-prompt internal checkpoint is + request-local and must not enter the prefix cache; the prompt-end + checkpoint keeps the existing publication behavior. Dense retention + (None) is unchanged. + """ + hash_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=32, + num_prefill_checkpoint_blocks=1, + ) + manager.coordinator.retention_interval = retention_interval + + request = make_request("producer", list(range(240)), hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(request) + + # Mid-prompt chunk: exports a transient checkpoint at 112. + assert manager.allocate_slots(request, 128, num_computed, computed_blocks) + transient_hash = request.block_hashes[112 // hash_block_size - 1] + transient_hit = manager.block_pool.get_cached_block(transient_hash, [1]) + assert (transient_hit is not None) == transient_published + + request.num_computed_tokens = 128 + manager.new_step_starts() + + # Final chunk: the prompt-end checkpoint at 224 stays published. + assert manager.allocate_slots(request, 112) is not None + end_hash = request.block_hashes[224 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(end_hash, [1]) is not None + + +def test_transient_checkpoint_evicts_retained_boundary_hash(): + """With retention_interval=0, the checkpoint slot can coincide with a + reachable boundary block that this step's full-block pass just hashed. + The checkpoint state overwrites the slot, so the stale boundary hash + must be evicted rather than left pointing at the wrong state. + """ + hash_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=32, + num_prefill_checkpoint_blocks=1, + ) + manager.coordinator.retention_interval = 0 + + request = make_request("producer", list(range(240)), hash_block_size, sha256) + # Block 2 ends at the shared-prefix junction 96, so the full-block pass + # hashes it as a retained boundary. + request.shared_prefix_boundary = 96 + computed_blocks, num_computed, _ = manager.get_computed_blocks(request) + assert manager.allocate_slots(request, 128, num_computed, computed_blocks) + + # The same slot is the transient checkpoint block (state@112): both the + # boundary hash and the checkpoint itself must stay unpublished. + boundary_hash = request.block_hashes[96 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(boundary_hash, [1]) is None + checkpoint_hash = request.block_hashes[112 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None + + def test_eagle_block_aligned_checkpoint_replaces_newer_hash(): hash_block_size = mamba_block_size = 32 manager = make_full_mamba_manager( diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index 7ff63d019c85..7e27a4517e03 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -334,7 +334,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): request = create_requests(num_requests=1, num_tokens=1)[0] request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens request.num_output_placeholders = 1 @@ -343,10 +342,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c6354cdc0fa0..86e66e1941f1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -279,6 +279,44 @@ def make_request( ) +@pytest.mark.parametrize("dcp", [1, 4]) +def test_effective_attention_block_size_matches_events(dcp): + from vllm.distributed.kv_events import BlockStored + from vllm.v1.engine.core import EngineCore + + config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], new_kv_cache_spec()), + ], + ) + manager = KVCacheManager( + generate_scheduler_kv_cache_config([config]), + max_model_len=256, + scheduler_block_size=16 * dcp, + hash_block_size=16 * dcp, + dcp_world_size=dcp, + enable_kv_cache_events=True, + ) + core = EngineCore.__new__(EngineCore) + core.vllm_config = SimpleNamespace(cache_config=CacheConfig(block_size=16)) + core.scheduler = SimpleNamespace(kv_cache_manager=manager) + core._initialize_effective_attention_block_size() + block_size = core.vllm_config.cache_config.effective_attention_block_size + assert block_size == 16 * dcp + + request = make_request( + "block-size", list(range(64)), block_size=16 * dcp, hash_fn=sha256 + ) + assert manager.allocate_slots(request, 64) is not None + assert [ + event.block_size + for event in manager.take_events() + if isinstance(event, BlockStored) + ] == [block_size] + + def new_kv_cache_spec( block_size=16, num_kv_heads=2, diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 3a972e682e4e..061c825fc895 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -51,6 +51,7 @@ ECConnectorOutput, KVConnectorOutput, ModelRunnerOutput, + RoutedExpertsLists, SamplingMaskLists, make_empty_encoder_model_runner_output, ) @@ -131,6 +132,44 @@ def test_add_requests(): assert len(scheduler.waiting) == i + 1 +def test_routed_experts_prompt_start_at_prompt_end(): + """A prompt-end offset should return empty routing data, not crash.""" + scheduler = create_scheduler() + (request,) = create_requests(num_requests=1) + request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens + scheduler.add_request(request) + scheduler_output = scheduler.schedule() + num_tokens = scheduler_output.num_scheduled_tokens[request.request_id] + + manager = Mock() + manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8) + manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8) + scheduler.enable_return_routed_experts = True + scheduler.routed_experts_mgr = manager + scheduler._re_block_ids = {request.request_id: []} + + outputs = scheduler.update_from_output( + scheduler_output, + ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[0]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + routed_experts=RoutedExpertsLists( + routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8), + slot_mapping=np.arange(num_tokens), + ), + ), + ) + + manager.get.assert_called_once_with( + [], request.num_prompt_tokens, token_start=request.num_prompt_tokens + ) + assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1) + + def test_finish_request(): scheduler = create_scheduler() requests = create_requests(num_requests=10) @@ -3635,7 +3674,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): ) request.structured_output_request = Mock() request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar) - request.structured_output_request.grammar.accept_tokens.return_value = False request.status = RequestStatus.RUNNING request.num_computed_tokens = request.num_tokens @@ -3643,10 +3681,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.connector = None scheduler.ec_connector = None scheduler.structured_output_manager = Mock() - scheduler.structured_output_manager.should_advance.return_value = True - scheduler.structured_output_manager.trim_reasoning_for_advance.side_effect = ( - lambda request, new_token_ids: new_token_ids - ) + scheduler.structured_output_manager.accept_tokens.return_value = False scheduler.requests = {request.request_id: request} scheduler.running = [request] scheduler.waiting = Mock() @@ -3695,8 +3730,8 @@ def free_request(req: Request, delay_free_blocks: bool = False): ) engine_core_outputs = scheduler.update_from_output(output, model_runner_output) - request.structured_output_request.grammar.accept_tokens.assert_called_once_with( - request.request_id, [123] + scheduler.structured_output_manager.accept_tokens.assert_called_once_with( + request, [123] ) assert request.resumable is False assert request.status == RequestStatus.FINISHED_ERROR @@ -6534,16 +6569,15 @@ def _decode_ready_request(scheduler): return request -def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): - """ngram_gpu pads unfilled draft slots with -1; update_draft_token_ids must +def test_update_draft_token_ids_strips_ngram_padding(): + """ngram_gpu pads unfilled draft slots with -1; manager validation must strip them before grammar.validate_tokens (which otherwise raises).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler.update_draft_token_ids( @@ -6554,16 +6588,15 @@ def test_update_draft_token_ids_strips_ngram_padding(monkeypatch): assert request.spec_token_ids == [10, 11] -def test_update_draft_token_ids_in_output_strips_padding(monkeypatch): +def test_update_draft_token_ids_in_output_strips_padding(): """Same guard on the output path; the -1 pad-back for the rejected count - is preserved (only the input to validate_tokens is stripped).""" + is preserved (only the input to manager validation is stripped).""" scheduler = create_scheduler(num_speculative_tokens=4) request = _decode_ready_request(scheduler) grammar = _RecordingGrammar() - request.structured_output_request = SimpleNamespace(grammar=grammar) - monkeypatch.setattr( - scheduler.structured_output_manager, "should_advance", lambda req: True + request.structured_output_request = SimpleNamespace( + grammar=grammar, reasoning_ended=True ) scheduler_output = SimpleNamespace( diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index aa1f730afbad..e388b73468d3 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -31,9 +31,8 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.worker import mamba_utils -from vllm.v1.worker.gpu_input_batch import CachedRequestState +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_model_runner import GPUModelRunner -from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch from vllm.v1.worker.mamba_utils import get_mamba_groups @@ -315,7 +314,7 @@ def check_copy_info( action: tuple[int, int], kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - input_batch: GPUInputBatch, + input_batch: InputBatch, ): assert copy_info is not None if action == (-1, -1): @@ -371,7 +370,7 @@ def fake_preprocess_mamba_fn( kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index a03718ab5a15..219850f5d489 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -14,6 +14,7 @@ import argparse import asyncio +import multiprocessing from http import HTTPStatus from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -34,6 +35,7 @@ from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import human_readable_int +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor @@ -66,6 +68,7 @@ def _make_async_llm( llm.output_processor = MagicMock() llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens + llm.admission_stats = None return llm @@ -84,6 +87,42 @@ def _make_scheduler_config(**kwargs) -> SchedulerConfig: ) +def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]: + counters = multiprocessing.RawArray( + "q", SharedAdmissionStats.num_counters(client_count) + ) + return [ + SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index) + for index in range(client_count) + ] + + +@pytest.mark.skip_global_cleanup +def test_shared_admission_stats_aggregate_api_servers(): + server_0, server_1 = _make_shared_stats() + + server_0.set_num_requests(2) + assert server_0.get_num_requests() == 2 + server_1.set_num_requests(3) + assert server_0.get_num_requests() == server_1.get_num_requests() == 5 + + server_0.set_num_requests(1) + assert server_0.get_num_requests() == 4 + + +@pytest.mark.skip_global_cleanup +def test_shared_request_snapshot_enforces_global_limit(): + server_0, server_1 = _make_shared_stats() + server_0.set_num_requests(2) + llm = _make_async_llm(max_num_queued_reqs=2) + llm.admission_stats = server_1 + + with pytest.raises(QueueOverflowError): + llm.check_admission() + + assert server_0.get_num_requests() == 2 + + # --------------------------------------------------------------------------- # Exception classes # --------------------------------------------------------------------------- diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index f97ed76c5c48..3f407457b22a 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -158,7 +158,10 @@ def setsockopt(self, *_args, **_kwargs): local_engines_only=False, enable_elastic_ep=False, ) - vllm_config = SimpleNamespace(parallel_config=parallel_config) + vllm_config = SimpleNamespace( + parallel_config=parallel_config, + cache_config=SimpleNamespace(effective_attention_block_size=None), + ) client = core_client_mod.MPClient( asyncio_mode=False, @@ -368,10 +371,15 @@ def test_dplb_finished_requests_release_inflight(): assert req.request_id not in client.reqs_in_flight -def test_apply_ready_response_syncs_block_size(): +@pytest.mark.parametrize( + ("effective_size", "other_size"), + [(None, None), (4224, 4224), (4224, 1056), (4224, None)], +) +def test_apply_ready_response_syncs_block_size(effective_size, other_size): import msgspec client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), @@ -399,14 +407,31 @@ def test_apply_ready_response_syncs_block_size(): max_loras=0, ) ) - client._apply_ready_response(payload) + fields = msgspec.msgpack.decode(payload) + if effective_size is None: + del fields["effective_attention_block_size"] + else: + fields["effective_attention_block_size"] = effective_size + client._apply_ready_response(msgspec.msgpack.encode(fields)) assert client.vllm_config.cache_config.block_size == 1056 + cache_config = client.vllm_config.cache_config + assert cache_config.effective_attention_block_size == effective_size + + fields["effective_attention_block_size"] = other_size + client._apply_ready_response(msgspec.msgpack.encode(fields)) + assert cache_config.effective_attention_block_size == ( + effective_size if effective_size == other_size else None + ) + + client._apply_ready_response(b"") + assert cache_config.effective_attention_block_size is None def test_apply_ready_response_syncs_mamba_block_size(): import msgspec client = object.__new__(MPClient) + client._effective_attention_block_sizes = set() client.vllm_config = SimpleNamespace( cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0), model_config=SimpleNamespace(max_model_len=8192), diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 26664e3ac550..89b77ae43fa8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, call import pytest @@ -74,24 +75,35 @@ make_offload_key, ) from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager -from vllm.v1.outputs import KVConnectorOutput +from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus -@pytest.mark.parametrize("boundary", [3840, 3904]) +@pytest.mark.parametrize( + "window,extra_retained", [(128, 0), (128, 3), (256, 1), (256, 3)] +) +@pytest.mark.parametrize("boundary", [3840, 3904, 3968]) @pytest.mark.parametrize("eagle", [False, True]) @pytest.mark.parametrize("left_state", ["present", "missing", "pending", "loading"]) -def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): +def test_swa_offload_window_covers_unaligned_hit( + boundary, eagle, left_state, window, extra_retained +): """A smaller SWA group can move the hit inside another group's CPU chunk.""" groups = [] - for i, (block_size, window) in enumerate([(256, None), (64, 128), (8, 128)]): + for i, (block_size, group_window) in enumerate( + [(256, None), (64, window), (8, 128)] + ): kwargs = dict( block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 ) kv_spec = ( FullAttentionSpec(**kwargs) - if window is None - else SlidingWindowSpec(**kwargs, sliding_window=window) + if group_window is None + else SlidingWindowSpec( + **kwargs, + sliding_window=group_window, + extra_retained_tokens=extra_retained if i == 1 else 0, + ) ) groups.append( KVCacheGroupSpec([f"layer{i}"], kv_spec, is_eagle_group=eagle and i == 1) @@ -152,7 +164,9 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): sched._chunks_being_loaded.add(left_key) external, load_async = sched.get_num_new_matched_tokens(request, 0) - if boundary == 3904 and left_state != "present": + first_retained_block = (boundary - window + 1 - extra_retained) // 64 + needs_left = first_retained_block // 4 < end - 1 + if needs_left and left_state != "present": assert external == (0 if left_state == "missing" else None) assert not load_async return @@ -176,8 +190,8 @@ def test_swa_offload_window_covers_unaligned_hit(boundary, eagle, left_state): sched.update_state_after_alloc(request, blocks, external) assert len(sched._jobs) == 1 job = next(iter(sched._jobs.values())) - if boundary == 3904: - assert {swa[14], swa[15]} <= job.keys + if needs_left: + assert {left_key, swa[end - 1]} <= job.keys manager.complete_load(job.keys, state.req_context) @@ -195,15 +209,17 @@ def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler: def _make_partial_tail_request( scheduler: OffloadingConnectorScheduler, + kv_transfer_params: dict[str, Any] | None = None, ) -> MagicMock: request = MagicMock() request.request_id = "req" - request.kv_transfer_params = None + request.kv_transfer_params = kv_transfer_params request.num_prompt_tokens = 30 request.num_tokens = 30 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)] request.all_token_ids = list(range(30)) request.lora_request = None + request.skip_reading_prefix_cache = False request.is_finished.return_value = False scheduler.on_new_request(request) return request @@ -427,6 +443,53 @@ def test_lookup_cap_stops_at_authoritative_prefix_boundary(): assert scheduler._req_status["req"].partial_tail_boundary == 20 +def test_max_load_tokens_zero_skips_partial_tail_lookup(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 0} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 16) == (0, False) + scheduler.manager.lookup.assert_not_called() + assert all( + state.num_hit_chunks == 1 for state in scheduler._req_status["req"].group_states + ) + + +def test_max_load_tokens_caps_tokens_beyond_gpu_prefix(): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 8} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 16) == (8, True) + assert scheduler._req_status["req"].partial_tail_boundary == 24 + + +def test_max_load_tokens_rounds_down_without_partial_tail(): + scheduler = _make_partial_tail_scheduler() + scheduler.config = scheduler.config._replace(supports_partial_tail=False) + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": 20} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 0) == (16, True) + + +@pytest.mark.parametrize("value", ["8", 8.5, -1, True]) +def test_invalid_max_load_tokens_is_ignored(value): + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request( + scheduler, kv_transfer_params={"max_load_tokens": value} + ) + scheduler.manager.lookup.return_value = LookupResult.HIT + + assert scheduler.get_num_new_matched_tokens(request, 0) == (28, True) + + def test_recurrent_group_unhashed_block_does_not_truncate_load_boundary(): """A recurrent group may hold unhashed blocks inside the computed prefix. @@ -1708,6 +1771,15 @@ def test_retry_stops_at_miss(self): class TestSlidingWindowLookup: + def test_pending_chunk_behind_gap_does_not_defer(self): + """A pending suffix too short for the window cannot improve the hit.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT, 4: LookupResult.HIT_PENDING} + ) + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 2 + ) + def test_all_hit_exact_window(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 @@ -1826,6 +1898,73 @@ def test_hit_pending_does_not_break_streak(self): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_swa_load_does_not_wait_for_unusable_pending_suffix( + request_runner, async_scheduling +): + """A CPU store beyond a missing SWA chunk must not block a ready prefix.""" + groups = [ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=4, num_kv_heads=1, head_size=1, dtype=torch.float32 + ), + ), + KVCacheGroupSpec( + ["swa"], + SlidingWindowSpec( + block_size=4, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=8, + ), + ), + ] + runner = request_runner( + block_size=4, + num_gpu_blocks=32, + async_scheduling=async_scheduling, + kv_cache_groups=groups, + ) + sched = runner.connector_scheduler + manager = CPUOffloadingManager(num_chunks=16) + sched.manager = manager + runner.new_request(token_ids=list(range(17))) + state = sched._req_status["0"] + state.update_offload_keys() + full, swa = (g.offload_keys for g in state.group_states) + ready = manager.prepare_store(full[:4] + swa[:2], state.req_context) + assert ready is not None + manager.complete_store(ready.keys_to_store, state.req_context) + pending = manager.prepare_store(swa[3:4], state.req_context) + assert pending is not None + + # The last SWA chunk is still being written; the preceding one is missing. + output = runner.scheduler.schedule() + metadata = output.kv_connector_metadata + assert isinstance(metadata, OffloadingConnectorMetadata) + [(job_id, load)] = metadata.load_jobs.items() + assert isinstance(load.dst_spec, GPULoadStoreSpec) + assert load.dst_spec.group_sizes == [2, 2] + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + runner.scheduler.update_from_output( + output, + ModelRunnerOutput.with_kv_conn_output_only( + KVConnectorOutput( + finished_recving={"0"}, + kv_connector_worker_meta=OffloadingWorkerMetadata( + completed_jobs={job_id: 1} + ), + ) + ), + ) + resumed = runner.scheduler.schedule() + assert resumed.num_scheduled_tokens == {"0": 9} + assert manager.lookup(swa[3], state.req_context) is LookupResult.HIT_PENDING + + @pytest.mark.parametrize("alignment_chunk_count", [4, 8, 64]) @pytest.mark.parametrize("sliding_window_chunks", [1, 2, 3]) @pytest.mark.parametrize("is_eagle_group", [False, True]) @@ -2503,27 +2642,14 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner assert req_status.transfer_jobs == pending_store_jobs +@pytest.mark.parametrize("extra_retained", [0, 1, 3]) @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_swa_alignment_skip(request_runner, async_scheduling: bool): - """SWA blocks unreachable by the load path are skipped during store. +def test_swa_alignment_skip(request_runner, async_scheduling: bool, extra_retained): + """Store the window a cold consumer must restore at each FA boundary. - Simulates a DeepSeek V4-like hybrid architecture where SWA groups have - much smaller block sizes than the full-attention (MLA) group, causing - most SWA blocks to be unreachable by the alignment-based load path. - - Setup: - - Group 0: full attention (MLA-like), block_size=16 - - Group 1: SWA, block_size=4, sliding_window=8 - - alignment_tokens = 16, so 16 / 4 = 4 SWA blocks per alignment segment. - sliding_window_size_in_chunks = ceil(8 / 4) = 2. - Within each segment of 4 SWA blocks, only the trailing 2 are stored. - - With 32 tokens (2 full-attn blocks, 8 SWA blocks): - - Group 0 stores: blocks 0, 1 (all full-attn blocks) - - Group 1 stores: blocks 2, 3, 6, 7 (skip 0,1,4,5) - - For real DeepSeek V4 (100K tokens), this reduces SWA stores by ~78%. + With FA blocks of 16 tokens and SWA blocks of 4 tokens, a window of 8 + needs the last two blocks per segment. Retaining three extra tokens for + MTP requires one more block; retaining just one token does not. """ full_attn_block_size = 16 swa_block_size = 4 @@ -2548,6 +2674,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): head_size=1, dtype=torch.float32, sliding_window=sliding_window, + extra_retained_tokens=extra_retained, ), ), ] @@ -2565,16 +2692,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): # Group 0: full attention assert kv_group_configs[0].sliding_window_size_in_chunks is None assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size - # Group 1: SWA with sliding_window_size of 2 chunks - assert kv_group_configs[1].sliding_window_size_in_chunks == 2 + # Group 1: SWA, including any additional retained blocks + assert kv_group_configs[1].sliding_window_size_in_chunks == ( + 3 if extra_retained > 1 else 2 + ) assert kv_group_configs[1].tokens_per_chunk == swa_block_size # alignment_tokens = full_attn_block_size = 16 assert runner.connector_scheduler.config.alignment_tokens == full_attn_block_size - # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks - # (block_size=4). Decode 1 token to kick off processing (stores are - # deferred to next step). - num_tokens = 32 + # Keep a partial last block so decoding does not reuse the producer's + # freed blocks before the harness records its deferred stores. + num_tokens = 33 runner.new_request(token_ids=[0] * num_tokens) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) @@ -2587,12 +2715,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): ) runner.run( decoded_tokens=[EOS_TOKEN_ID], - # Group 0 (full attn, block_size=16): 2 offloaded chunks - # -> GPU blocks (0, 0) and (0, 1) - # Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2 - # per segment of 4: - # Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3) - # Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7) + # Retain each segment's tail, including MTP history when needed. expected_stored=( (0, 0), (0, 1), @@ -2600,27 +2723,26 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): (1, 3), (1, 6), (1, 7), - ), + ) + + (((1, 1), (1, 5)) if extra_retained > 1 else ()), ) # Verify that loads still work correctly for the stored SWA blocks. runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * num_tokens + [1]) - runner.manager.lookup.return_value = LookupResult.HIT - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2 + runner.manager.lookup.side_effect = lambda key, ctx: ( + LookupResult.HIT if (key, 0) in runner.offloaded else LookupResult.MISS + ) runner.run( decoded_tokens=[EOS_TOKEN_ID], - # Group 0: full prefix lookup hits 2 offloaded chunks - # -> loads GPU blocks (0, 0), (0, 1) - # Group 1: sliding window lookup finds trailing 2 from last segment - # (blocks 6, 7 which were stored) - # -> loads GPU blocks (1, 6), (1, 7) + # The consumer restores the full FA prefix and the final SWA tail. expected_loaded=( (0, 0), (0, 1), (1, 6), (1, 7), - ), + ) + + (((1, 5),) if extra_retained > 1 else ()), ) @@ -2888,6 +3010,47 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): runner.manager.lookup.assert_not_called() +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_max_load_tokens_limits_external_load(request_runner, async_scheduling: bool): + """The load cap limits external reads without changing the store path.""" + block_size = 4 + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + blocks_per_chunk=blocks_per_chunk, + ) + + runner.new_request(token_ids=[0] * tokens_per_chunk) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0, 1, 2), + ) + + runner.scheduler.reset_prefix_cache() + runner.manager.lookup.reset_mock() + runner.new_request( + token_ids=[0] * tokens_per_chunk, + kv_transfer_params={"max_load_tokens": 0}, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_loaded=(), + expected_stored=(0, 1, 2), + ) + + runner.manager.lookup.assert_not_called() + + # --------------------------------------------------------------------------- # Eagle/MTP test class # --------------------------------------------------------------------------- diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 9b54f626f265..39fb0bb54fae 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -4,6 +4,7 @@ import contextlib import inspect import os +import queue import tempfile import textwrap import time @@ -341,6 +342,10 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv(): assert req_meta.remote.request_id == f"prefill-{42}" # do_remote_prefill is consumed by request_finished to prevent re-issuing. assert request.kv_transfer_params["do_remote_prefill"] is False + # The scheduler is not waiting on this recv -- the request is already gone + # from self.requests, so reporting it would trip `assert req_id in + # self.requests` in _update_from_kv_xfer_finished. + assert req_meta.awaiting_kvs is False @patch( @@ -1948,17 +1953,43 @@ def test_mixed_memory_local_descriptors_split_by_memory_type(): assert memory_types == ["DRAM", "VRAM"] -def test_mixed_memory_read_notifies_after_both_transfers_finish(): +@pytest.fixture +def recv_worker(): + """Receive lifecycle state without distributed or device initialization.""" worker = object.__new__(NixlConnectorWorker) - worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} - worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} - worker._dram_src_handles_by_block_size = {16: 10} - worker._recving_metadata = {"request": MagicMock()} + worker.transfer_topo = MagicMock() + worker.transfer_topo.block_size_ratio.return_value = 1 + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace( + remote_block_size=16, remote_physical_blocks_per_logical=1 + ) + worker._physical_blocks_per_logical_kv_block = 1 + worker._recving_metadata = {"request": MagicMock(local_block_ids=([1, 2, 3],))} worker._recving_transfers = defaultdict(list) + worker._failed_recv_reqs = queue.Queue() + worker._recv_failures = set() + worker._replicated_pcp_done_sending = set() + worker._invalid_block_ids = queue.Queue() worker._pending_recv_notifs = {} - worker._failed_inflight_recvs = set() - worker.xfer_stats = MagicMock() + worker._reqs_to_send = {} + worker._replicated_pcp_done_sending = set() + worker._is_hma_required = False + worker._has_mamba = False + worker.use_host_buffer = False + worker.enable_permute_local_kv = False + worker.enable_heterogeneous_attn_post_process = False + worker.tp_rank = 0 + worker._log_failure = MagicMock() # type: ignore[method-assign] + worker.xfer_stats = NixlKVConnectorStats() worker.nixl_wrapper = MagicMock() + worker.nixl_wrapper.get_xfer_telemetry.return_value = get_default_xfer_telemetry() + return worker + + +def test_mixed_memory_read_notifies_after_both_transfers_finish(recv_worker): + worker = recv_worker + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])} + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])} + worker._dram_src_handles_by_block_size = {16: 10} worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102] worker.nixl_wrapper.check_xfer_state.return_value = "DONE" @@ -1985,27 +2016,20 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish(): np.testing.assert_array_equal(device_read.args[4], [7]) worker.nixl_wrapper.send_notif.assert_not_called() - assert worker._pop_done_transfers(worker._recving_transfers) == {"request"} + assert worker.get_finished() == (set(), {"request"}) worker.nixl_wrapper.send_notif.assert_called_once_with( "prefill", notif_msg=b"request:1" ) -def test_mixed_memory_read_failure_does_not_notify_producer(): +def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker): """A failed half of a split READ must suppress its success notification.""" - worker = object.__new__(NixlConnectorWorker) - worker._recving_metadata = {"request": MagicMock()} + worker = recv_worker worker._recving_transfers = {"request": [101, 102]} worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} - worker._failed_inflight_recvs = set() - worker._is_hma_required = True - worker._failed_recv_reqs = MagicMock() - worker._log_failure = MagicMock() # type: ignore[method-assign] - worker.xfer_stats = MagicMock() - worker.nixl_wrapper = MagicMock() worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"] - worker._pop_done_transfers(worker._recving_transfers) + assert worker.get_finished() == (set(), {"request"}) assert "request" not in worker._pending_recv_notifs worker.nixl_wrapper.send_notif.assert_not_called() @@ -2695,6 +2719,74 @@ def check_xfer_state(self, handle: int) -> str: return super().check_xfer_state(handle) +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FailingNixlWrapper, +) +@pytest.mark.parametrize("awaiting_kvs", [True, False]) +@pytest.mark.parametrize("notification_fails", [False, True]) +def test_empty_recv_is_reported_only_when_awaited( + default_vllm_config, dist_init, awaiting_kvs, notification_fails +): + """An empty recv is reported iff the scheduler parked the request on it. + + _read_blocks returns early when the local block list is empty -- D already + holds the KV and only the producer needs telling. Two kinds of caller reach + that path, and they need opposite handling: + + Parked (awaiting_kvs=True, a full prefix cache hit). Without a report the + request is named in neither _recving_transfers nor _failed_recv_reqs and + never reaches finished_recving. The scheduler has no other way to release a + WAITING_FOR_REMOTE_KVS request, and there is no timeout, so it sits in + skipped_waiting holding its blocks for the life of the process. + + Notify-only (awaiting_kvs=False): request_finished seeding an empty recv to + free P's blocks for a request aborted before it was scheduled, or a readback + on a request that stays RUNNING. Reporting either one crashes the engine + core -- _update_from_kv_xfer_finished asserts the id is still in + self.requests and that the request is parked or finished. + + A failed notification changes neither: the KV is local either way, and the + producer frees its own blocks on a timeout. + """ + vllm_config = create_vllm_config() + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, + connector.engine_id, + hand_shake_latency=0.0, + kv_cache_config=connector._kv_cache_config, + ) + connector.connector_worker.nixl_wrapper.fail_send_notif = notification_fails + + request_id = f"test_empty_recv_awaited_{awaiting_kvs}" + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id=request_id, + local_block_ids=(), # empty: the whole prompt is already cached locally + kv_transfer_params={ + "remote_block_ids": [[20, 21, 22]], + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + "remote_request_id": f"prefill-{request_id}", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + }, + awaiting_kvs=awaiting_kvs, + ) + connector.bind_connector_metadata(metadata) + dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) + connector.start_load_kv(dummy_ctx) + connector.bind_connector_metadata(NixlConnectorMetadata()) + time.sleep(0.1) + connector.start_load_kv(dummy_ctx) + + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert (request_id in done_recving) is awaiting_kvs + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, @@ -2897,13 +2989,10 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): # Wait for handshake to fail time.sleep(0.3) - # Check that blocks were marked invalid - invalid_blocks = connector.get_block_ids_with_load_errors() - assert invalid_blocks == {1, 2, 3} - - # Check that request appears in get_finished + assert connector.get_block_ids_with_load_errors() == set() _, done_recving = connector.get_finished(finished_req_ids=set()) assert request_id in done_recving + assert connector.get_block_ids_with_load_errors() == {1, 2, 3} @patch( @@ -2954,6 +3043,7 @@ def test_transfer_setup_failure_returns_finished( time.sleep(0.1) connector.start_load_kv(dummy_ctx) + assert connector.get_block_ids_with_load_errors() == set() results = connector.get_transfer_results(finished_req_ids=set()) assert request_id in results.finished_recving assert results.failed_recving == {request_id} @@ -3176,74 +3266,49 @@ def test_failed_request_skips_kv_postprocessing( assert invalid_blocks == {1, 2, 3} -@patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", - FailingNixlWrapper, -) -def test_handles_failing_in_separate_polls_do_not_kill_the_engine( - default_vllm_config, dist_init -): - """Failures across polls report once, after the last writer is terminal.""" - vllm_config = create_vllm_config() - connector = NixlConnector( - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) - ) - connector.connector_worker = FakeNixlConnectorWorker( - vllm_config, connector.engine_id, hand_shake_latency=0 - ) - worker = connector.connector_worker +@pytest.mark.parametrize("sibling_state", ["ERR", "DONE"]) +def test_recv_failure_waits_for_sibling_transfer(recv_worker, sibling_state): + """Failure must not release blocks while a sibling can still write to them.""" + worker = recv_worker + worker._recving_transfers["request"] = [101, 102] + worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]} + worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state] - request_id = "test_two_handles_failing_in_separate_polls" - metadata = NixlConnectorMetadata() - metadata.add_new_req_to_recv( - request_id=request_id, - local_block_ids=([1, 2, 3],), - kv_transfer_params={ - "remote_block_ids": ([4, 5, 6],), - "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, - "remote_request_id": f"prefill-{request_id}", - "remote_host": "localhost", - "remote_port": 1234, - "remote_tp_size": 1, - }, - ) - connector.bind_connector_metadata(metadata) - dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}) - connector.start_load_kv(dummy_ctx) - connector.bind_connector_metadata(NixlConnectorMetadata()) - time.sleep(0.1) - connector.start_load_kv(dummy_ctx) + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() + assert "request" in worker._recving_metadata + worker.nixl_wrapper.release_xfer_handle.assert_called_once_with(101) - # Give the request a second handle, as a second remote rank would. - handles = worker._recving_transfers[request_id] - assert len(handles) == 1 - first, second = handles[0], object() - worker._recving_transfers[request_id] = [first, second] + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == {"request"} + assert worker.get_block_ids_with_load_errors() == {1, 2, 3} + worker.nixl_wrapper.send_notif.assert_not_called() - # Poll 1: the first handle has failed, the second is still running. - def first_failed(handle): - return "ERR" if handle is first else "PROC" + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() - with patch.object( - worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed - ): - _, done_recving = connector.get_finished(finished_req_ids=set()) - assert request_id not in done_recving - assert request_id in worker._recving_metadata - assert worker._recving_transfers[request_id] == [second] - assert connector.get_block_ids_with_load_errors() == set() +def test_recv_failure_waits_for_unpollable_handle_to_be_released(recv_worker): + """A polling exception and failed release do not prove that DMA has stopped.""" + worker = recv_worker + worker._recving_transfers["request"] = [101] + worker.nixl_wrapper.check_xfer_state.side_effect = [RuntimeError("poll"), "DONE"] + worker.nixl_wrapper.release_xfer_handle.side_effect = [ + RuntimeError("release"), + None, + ] - # Only the last terminal handle permits the scheduler to reuse the blocks. - with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"): - _, done_recving = connector.get_finished(finished_req_ids=set()) + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == set() + assert worker.get_block_ids_with_load_errors() == set() + assert "request" in worker._recving_metadata - assert done_recving == {request_id} - assert request_id not in worker._recving_transfers - assert request_id not in worker._recving_metadata - assert connector.get_block_ids_with_load_errors() == {1, 2, 3} - _, done_recving = connector.get_finished(finished_req_ids=set()) - assert done_recving == set() + results = worker.get_transfer_results() + assert results.finished_recving == results.failed_recving == {"request"} + assert worker.get_block_ids_with_load_errors() == {1, 2, 3} + assert worker.nixl_wrapper.release_xfer_handle.call_count == 2 def _set_test_speculative_config( diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 259e9536c1ca..ae2ddece7f77 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -149,8 +149,49 @@ def uncached(block_id): scheduler.update_state_after_alloc(request, blocks, num_external_tokens=2) - _, _, local_num_computed_blocks = scheduler._reqs_need_recv[request.request_id] + _, _, local_num_computed_blocks, awaiting_kvs = scheduler._reqs_need_recv[ + request.request_id + ] assert local_num_computed_blocks == (2, 1) + # There are external tokens to pull, so the scheduler parks the request and + # is waiting to see it in finished_recving. + assert awaiting_kvs is True + + +@pytest.mark.cpu_test +def test_full_local_hit_is_not_awaited_by_the_scheduler(): + """A recv with num_external_tokens == 0 must not be flagged as awaited. + + The request keeps running -- there is nothing to pull, only a notification + owed to the producer -- so the worker must not report it in + finished_recving. Doing so trips `assert RequestStatus.is_finished` in + _update_from_kv_xfer_finished, which only tolerates a finished recv for a + request that is parked or already done.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, + ) + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.core.kv_cache_utils import KVCacheBlock + + scheduler = object.__new__(NixlPullConnectorScheduler) + scheduler._reqs_in_batch = set() + scheduler._reqs_need_save = {} + scheduler._reqs_need_recv = {} + scheduler.use_host_buffer = False + scheduler.is_bidirectional_kv_xfer_enabled = False + scheduler._is_hma_required = False + scheduler.kv_cache_config = MagicMock( + select_transfer_block_ids=lambda block_ids: block_ids + ) + + blocks = KVCacheBlocks(blocks=([KVCacheBlock(block_id=0, _block_hash=object())],)) + request = create_request(do_remote_prefill=True) + + scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0) + + _, local_block_ids, _, awaiting_kvs = scheduler._reqs_need_recv[request.request_id] + assert not local_block_ids # full local hit: nothing to pull + assert awaiting_kvs is False @pytest.mark.cpu_test @@ -877,6 +918,7 @@ def _make_mock_worker_for_desc_ids( worker = MagicMock(spec=NixlConnectorWorker) worker.num_regions = num_regions worker._uses_region_group_mapping = False + worker._transfer_layer_group_ids = () worker._has_mamba = has_mamba worker._group_spec_types = group_spec_types worker.block_len_per_layer = block_len_per_layer or [100] diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index bfec51da6901..a4bfa9ab83ae 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -12,6 +12,7 @@ """ from collections import defaultdict +from threading import Event, Lock from unittest.mock import MagicMock, patch import numpy as np @@ -112,6 +113,7 @@ def test_local_descriptors_follow_each_region_pool_capacity(): worker.block_len_per_layer = [16, 16] worker.block_stride_per_layer = [16, 16] worker.region_num_blocks = [2, 3] + worker._transfer_layer_region_indices = () descriptors = worker._build_fa_local([100, 1000], block_size_ratio=1) @@ -119,7 +121,8 @@ def test_local_descriptors_follow_each_region_pool_capacity(): @pytest.mark.cpu_test -def test_overlaid_transfer_groups_share_region_geometry(): +@pytest.mark.parametrize("push_pp", [False, True]) +def test_overlaid_transfer_groups_share_region_geometry(push_pp): """Groups overlaid on one allocation share its transfer region.""" import msgspec @@ -127,6 +130,9 @@ def test_overlaid_transfer_groups_share_region_geometry(): from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlAgentMetadata, ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( + NixlPushConnectorWorker, + ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -153,7 +159,14 @@ def test_overlaid_transfer_groups_share_region_geometry(): } groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches] - worker = object.__new__(NixlConnectorWorker) + worker_cls = NixlPushConnectorWorker if push_pp else NixlConnectorWorker + worker = object.__new__(worker_cls) + if push_pp: + worker._push_writer_stop = Event() + worker._push_writer_wake = Event() + worker._push_writer_thread = MagicMock() + worker._sending_transfers_lock = Lock() + worker._sending_transfers = defaultdict(list) worker.tp_rank = 0 worker.world_size = 1 worker.transfer_tp_rank = 0 @@ -198,6 +211,9 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker._desc_is_dram_by_block_size = {} worker._desc_pos_by_block_size = {} worker._dram_src_handles_by_block_size = {} + worker._transfer_layer_names = () + worker._transfer_layer_region_indices = () + worker._transfer_layer_group_ids = () worker._region_is_mla = [] worker.block_len_per_layer = [] worker.block_stride_per_layer = [] @@ -205,7 +221,8 @@ def test_overlaid_transfer_groups_share_region_geometry(): worker.use_host_buffer = False worker.host_xfer_buffers = {} worker.device_kv_caches = {} - worker.pp_size = 1 + worker.pp_size = 2 if push_pp else 1 + worker._is_hma_required = True worker.dcp_size = 1 worker.pcp_size = 1 worker.kv_buffer_device = "cuda" @@ -240,7 +257,13 @@ def test_overlaid_transfer_groups_share_region_geometry(): expected_addrs = [ backing.data_ptr() + block * block_stride for block in range(num_blocks) ] - assert worker.src_blocks_data[:, 0].tolist() == expected_addrs + num_desc_regions = 2 if push_pp else 1 + assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions + assert worker.num_descs == num_blocks * num_desc_regions + assert ( + worker.dst_region_num_blocks[worker.engine_id] + == [num_blocks] * num_desc_regions + ) metadata = msgspec.msgpack.decode( worker.xfer_handshake_metadata.agent_metadata_bytes, @@ -248,6 +271,7 @@ def test_overlaid_transfer_groups_share_region_geometry(): ) assert metadata.region_group_ids == [-1] assert metadata.region_num_blocks == [num_blocks] + assert metadata.region_members == ([["layer.0", "layer.1"]] if push_pp else []) assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]] diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 2c7f0cb866bf..7ce84bf685b1 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -32,9 +32,13 @@ import msgspec import pytest +from vllm.distributed.kv_transfer.kv_connector.utils import TransferTopology from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorTransferResults, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import ( + NixlBaseConnectorWorker, +) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( PUSH_REG_NOTIF_PREFIX, NixlAgentMetadata, @@ -43,6 +47,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import ( NixlPushConnectorWorker, ) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( get_base_request_id, ) @@ -133,7 +138,7 @@ def test_p_side_mamba_truncates_before_cache_lookup(self): with patch.object( sched, - "_truncate_mamba_request_for_prefill", + "_truncate_request_for_prefill", side_effect=AssertionError("must not truncate after cache lookup"), ): assert sched.get_num_new_matched_tokens(request, 0) == (0, False) @@ -351,8 +356,10 @@ def fresh(cls) -> _StubWriterWorker: # Base worker fields touched by start_load_kv / _get_new_notifs. w._recving_metadata = {} w._pending_recv_notifs = {} - w._failed_inflight_recvs = set() + w._failed_recv_reqs = queue.Queue() + w._recv_failures = set() w._recving_transfers = defaultdict(list) + w._is_hma_required = False w._reqs_to_process = set() w._reqs_to_send = {} w.consumer_notification_counts_by_req = defaultdict(int) @@ -360,12 +367,16 @@ def fresh(cls) -> _StubWriterWorker: w.pcp_rank = 0 w.pcp_dcp_sharded = False w.world_size = 1 + w.pp_size = 1 w.engine_id = "test-decode-engine" w._remote_agents = {} w._handshake_lock = threading.RLock() w._physical_blocks_per_logical_kv_block = 1 w._uses_region_group_mapping = False w.region_group_ids = [0] + w._transfer_layer_names = () + w._transfer_layer_region_indices = () + w._transfer_layer_group_ids = () w.dst_region_num_blocks = {} w.dst_region_group_ids = {} w.dst_uses_region_group_mapping = {} @@ -455,12 +466,14 @@ def test_sends_to_all_agents_with_prefix_and_payload(self): list(group) for group in rd["local_block_ids"] ] - def test_fails_request_when_engine_not_handshaked(self): + def test_queues_failed_request_when_engine_not_handshaked(self): w, failures = self._reg_send_worker() w._do_send_reg_notif("req-1", _registration_data("req-1")) - assert failures == [("req-1", None)] + assert failures == [] + assert w._failed_recv_reqs.get_nowait() == "req-1" + assert w._failed_recv_reqs.empty() w.nixl_wrapper.send_notif.assert_not_called() def test_agent_read_is_serialized_with_handshake_lock(self): @@ -877,6 +890,81 @@ def test_get_transfer_results_evicts_completed_state(self): assert evicted == ["req-done"] assert w._push_writer_wake.is_set() + @staticmethod + def _pollable_worker() -> _StubWriterWorker: + """Stub worker with the extra base-worker state the real + ``get_transfer_results`` needs to poll ``_sending_transfers``.""" + w = _StubWriterWorker.fresh() + w.transfer_topo = MagicMock() + w.nixl_wrapper = MagicMock() + w.xfer_stats = MagicMock() + w._log_failure = MagicMock() # type: ignore[method-assign] + w._invalid_block_ids = queue.Queue() + w._pending_recv_notifs = {} + w._replicated_pcp_done_sending = set() + w.use_host_buffer = False + return w + + def _make_sending_req(self, w: _StubWriterWorker) -> str: + request_id = "req-send-1" + w._sending_transfers[request_id] = [101, 102] + w._reqs_to_send[request_id] = time.perf_counter() + 60 + w._reqs_to_process.add(request_id) + return request_id + + def test_failed_push_send_is_not_reported_done_sending(self): + """A failed P-side WRITE must not surface as finished sending. + + Reporting it would free P-side blocks as if the KV had been + delivered while D still waits for the data; the request must be + left to the lease / watchdog instead. + """ + w = self._pollable_worker() + request_id = self._make_sending_req(w) + w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"] + + results = w.get_transfer_results() + + assert request_id not in results.finished_sending + # Lease tracking is preserved so the watchdog can reschedule. + assert request_id in w._reqs_to_send + assert request_id in w._reqs_to_process + # Both handles were still cleaned up. + assert request_id not in w._sending_transfers + + @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)]) + def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma): + """P-side sends are completed by consumer notifs / lease expiry, + not by local WRITE-handle completion.""" + w = self._pollable_worker() + w.pp_size = pp_size + w._is_hma_required = is_hma + request_id = self._make_sending_req(w) + w.nixl_wrapper.check_xfer_state.side_effect = ["DONE", "PROC", "DONE"] + + results = w.get_transfer_results() + assert request_id not in results.finished_sending + assert w._sending_transfers[request_id] == [102] + assert request_id in w._reqs_to_send + + results = w.get_transfer_results() + + assert request_id not in results.finished_sending + assert request_id in w._reqs_to_send + assert request_id not in w._sending_transfers + + # The consumer notif is what completes the send. + w._pending_completion_notifs.put(f"{request_id}:1".encode()) + results = w.get_transfer_results() + assert request_id in results.finished_sending + assert request_id not in w._reqs_to_send + assert request_id not in w._reqs_to_process + assert w._evict_finished_inbox.get_nowait() == request_id + assert w.get_transfer_results().finished_sending == set() + assert w._evict_finished_inbox.empty() + assert w.nixl_wrapper.release_xfer_handle.call_count == 2 + w.xfer_stats.record_kv_expired_req.assert_not_called() + # ----------------------------------------------------------------- # # Negative / error-path tests # @@ -1079,7 +1167,7 @@ def test_empty_transfer_results_do_not_enqueue_eviction(self): # Wake set so the writer drains NIXL notifs even when idle. assert w._push_writer_wake.is_set() - def test_failed_push_transfer_finishes_sending(self): + def test_failed_push_transfer_does_not_evict_writer_state(self): w = _StubWriterWorker.fresh() w.nixl_wrapper = MagicMock() w.nixl_wrapper.check_xfer_state.return_value = "ERR" @@ -1096,11 +1184,12 @@ def test_failed_push_transfer_finishes_sending(self): ): results = w.get_transfer_results() - assert results.finished_sending == {"req-failed"} + assert results.finished_sending == set() assert results.finished_recving == set() assert "req-failed" not in w._sending_transfers - assert "req-failed" not in w._reqs_to_send - assert "req-failed" not in w._reqs_to_process + assert "req-failed" in w._reqs_to_send + assert "req-failed" in w._reqs_to_process + assert w._evict_finished_inbox.empty() w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17) def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog): @@ -1511,3 +1600,375 @@ def test_no_prefix_hit_leaves_blocks_untrimmed(self): local, remote = self._written_block_ids(w) assert local == [10, 11, 12] assert remote == [500, 501, 502] + + @pytest.mark.parametrize("layer_name_routing", [False, True]) + def test_different_region_groups_require_layer_routing(self, layer_name_routing): + """A layer mapped to a pooled remote region can still receive its blocks.""" + w, engine_id = self._worker_driving_xfer() + w.dst_region_group_ids[engine_id] = [-1] + if layer_name_routing: + w._transfer_layer_group_ids = (0,) + reg = _registration_data("req-groups", local_block_ids=([500, 501],)) + + if not layer_name_routing: + with pytest.raises(NotImplementedError, match="cache-group layouts"): + NixlPushConnectorWorker._do_start_push_kv( + w, "req-groups", ([10, 11],), reg + ) + w.nixl_wrapper.make_prepped_xfer.assert_not_called() + return + + NixlPushConnectorWorker._do_start_push_kv(w, "req-groups", ([10, 11],), reg) + + local, remote = self._written_block_ids(w) + assert local == [10, 11] + assert remote == [500, 501] + + +def _agent_metadata( + region_layers: list[list[str]], + base_addresses: list[int], + block_lens: list[int], + block_strides: list[int] | None = None, +) -> NixlAgentMetadata: + return NixlAgentMetadata( + engine_id="remote-engine", + agent_metadata=b"agent", + kv_caches_base_addr=base_addresses, + device_id=7, + num_blocks=2, + block_lens=block_lens, + block_strides=list(block_lens if block_strides is None else block_strides), + kv_cache_layout="LBHNC", + block_size=16, + ssm_sizes=(0, 0), + attn_backend_name="FLASH_ATTN", + physical_blocks_per_logical_kv_block=1, + region_members=region_layers, + ) + + +def _layer_routing_worker( + region_layers: list[list[str]], + group_by_layer: dict[str, int], + pp_size: int = 2, + is_hma: bool = True, +) -> _StubWriterWorker: + worker = _StubWriterWorker.fresh() + worker.pp_size = pp_size + worker.dcp_size = 1 + worker.block_size = 16 + worker._has_mamba = False + worker._is_hma_required = is_hma + worker.kv_cache_config = MagicMock(transfer_group_index_by_layer=group_by_layer) + worker.transfer_topo = MagicMock() + worker._set_region_layers(region_layers) + return worker + + +@pytest.mark.parametrize( + ("region_num_blocks", "expected"), + [(None, [1, 2, 15, 21, 22]), ([4, 7, 5], [1, 2, 9, 12, 13])], +) +def test_layer_group_ids_route_descriptor_blocks(region_num_blocks, expected): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + desc_ids = worker._compute_desc_ids( + block_ids=[[1, 2], [5]], + dst_num_blocks=10, + block_size_ratio=None, + physical_blocks_per_logical=1, + region_num_blocks=region_num_blocks, + ) + assert desc_ids.tolist() == expected + + +def test_layer_metadata_round_trip(): + metadata = _agent_metadata([["L0", "L1"]], [0x10000], [256]) + metadata.region_num_blocks = [4] + metadata.region_group_ids = [-1] + metadata.region_names = ["L0"] + metadata.region_mem_types = ["VRAM"] + + encoded = msgspec.msgpack.encode(metadata) + assert msgspec.msgpack.Decoder(NixlAgentMetadata).decode(encoded) == metadata + + +def test_layer_identity_gate_preserves_the_non_hma_path(): + assert _layer_routing_worker([["a"]], {"a": 0})._transfer_layer_region_indices == ( + 0, + ) + + # A base worker (pull) never routes by layer name. + pull = object.__new__(NixlBaseConnectorWorker) + pull._transfer_layer_names = () + pull._transfer_layer_region_indices = () + pull._transfer_layer_group_ids = () + pull._set_region_layers([["a"]]) + assert pull._transfer_layer_region_indices == () + + # A non-HMA local layout does not require layer-name routing. + assert ( + _layer_routing_worker( + [["a"]], {"a": 0}, is_hma=False + )._transfer_layer_region_indices + == () + ) + + # PP=1 has congruent local/remote regions and keeps the region route even + # when its allocator is hybrid. + assert ( + _layer_routing_worker( + [["a"]], {"a": 0}, pp_size=1 + )._transfer_layer_region_indices + == () + ) + + +def test_layer_alignment_fails_loud_when_remote_omits_layers(): + # Falling back to region-index routing here would silently transfer stale + # KV, so an unannounced peer must fail the handshake instead. + worker = _layer_routing_worker([["a"]], {"a": 0}) + with pytest.raises(AssertionError, match="no region_members"): + worker._align_remote_regions_by_layer(_agent_metadata([], [0xA000], [128])) + + +def test_layer_alignment_expands_pooled_regions(): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + assert worker._transfer_layer_names == ("a", "a.swa", "b") + assert worker._transfer_layer_region_indices == (0, 0, 1) + assert worker._transfer_layer_group_ids == (0, 1, 0) + + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128]) + metadata.region_num_blocks = [4, 6] + metadata.region_group_ids = [-1, 0] + metadata.region_names = ["a", "b"] + metadata.region_mem_types = ["VRAM", "VRAM"] + worker._align_remote_regions_by_layer(metadata) + + assert metadata.kv_caches_base_addr == [0xA000, 0xA000, 0xB000] + assert metadata.block_lens == [128, 128, 128] + assert metadata.block_strides == [128, 128, 128] + assert metadata.region_members == [["a"], ["a.swa"], ["b"]] + assert metadata.region_num_blocks == [4, 4, 6] + assert metadata.region_group_ids == [-1, -1, 0] + assert metadata.region_names == ["a", "a.swa", "b"] + assert metadata.region_mem_types == ["VRAM", "VRAM", "VRAM"] + + +def test_layer_alignment_filters_and_reorders_a_pp_stage(): + worker = _layer_routing_worker([["l2"], ["l3"]], {"l2": 0, "l3": 1}) + metadata = _agent_metadata( + [["l2"], ["l0"], ["l3"], ["l1"]], + [0xC000, 0xA000, 0xD000, 0xB000], + [65536, 65536, 32768, 32768], + [131072, 131072, 65536, 65536], + ) + metadata.region_num_blocks = [4, 6, 8, 10] + metadata.region_group_ids = [0, 0, 1, 1] + metadata.region_names = ["l2", "l0", "l3", "l1"] + metadata.region_mem_types = ["VRAM"] * 4 + + worker._align_remote_regions_by_layer(metadata) + + assert worker._transfer_layer_region_indices == (0, 1) + assert worker._transfer_layer_group_ids == (0, 1) + assert metadata.kv_caches_base_addr == [0xC000, 0xD000] + assert metadata.block_lens == [65536, 32768] + assert metadata.block_strides == [131072, 65536] + assert metadata.region_num_blocks == [4, 8] + assert metadata.region_group_ids == [0, 1] + assert metadata.region_names == ["l2", "l3"] + assert metadata.region_mem_types == ["VRAM", "VRAM"] + + +@pytest.mark.parametrize( + ("local_num_blocks", "remote_num_blocks"), + [([4, 4], [4, 4, 4]), ([4, 6], [3, 5, 7])], +) +def test_layer_descriptors_pair_layers_across_asymmetric_pp_split( + local_num_blocks, remote_num_blocks +): + """A PP split can leave each stage a different mix of attention types. + + Stage 1 owns L3 (full) pooled with L4 (sliding) in region 0, and L5 (full) + alone in region 1. The consumer holds every layer and pools them + differently, so region indices cannot be paired positionally. + """ + worker = _layer_routing_worker([["L3", "L4"], ["L5"]], {"L3": 0, "L4": 1, "L5": 0}) + worker.block_len_per_layer = [128, 128] + worker.block_stride_per_layer = [256, 512] + worker._region_is_mla = [False, False] + worker.num_blocks = 4 + worker.region_num_blocks = local_num_blocks + worker.device_id = 0 + + consumer = _agent_metadata( + [["L0", "L1"], ["L3", "L2"], ["L5", "L4"]], + [0xA000, 0xB000, 0xC000], + [128, 128, 128], + [256, 512, 1024], + ) + consumer.num_blocks = 4 + consumer.region_num_blocks = remote_num_blocks + worker._align_remote_regions_by_layer(consumer) + + # L3 -> remote region 1, L4 -> 2, L5 -> 2. Pairing by index would have sent + # L5 to region 1 and L4 to region 2's sibling. + assert consumer.kv_caches_base_addr == [0xB000, 0xC000, 0xC000] + + plan = TPMapping(((0,), (0,)), (0,), {0: 0}, 0) + local_descs = worker._build_fa_local([0x1000, 0x2000], block_size_ratio=1) + remote_descs = worker._build_fa_remote(plan, consumer, block_size_ratio=1) + local_counts = [local_num_blocks[i] for i in worker._transfer_layer_region_indices] + local_ids = worker._compute_desc_ids( + [[1, 2], [3]], 4, None, 1, region_num_blocks=local_counts + ) + remote_ids = worker._compute_desc_ids( + [[0, 1], [2]], 4, None, 1, region_num_blocks=consumer.region_num_blocks + ) + assert len(local_descs) == sum(local_counts) + assert len(remote_descs) == sum(consumer.region_num_blocks) + + # Each pair addresses the same layer's blocks despite different pooling + # and strides. L3 and L5 use group 0; L4 uses group 1. + assert local_descs[local_ids].tolist() == [ + [0x1100, 128, 0], + [0x1200, 128, 0], + [0x1300, 128, 0], + [0x2200, 128, 0], + [0x2400, 128, 0], + ] + assert remote_descs[remote_ids].tolist() == [ + [0xB000, 128, 7], + [0xB200, 128, 7], + [0xC800, 128, 7], + [0xC000, 128, 7], + [0xC400, 128, 7], + ] + + +def test_layer_alignment_is_canonical_across_remote_orderings(): + local, groups = [["x"], ["y"]], {"x": 0, "y": 1} + rank0 = _agent_metadata([["x"], ["y"]], [0x1000, 0x2000], [64, 128], [256, 512]) + rank1 = _agent_metadata([["y"], ["x"]], [0x2000, 0x1000], [128, 64], [512, 256]) + + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank0) + _layer_routing_worker(local, groups)._align_remote_regions_by_layer(rank1) + + assert rank0.kv_caches_base_addr == rank1.kv_caches_base_addr == [0x1000, 0x2000] + assert rank0.block_lens == rank1.block_lens == [64, 128] + assert rank0.block_strides == rank1.block_strides == [256, 512] + + +def test_layer_alignment_is_idempotent(): + worker = _layer_routing_worker( + [["a", "a.swa"], ["b"]], {"a": 0, "a.swa": 1, "b": 0} + ) + metadata = _agent_metadata([["a", "a.swa"], ["b"]], [0xA000, 0xB000], [128, 128]) + metadata.region_num_blocks = [4, 6] + metadata.region_group_ids = [-1, 0] + metadata.region_names = ["a", "b"] + metadata.region_mem_types = ["VRAM", "VRAM"] + + worker._align_remote_regions_by_layer(metadata) + aligned = msgspec.msgpack.encode(metadata) + worker._align_remote_regions_by_layer(metadata) + + assert msgspec.msgpack.encode(metadata) == aligned + + +def test_layer_alignment_rejects_missing_local_layer(): + worker = _layer_routing_worker([["l0"], ["l1"]], {"l0": 0, "l1": 1}) + metadata = _agent_metadata([["l0"]], [0xA000], [128]) + + with pytest.raises(AssertionError, match="missing locally owned layers"): + worker._align_remote_regions_by_layer(metadata) + + +def test_layer_alignment_rejects_duplicate_remote_layer(): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["a"]], [0xA000, 0xB000], [128, 128]) + + with pytest.raises(AssertionError, match="in multiple regions"): + worker._align_remote_regions_by_layer(metadata) + + +def test_layer_alignment_rejects_inconsistent_remote_metadata(): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["b"]], [0xA000], [128]) + + with pytest.raises(AssertionError, match="lengths disagree"): + worker._align_remote_regions_by_layer(metadata) + + +@pytest.mark.parametrize( + "field", + ["region_num_blocks", "region_group_ids", "region_names", "region_mem_types"], +) +def test_layer_alignment_rejects_inconsistent_region_geometry(field): + worker = _layer_routing_worker([["a"]], {"a": 0}) + metadata = _agent_metadata([["a"], ["b"]], [0xA000, 0xB000], [128, 128]) + setattr(metadata, field, []) + original = msgspec.msgpack.encode(metadata) + + with pytest.raises(AssertionError, match="lengths disagree"): + worker._align_remote_regions_by_layer(metadata) + assert msgspec.msgpack.encode(metadata) == original + + +def test_set_region_layers_rejects_duplicate_local_layer(): + with pytest.raises(AssertionError, match="spans multiple NIXL regions"): + _layer_routing_worker([["a"], ["a"]], {"a": 0}) + + +def test_set_region_layers_rejects_layer_outside_any_kv_group(): + with pytest.raises(AssertionError, match="outside any local group"): + _layer_routing_worker([["a"], ["b"]], {"a": 0}) + + +@pytest.mark.parametrize( + ("local_block_size", "remote_block_size", "remote_tp_size", "error"), + [ + (32, 16, 1, "identical P/D block sizes"), + (16, 32, 1, "identical P/D block sizes"), + (16, 16, 2, "decode TP greater"), + ], +) +def test_layer_handshake_rejects_unsupported_geometry( + local_block_size: int, remote_block_size: int, remote_tp_size: int, error: str +): + """Reject unsupported peers without registering agents or transfer state.""" + metadata = _agent_metadata([["a"]], [0xA000], [128]) + metadata.block_size = remote_block_size + worker = _layer_routing_worker([["a"]], {"a": 0}) + worker.block_size = local_block_size + worker.block_len_per_layer = [128] + worker.use_mla = False + worker.nixl_wrapper = MagicMock() + worker.kv_caches_base_addr = defaultdict(dict) + worker.dst_num_blocks = {} + worker.tp_mappings = {} + worker.transfer_topo = TransferTopology( + tp_rank=0, + tp_size=1, + block_size=local_block_size, + engine_id=worker.engine_id, + is_mla=False, + is_mamba=False, + total_num_kv_heads=8, + attn_backends=[], + ) + + with pytest.raises(NotImplementedError, match=error): + worker.add_remote_agent(metadata, remote_tp_size=remote_tp_size) + worker.nixl_wrapper.add_remote_agent.assert_not_called() + worker.nixl_wrapper.prep_xfer_dlist.assert_not_called() + assert not worker.tp_mappings + assert not worker.dst_num_blocks + assert not worker.kv_caches_base_addr + with pytest.raises(KeyError): + worker.transfer_topo.get_engine_info(metadata.engine_id) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 1ee1e817766a..98633c233bb2 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -521,6 +521,8 @@ def make_nixl_scheduler( Only sets the flags needed by the tests. When *heartbeat=True* the scheduler-side heartbeat bookkeeping fields are also initialised. """ + from types import SimpleNamespace + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) @@ -532,6 +534,7 @@ def make_nixl_scheduler( block_size=16, mamba_enabled=has_mamba, ) + sched.vllm_config = SimpleNamespace(num_prefill_lookahead_tokens=0) if heartbeat: sched._heartbeat_by_engine = {} @@ -598,9 +601,11 @@ def make_nixl_push_scheduler( mamba_enabled=has_mamba, ) - # vllm_config is consulted for parallel_config.tensor_parallel_size. + # vllm_config is consulted for parallel_config.tensor_parallel_size, and by + # `_prefill_backoff` on both the P and D prefill paths. vllm_config = MagicMock() vllm_config.parallel_config.tensor_parallel_size = 1 + vllm_config.num_prefill_lookahead_tokens = 0 sched.vllm_config = vllm_config # Push-specific state. diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index dc0e384d7c2f..aceec1cd8fa1 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager: mgr._sessions = {} mgr._kv_to_session = {} mgr._unbound_stores = {} + mgr._reaped_stores = {} + mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S mgr._failed_serve_ctxs = [] return mgr @@ -759,9 +761,9 @@ def test_unbound_store_kept_within_timeout(self): assert "req-fresh" in mgr._unbound_stores def test_unbound_store_reaped_after_timeout(self): - """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed - and their kv_request_id lands in _failed_req_ids so a late - FetchMsg/lookup doesn't try to satisfy them.""" + """Unbound stores past the reap deadline surface as failed and their + kv_request_id lands in _reaped_stores so a late FetchMsg is rejected + instead of parking demand nothing can satisfy.""" from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch mgr = self._make() @@ -779,7 +781,7 @@ def test_unbound_store_reaped_after_timeout(self): # 2 baseline + 2 buffered stores assert JobResult(job_id=10, success=False) in results assert JobResult(job_id=11, success=False) in results - assert "req-stale" in mgr._failed_req_ids + assert "req-stale" in mgr._reaped_stores def test_submit_store_parks_unbound_batch(self): """submit_store on an unbound id appends a batch with a fresh @@ -1204,6 +1206,75 @@ def test_both_loads_succeed(self): assert 201 in b_ok, f"B loads succeeded: {b_ok}" assert 200 in b_ok, f"B stores succeeded: {b_ok}" + def test_late_fetch_after_reap_fails_immediately(self): + """A fetch for a reaped kv_request_id fails in one round trip. + + The producer drops parked blocks once unbound_store_timeout_s + expires. A consumer whose fetch arrives after that must learn on the + next tick via TransferDoneMsg(success=False) — not stall for + _LOAD_TIMEOUT_S and then take the abort path. + """ + from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S + + mgr_a, mgr_b = _build_paired_managers() + kv_id = "req-late" + + # Record every message type leaving the consumer, so the assertion + # below can prove no abort was ever needed. + sent_from_a: list[str] = [] + drain = mgr_a._control._drain_outbound_to + + def recording_drain(peer_local_id: str): + out = drain(peer_local_id) + sent_from_a.extend(msg.get("type") for _, msg in out) + return out + + mgr_a._control._drain_outbound_to = recording_drain # type: ignore[method-assign] + + # Producer parks the blocks, then reaps them before any fetch lands. + mgr_b.submit_store( + _job_metadata( + job_id=200, + keys=[b"b-block"], + chunk_ids=[0], + kv_params={"remote_decoder": {"kv_request_id": kv_id}}, + ) + ) + mgr_b._unbound_store_timeout_s = 0.0 + reap_results = list(mgr_b.get_finished_jobs()) + assert JobResult(job_id=200, success=False) in reap_results + assert kv_id in mgr_b._reaped_stores + + # Consumer now asks for the blocks that no longer exist. + consumer_params = { + "remote_prefiller": { + "kv_request_id": kv_id, + "remote_host": "B", + "remote_port": 2, + }, + } + mgr_a.on_new_request(_req_context(consumer_params)) + mgr_a.submit_load( + _job_metadata( + job_id=101, + keys=[b"b-block"], + chunk_ids=[0], + kv_params=consumer_params, + ) + ) + + started = time.monotonic() + all_a: list[JobResult] = [] + for _ in range(8): + all_a.extend(list(mgr_a.get_finished_jobs())) + list(mgr_b.get_finished_jobs()) + elapsed = time.monotonic() - started + + assert JobResult(job_id=101, success=False) in all_a, all_a + assert kv_id in mgr_a._failed_req_ids + assert "abort_fetch" not in sent_from_a, sent_from_a + assert elapsed < _LOAD_TIMEOUT_S + # --------------------------------------------------------------------------- # _accept_new_peers — duplicate connection rejection diff --git a/tests/v1/kv_offload/tiering/test_factory.py b/tests/v1/kv_offload/tiering/test_factory.py index f5a576d25429..c391ac201631 100644 --- a/tests/v1/kv_offload/tiering/test_factory.py +++ b/tests/v1/kv_offload/tiering/test_factory.py @@ -10,6 +10,7 @@ 3. Error paths — missing tier_type, unknown tier_type, duplicate registration. """ +from importlib.util import find_spec from unittest.mock import MagicMock import pytest @@ -42,8 +43,11 @@ def _make_mock_args(): def test_pre_registered_tiers_can_be_imported(): - """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager.""" + """CI sentinel: registered tiers import and yield SecondaryTierManager.""" for tier_type in SecondaryTierFactory._registry: + # KVCR is an optional external dependency and may not be installed. + if tier_type == "kvcr" and find_spec("kvcr") is None: + continue cls = SecondaryTierFactory._registry[tier_type]() assert issubclass(cls, SecondaryTierManager) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py new file mode 100644 index 000000000000..38a2ae044ce2 --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Collection, Iterable, Mapping +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest + +pytest.importorskip("kvcr") + +from kvcr import ROUTER_HINT_KEY, KVCRBindings +from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig +from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpEntryResult, + OpEntryStatus, + OpHandle, + QueryStatus, +) + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadKey, + ReqContext, + make_offload_key, +) +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob +from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager +from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager + + +def _op_entries( + entries: Mapping[BlockKey, bool], +) -> dict[BlockKey, OpEntryResult]: + return { + key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED) + for key, success in entries.items() + } + + +class _StubControlChannel: + def __init__(self, endpoint: str) -> None: + self.endpoint = endpoint + + def send(self, endpoint: str, message: bytes) -> bool: + return True + + def recv(self) -> list[bytes]: + return [] + + def close(self) -> None: + pass + + +class RecordingKVCR: + def __init__(self) -> None: + self.config: KVCRConfig | None = None + self.guard_config: KVCRGuardConfig | None = None + self.backend_configs = KVCRBackendConfigs() + self.constructor_bindings: KVCRBindings | None = None + self.nixl_agent_name = "recording" + self.framework_control: _StubControlChannel | None = None + self.inventory_sink = None + self.query_status = QueryStatus.MISS + self.stats: OffloadingConnectorStats | None = None + self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.discard_hint_calls: list[str] = [] + self.deliver_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None] + ] = [] + self.deposit_calls: list[ + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]] + ] = [] + self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = [] + self._next_op_handle = 1 + + def submit_hint(self, *args: Any, **kwargs: Any) -> None: + self.submit_hint_calls.append((args, kwargs)) + + def discard_hint(self, request_id: str) -> None: + self.discard_hint_calls.append(request_id) + + def query( + self, + keys: Collection[BlockKey], + request_id: str | None = None, + ) -> list[tuple[QueryStatus, CacheTier | None]]: + tier = { + QueryStatus.HIT: CacheTier.LOCAL_G2, + QueryStatus.FETCHING: CacheTier.LOCAL_G2, + QueryStatus.FETCHABLE: CacheTier.REMOTE_G2, + QueryStatus.MISS: None, + }[self.query_status] + return [(self.query_status, tier) for _ in keys] + + def deliver( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + request_id: str | None = None, + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deliver_calls.append((op_handle, dict(blocks), request_id)) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def deposit( + self, + blocks: Mapping[BlockKey, list[MemDescriptor]], + ) -> OpHandle: + op_handle = self._next_op_handle + self._next_op_handle += 1 + self.deposit_calls.append((op_handle, dict(blocks))) + self.complete(op_handle, {key: True for key in blocks}) + return op_handle + + def complete( + self, + op_handle: OpHandle, + entries: Mapping[BlockKey, bool], + ) -> None: + self.completed.append((op_handle, _op_entries(entries))) + + def poll_completed( + self, + ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]: + completed = self.completed + self.completed = [] + return completed + + def get_stats(self) -> OffloadingConnectorStats | None: + stats = self.stats + self.stats = None + return stats + + def close(self) -> None: + pass + + +class _ExternalPolicy(FIFOPolicy): + pass + + +def _job( + job_id: int, + req_context: ReqContext, + key: OffloadKey | None = None, + chunk_id: int = 0, +) -> TransferJob: + return TransferJob( + job_id=job_id, + keys=[key if key is not None else OffloadKey(b"k0")], + chunk_ids=np.array([chunk_id], dtype=np.int64), + is_promotion=True, + req_context=req_context, + ) + + +def _make_tier( + monkeypatch, + kvcr: RecordingKVCR, + *, + enable_telemetry: bool = False, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + enable_kv_cache_events: bool = False, + self_describing_kv_events: bool = False, + policy: str | None = None, + g3: dict[str, object] | None = None, + control_ports: list[int] | None = None, + data_parallel_rank_local: int | None = None, +) -> KVCRSecondaryTierManager: + def make_control(_bind_host, bind_port, advertise_host): + return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}") + + def make_kvcr(config, bindings, backend_configs, guard_config): + kvcr.config = config + kvcr.guard_config = guard_config + kvcr.nixl_agent_name = config.nixl_agent_name + kvcr.backend_configs = backend_configs + kvcr.constructor_bindings = bindings + kvcr.framework_control = bindings.framework_control + kvcr.inventory_sink = bindings.inventory_sink + return kvcr + + monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr) + monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control) + return KVCRSecondaryTierManager( + offloading_spec=SimpleNamespace( + config=SimpleNamespace( + parallel=SimpleNamespace( + data_parallel_rank_local=data_parallel_rank_local, + ) + ), + kv_events_config=SimpleNamespace( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=self_describing_kv_events, + ), + ), + primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)), + tier_type="kvcr", + router_capabilities=["router_hint"], + control_host="127.0.0.1", + control_ports=control_ports if control_ports is not None else [7777], + control_advertise_host="127.0.0.1", + enable_telemetry=enable_telemetry, + secondary_g2_slots=secondary_g2_slots, + kvcr_service_socket_path=kvcr_service_socket_path, + compatibility_digest=compatibility_digest, + policy=policy, + g3=g3, + local_dram_backend="UCX", + remote_fw_dram_backend="UCX", + ) + + +def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch): + """Keep the control endpoint and guard pool aligned to the local DP rank.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + kvcr_service_socket_path="/tmp/kvcr.sock", + compatibility_digest="Opaque-Digest", + secondary_g2_slots=1, + control_ports=[7001, 7002], + data_parallel_rank_local=1, + ) + + assert kvcr.framework_control is not None + assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002" + assert kvcr.guard_config == KVCRGuardConfig( + kvcr_service_socket_path="/tmp/kvcr.sock", + guard_index=1, + compatibility_digest="Opaque-Digest", + ) + assert kvcr.config is not None + assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)] + assert kvcr.backend_configs.local_dram is None + + +def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path): + """Convert user-provided G3 paths to KVCR's typed configuration.""" + kvcr = RecordingKVCR() + path = tmp_path / "g3.data" + + _make_tier( + monkeypatch, + kvcr, + g3={"paths": [str(path)], "capacity_bytes_per_file": 64}, + ) + + assert kvcr.backend_configs.g3 == G3Options( + paths=(path,), + capacity_bytes_per_file=64, + ) + + +def test_kvcr_tier_adapts_request_and_load(monkeypatch): + """Check hint forwarding, key conversion, load descriptors, and cleanup.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + router_hint = {"opaque": True} + ctx = ReqContext( + req_id="req", + kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()}, + ) + key = make_offload_key((123).to_bytes(8, "big"), 0) + same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7) + other_key = make_offload_key((124).to_bytes(8, "big"), 0) + + tier.on_new_request(ctx) + + assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})] + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.key_adapter is not None + monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0") + decode = bindings.key_adapter.decode + hash_123 = (123).to_bytes(8, "big") + assert decode(BlockKey(bytes(key))) == hash_123 + assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123 + assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big") + + tier.submit_load(_job(7, ctx, key=key, chunk_id=2)) + + # Check load destinations and job completion independently of hint parsing. + assert len(kvcr.submit_hint_calls) == 1 + _, blocks, request_id = kvcr.deliver_calls[0] + assert request_id == "req" + assert list(blocks) == [key] + (descriptor,) = blocks[key] + assert descriptor.end_point_name == kvcr.nixl_agent_name + assert descriptor.addr == tier._primary_base_addr + 2 * 16 + assert descriptor.size == 16 + assert list(tier.get_finished_jobs()) == [JobResult(7, True)] + + # Here we verify that request cleanup discards the request-scoped hint in + # KVCR. + tier.on_request_finished(ctx) + assert kvcr.discard_hint_calls == ["req"] + + +def test_kvcr_tier_allows_request_without_router_hint(monkeypatch): + """Keep router hints optional for requests from non-hint-aware routers.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + + tier.on_new_request(ReqContext(req_id="req")) + + assert kvcr.submit_hint_calls == [] + + +@pytest.mark.parametrize( + ("status", "expected"), + [ + (QueryStatus.MISS, LookupResult.MISS), + (QueryStatus.HIT, LookupResult.HIT), + (QueryStatus.FETCHABLE, LookupResult.HIT), + (QueryStatus.FETCHING, LookupResult.RETRY), + ], +) +def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): + """Map KVCR cache states to the scheduler's lookup contract.""" + kvcr = RecordingKVCR() + kvcr.query_status = status + tier = _make_tier(monkeypatch, kvcr) + + assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected + + +def test_kvcr_tier_serves_primary_pin_request(monkeypatch): + """Hold primary-tier hits until KVCR releases the corresponding pin.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) + hit_keys = (keys[0], keys[2]) + chunk_ids = {keys[0]: 1, keys[2]: 5} + lifecycle: list[str] = [] + + class Parent: + def on_new_request(self, req_context): + lifecycle.append("new") + return SimpleNamespace() + + def lookup(self, key, req_context): + return LookupResult.HIT if key in hit_keys else LookupResult.MISS + + def create_store_job(self, requested_keys, req_context): + return TransferJob( + job_id=11, + keys=requested_keys, + chunk_ids=np.array([chunk_ids[key] for key in requested_keys]), + is_promotion=True, + req_context=req_context, + ) + + def on_request_finished(self, req_context): + lifecycle.append("finished") + + bindings = kvcr.constructor_bindings + assert bindings is not None + request = bindings.request_pin(keys) + + tier.serve_external_requests(Parent()) + + [(queued_request, result)] = bindings.poll_pin_results() + assert queued_request == request + assert result is not None + pin_handle, descriptors = result + assert descriptors[keys[1]] is None + block_descriptors = descriptors[keys[2]] + assert block_descriptors is not None + (descriptor,) = block_descriptors + assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride) + assert lifecycle == ["new", "finished"] + + polls = 0 + + def release_pin(): + nonlocal polls + polls += 1 + bindings.release_pin(pin_handle) + return [] + + monkeypatch.setattr(kvcr, "poll_completed", release_pin) + tier.drain_jobs() + + assert polls == 1 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + +def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch): + """Keep telemetry opt-in and namespace metrics only at the vLLM boundary.""" + assert KVCRSecondaryTierManager.build_metric_definitions({}) == {} + definitions = KVCRSecondaryTierManager.build_metric_definitions( + {"enable_telemetry": True} + ) + assert "vllm:kvcr_duration_seconds" in definitions + assert "vllm:kvcr_transfer_blocks" in definitions + + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True) + bindings = kvcr.constructor_bindings + assert bindings is not None + assert bindings.stats_factory is not None + stats = bindings.stats_factory() + stats.increase_counter( + "kvcr_transfer_blocks", + 2, + ("remote_deliver",), + ) + kvcr.stats = stats + + returned = tier.get_stats() + + assert returned is stats + assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2} + + +@pytest.mark.parametrize( + ("policy", "expected_type"), + [ + ("fifo", FIFOPolicy), + ("lru", LRUPolicy), + ("g3_fifo", G3FIFOPolicy), + ("g3_lru", G3LRUPolicy), + (f"{__name__}._ExternalPolicy", _ExternalPolicy), + ], +) +def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type): + """Resolve every built-in and fully qualified external policy.""" + kvcr = RecordingKVCR() + _make_tier(monkeypatch, kvcr, policy=policy) + + bindings = kvcr.constructor_bindings + assert bindings is not None + assert type(bindings.policy) is expected_type + + +@pytest.mark.parametrize( + ("socket_path", "digest"), + [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")], +) +def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest): + """Reject partial guard configuration before connecting to the service.""" + with pytest.raises(ValueError, match="configured together"): + _make_tier( + monkeypatch, + RecordingKVCR(), + kvcr_service_socket_path=socket_path, + compatibility_digest=digest, + ) + + +def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): + """Cover store descriptors and inventory translation at the KVCR boundary.""" + kvcr = RecordingKVCR() + tier = _make_tier( + monkeypatch, + kvcr, + secondary_g2_slots=2, + enable_kv_cache_events=True, + self_describing_kv_events=True, + ) + local_dram = kvcr.backend_configs.local_dram + assert local_dram is not None + assert [(name, size) for name, _, size in local_dram.pools] == [ + ("", 2 * tier._primary_row_stride) + ] + + key = OffloadKey(b"k0") + tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2)) + + _, blocks = kvcr.deposit_calls[0] + assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)] + + assert kvcr.inventory_sink is not None + kvcr.inventory_sink( + InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False) + ) + kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False)) + events = list(tier.take_events()) + assert [ + (event.keys, event.medium, event.ownership, event.removed) for event in events + ] == [ + ([key], Medium.CPU, "kvcr", False), + ([key], Medium.STORAGE, "kvcr", False), + ] + assert all(event.removal_expected for event in events) + tier.shutdown() + + +def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): + """Require tier-aware events whenever KVCR owns local cache inventory.""" + with pytest.raises(ValueError, match="self_describing_kv_events"): + _make_tier( + monkeypatch, + RecordingKVCR(), + secondary_g2_slots=1, + enable_kv_cache_events=True, + ) + + +def test_kvcr_tier_accumulates_block_results(monkeypatch): + """Wait for every block result and preserve partial success while draining.""" + kvcr = RecordingKVCR() + tier = _make_tier(monkeypatch, kvcr) + keys = [OffloadKey(b"k0"), OffloadKey(b"k1")] + tier.submit_load( + TransferJob( + job_id=13, + keys=keys, + chunk_ids=np.array([0, 1], dtype=np.int64), + is_promotion=True, + req_context=ReqContext(req_id="req"), + ) + ) + + op_handle, _, _ = kvcr.deliver_calls[0] + kvcr.completed.clear() + kvcr.complete(op_handle, {keys[0]: True}) + assert list(tier.get_finished_jobs()) == [] + + kvcr.complete(op_handle, {keys[1]: False}) + tier.drain_jobs() + + [result] = tier.get_finished_jobs() + assert result.job_id == 13 + assert not result.success + assert result.successful_keys == {keys[0]} diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 330e10438f0b..4e12a6f49dd8 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""grammar_bitmask under spec-decode draft padding (#44006).""" +"""Backend and engine-reasoner checks that are not manager-flow tests.""" import pytest from transformers import AutoTokenizer @@ -49,221 +49,6 @@ def _make_manager_and_request(backend: str, prompt_str: str = '{"a": "b"}'): return tokenizer, manager, request, prompt -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_with_padded_invalid_drafts(backend): - """Bitmask handles -1 padded drafts and returns N+1 rows.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid_drafts = [tokenizer.encode(":")[0], tokenizer.encode(' "')[0]] - padded = valid_drafts + [-1, -1] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_when_grammar_terminates_mid_window(backend): - """Drafts following an EOS that terminates the grammar are a no-op.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - eos = tokenizer.eos_token_id - drafts = [eos] + [tokenizer.encode(" ")[0]] * (NUM_SPEC_TOKENS - 1) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == NUM_SPEC_TOKENS + 1 - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_idempotent_across_calls(backend): - """Repeated calls with the same input return the same bitmask.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - drafts = [tokenizer.encode(":")[0], -1, -1, -1] - - first = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - second = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert first is not None and second is not None - assert (first == second).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bonus_position_constrained_after_invalid_drafts(backend): - """Regression for #44006: bonus row stays constrained after -1 padding.""" - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str='{"a"' - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - valid = tokenizer.encode(":")[0] - drafts = [valid, -1, -1, -1] - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_constrained_when_reasoning_ends_midwindow(backend): - """Drafts after a mid-window reasoning-end marker stay constrained.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, marker, post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_post_reasoning_end_drafts_skip_grammar_advance(backend, caplog): - """Post-marker drafts predate the bitmask and may be grammar-invalid; - grammar_bitmask must skip the grammar advance instead of asserting. - """ - tokenizer, manager, request, prompt = _make_manager_and_request( - backend, prompt_str="{" - ) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - assert not grammar.is_terminated() - - marker = tokenizer.encode("\n")[0] - - class StubReasoner: - def __init__(self, *_, **__): - self.end_token_id = marker - - def is_reasoning_end(self, input_ids): - return marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - manager.reasoner_cls = StubReasoner - request.structured_output_request.reasoner = StubReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - # A token that the JSON grammar would reject as the first post-marker - # token; without the fix grammar.accept_tokens fires the assertion. - invalid_post = tokenizer.encode("z")[0] - drafts = [pre, marker, invalid_post] - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - # Post-marker position is still bitmask-constrained. - assert not (bitmask[2] == -1).all() - # Grammar must not have advanced through the unvalidated draft. - assert not grammar.is_terminated() - assert "Failed to advance FSM" not in caplog.text - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_validate_tokens_then_bitmask_round_trip(backend): - """validate_tokens -> pad with -1 -> grammar_bitmask must not assert.""" - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - raw_drafts = [tokenizer.encode(",")[0], 99999, 12345, 67890] - accepted = grammar.validate_tokens(raw_drafts) - assert len(accepted) <= len(raw_drafts) - - padded = accepted + [-1] * (len(raw_drafts) - len(accepted)) - assert len(padded) == len(raw_drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: padded}, - ) - assert bitmask is not None - assert bitmask.shape[0] == len(padded) + 1 - assert not grammar.is_terminated() - - def test_xgrammar_accept_tokens_stops_at_termination(capfd): """Tokens after a terminating EOS do not reach the matcher.""" tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") @@ -302,7 +87,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert grammar.validate_tokens([eos, trailing]) == [eos] assert "trying to accept new token" not in capfd.readouterr().err - # Check matcher state directly to verify validation rolled it back. assert not grammar.matcher.is_terminated() assert grammar.accept_tokens(request.request_id, [eos]) @@ -312,87 +96,6 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert "trying to accept new token" not in capfd.readouterr().err -class _MarkerReasoner: - """Stub reasoner whose reasoning-end marker is a single fixed token.""" - - def __init__(self, marker: int): - self.marker = marker - - def is_reasoning_end(self, input_ids): - return self.marker in list(input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return self.marker in list(delta_ids) - - -def _setup_boundary_request(backend: str): - """Request with a structural-tag key and reasoning not yet ended.""" - from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - marker = tokenizer.encode("\n")[0] - structured_req = request.structured_output_request - # The grammar itself is JSON (cheap to build); only the key kind matters - # for the should_advance structural-tag branch, so pre-seed the cached - # property. - structured_req.__dict__["structured_output_key"] = ( - StructuredOutputOptions.STRUCTURAL_TAG, - "", - ) - manager.reasoner_cls = _MarkerReasoner - structured_req.reasoner = _MarkerReasoner(marker) - structured_req.reasoning_ended = False - return tokenizer, manager, request, prompt, marker - - -def test_should_advance_records_reasoning_end_index(): - """Regression for #44006 on post-#42452 main: the boundary step must - record where reasoning ends so the scheduler can trim before advancing. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - request.append_output_token_ids([pre, marker, post]) - - assert manager.should_advance(request) - assert structured_req.reasoning_ended - # Marker sits at absolute index len(prompt) + 1. - assert structured_req.reasoning_end_token_index == len(prompt) + 1 - - -def test_trim_reasoning_for_advance(): - """trim drops the marker and everything before it; later steps and - requests without a recorded boundary pass through unchanged. - """ - tokenizer, manager, request, prompt, marker = _setup_boundary_request("xgrammar") - structured_req = request.structured_output_request - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode("{")[0] - - # No boundary recorded yet: pass-through. - assert manager.trim_reasoning_for_advance(request, [pre]) == [pre] - - # Boundary step: marker mid-step keeps only the suffix. - step_tokens = [pre, marker, post] - request.append_output_token_ids(step_tokens) - assert manager.should_advance(request) - assert manager.trim_reasoning_for_advance(request, step_tokens) == [post] - - # Boundary step variant: marker last (the #44006 crash shape - # [198, ]) trims to empty -> scheduler skips accept_tokens. - structured_req.reasoning_end_token_index = len(request.all_token_ids) - 1 - assert manager.trim_reasoning_for_advance(request, step_tokens) == [] - - # Later steps: tokens are past the boundary, returned unchanged. - structured_req.reasoning_end_token_index = len(prompt) + 1 - next_step = [post, post] - request.append_output_token_ids(next_step) - assert manager.trim_reasoning_for_advance(request, next_step) == next_step - - class _EngineReasonerStub(ParserEngineReasoningAdapter): """Adapter-typed reasoner with a fixed end-token set and no real engine.""" @@ -409,7 +112,7 @@ def find_reasoning_end_offset(self, token_ids): for offset, token in enumerate(token_ids): if token in self._end_token_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids): return any(token in self._end_token_ids for token in input_ids) @@ -449,61 +152,3 @@ def test_bitmask_engine_reasoner_ends_midwindow_with_padding(backend): assert not (bitmask[2] == -1).all() assert reasoner.windows == [[pre, marker, post]] assert not grammar.is_terminated() - - -@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) -def test_bitmask_legacy_order_sensitive_predicate_probed_per_token(backend): - """Legacy predicates are probed one draft at a time, never on the whole - window. A predicate that reports which marker is *newest* (KimiK3) would - otherwise answer False for a window that closes and then reopens - reasoning, leaving the post-marker drafts and the bonus row - unconstrained. - """ - tokenizer, manager, request, prompt = _make_manager_and_request(backend) - grammar = request.structured_output_request.grammar - - assert grammar.accept_tokens(request.request_id, prompt) - - close = tokenizer.encode("\n")[0] - reopen = tokenizer.encode("#")[0] - assert close != reopen - - class NewestMarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end(self, input_ids): - return self.is_reasoning_end_streaming(input_ids, input_ids) - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - for token in reversed(list(delta_ids)): - if token == close: - return True - if token == reopen: - return False - return False - - manager.reasoner_cls = NewestMarkerReasoner - request.structured_output_request.reasoner = NewestMarkerReasoner() - request.structured_output_request.reasoning_ended = False - - pre = tokenizer.encode(" ")[0] - post = tokenizer.encode(",")[0] - drafts = [pre, close, post, reopen] - # Sanity: the whole window says "reopened", so only per-token probing - # can find the close marker. - assert not NewestMarkerReasoner().is_reasoning_end_streaming([], drafts) - - bitmask = manager.grammar_bitmask( - requests={request.request_id: request}, - structured_output_request_ids=[request.request_id], - scheduled_spec_decode_tokens={request.request_id: drafts}, - ) - - assert bitmask is not None - assert bitmask.shape[0] == len(drafts) + 1 - assert (bitmask[0] == -1).all() - assert (bitmask[1] == -1).all() - assert not (bitmask[2] == -1).all() - assert not (bitmask[-1] == -1).all() - assert not grammar.is_terminated() diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py deleted file mode 100644 index a57e5da72b26..000000000000 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ /dev/null @@ -1,557 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Unit tests for reasoning-aware structured output functionality (PR #25515).""" - -from collections.abc import Sequence -from unittest.mock import Mock - -import pytest - -from vllm.config import ModelConfig, SchedulerConfig, VllmConfig -from vllm.v1.request import Request -from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.structured_output.backend_types import StructuredOutputOptions - - -class UntouchableTokenIds(Sequence[int]): - """A token history that fails the test if anything reads it.""" - - def __init__(self, length: int) -> None: - self._length = length - - def __len__(self) -> int: - return self._length - - def __getitem__(self, item): - raise AssertionError("token history must not be read") - - def __iter__(self): - raise AssertionError("token history must not be iterated") - - -def make_qwen3_reasoner(): - """A real engine-based reasoning adapter over a tiny mock vocabulary.""" - from tests.parser.engine.conftest import make_mock_tokenizer - from vllm.parser.engine.registered_adapters import Qwen3ParserReasoningAdapter - - tokenizer = make_mock_tokenizer( - {"": 200, "": 201, "": 202, "": 203} - ) - return Qwen3ParserReasoningAdapter(tokenizer) - - -class MockReasoner: - def __init__(self, tokenizer): - self.is_reasoning_end = Mock(return_value=False) - self.is_reasoning_end_streaming = Mock(return_value=False) - - -class TestReasoningStructuredOutput: - """Test reasoning-aware structured output functionality.""" - - @pytest.fixture - def mock_model_config(self): - """Create a mock ModelConfig.""" - config = Mock(spec=ModelConfig) - config.skip_tokenizer_init = True # Skip tokenizer init to avoid network calls - config.get_vocab_size = Mock(return_value=50000) - # Add missing runner_type attribute that tokenizer initialization expects - config.runner_type = "generate" - # Add other attributes that tokenizer initialization might need - config.tokenizer = "test-tokenizer" - config.tokenizer_mode = "auto" - config.trust_remote_code = False - config.tokenizer_revision = None - return config - - @pytest.fixture - def mock_scheduler_config(self): - """Create a mock SchedulerConfig.""" - config = Mock(spec=SchedulerConfig) - config.max_num_seqs = 128 - return config - - @pytest.fixture - def mock_vllm_config(self, mock_model_config, mock_scheduler_config): - """Create a mock VllmConfig.""" - config = Mock(spec=VllmConfig) - config.model_config = mock_model_config - config.scheduler_config = mock_scheduler_config - config.structured_outputs_config = Mock() - config.structured_outputs_config.reasoning_parser = None - config.structured_outputs_config.enable_in_reasoning = False - config.speculative_config = None - return config - - @pytest.fixture - def mock_request_with_structured_output(self): - """Create a mock request with structured output.""" - request = Mock(spec=Request) - request.structured_output_request = Mock() - request.structured_output_request.reasoning_ended = None - request.structured_output_request.grammar = Mock() - request.structured_output_request.reasoning_parser_kwargs = None - request.structured_output_request.reasoner = None - request.structured_output_request.grammar.is_terminated = Mock( - return_value=False - ) - request.use_structured_output = True - request.prompt_token_ids = [1, 2, 3, 4, 5] - request.all_token_ids = [1, 2, 3, 4, 5, 6, 7, 8] - request.num_computed_tokens = 5 - request.num_output_placeholders = 0 - request.request_id = "mock_req" - return request - - @pytest.fixture - def manager_with_reasoner(self, mock_vllm_config): - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = MockReasoner - manager.tokenizer = Mock() - return manager - - def test_should_fill_bitmask_with_enable_in_reasoning( - self, manager_with_reasoner, mock_request_with_structured_output - ): - """Test should_fill_bitmask when enable_in_reasoning is True.""" - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_fill_bitmask_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """An active grammar does not need a request-local reasoner.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - assert result is True - assert structured_req.reasoner is None - - def test_grammar_bitmask_skips_reasoner_when_already_active( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Bitmask generation skips reasoning-boundary detection.""" - manager_with_reasoner.vllm_config.num_speculative_tokens = 1 - manager_with_reasoner._grammar_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.shape = (1, 1) - expected_bitmask = Mock() - manager_with_reasoner._grammar_bitmask.numpy.return_value = expected_bitmask - manager_with_reasoner._fill_bitmasks = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = True - request_id = mock_request_with_structured_output.request_id - - bitmask = manager_with_reasoner.grammar_bitmask( - requests={request_id: mock_request_with_structured_output}, - structured_output_request_ids=[request_id], - scheduled_spec_decode_tokens={}, - ) - - assert bitmask is expected_bitmask - assert structured_req.reasoner is None - - def test_should_fill_bitmask_without_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_fill_bitmask when enable_in_reasoning is False.""" - # Keep enable_in_reasoning as False (default) - config = manager_with_reasoner.vllm_config.structured_outputs_config - assert config.enable_in_reasoning is False - - result = manager_with_reasoner.should_fill_bitmask( - mock_request_with_structured_output - ) - - # Should set reasoning_ended and return its value - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is False - ) - assert result is False - - def test_should_fill_bitmask_no_reasoner( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test should_fill_bitmask when no reasoner is configured.""" - manager = StructuredOutputManager(mock_vllm_config) - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - # Should default to True when no reasoner - assert result is True - - def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs( - self, mock_vllm_config, mock_request_with_structured_output - ): - """Test request-level parser kwargs override the default reasoner.""" - - class KwargReasoner: - def __init__(self, tokenizer, chat_template_kwargs=None): - self.chat_template_kwargs = chat_template_kwargs or {} - - def is_reasoning_end(self, input_ids): - return not self.chat_template_kwargs.get("enable_thinking", False) - - manager = StructuredOutputManager(mock_vllm_config) - manager.reasoner_cls = KwargReasoner - manager.tokenizer = Mock() - - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_parser_kwargs = { - "chat_template_kwargs": {"enable_thinking": True} - } - - result = manager.should_fill_bitmask(mock_request_with_structured_output) - - assert result is False - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is not None - ) - - def test_should_advance_with_enable_in_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when enable_in_reasoning is True.""" - # Enable enable_in_reasoning - manager_with_reasoner.enable_in_reasoning = True - - # Should always return True when enable_in_reasoning is enabled - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_reasoning_not_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has not ended.""" - # Set reasoning as not ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return False since reasoning hasn't ended - assert result is False - - def test_should_advance_reasoning_just_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning ends in current step.""" - # Set reasoning as not ended initially, but ends in this step - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = True - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoner = reasoner - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # The scheduler trims the reasoning prefix before advancing the grammar. - assert ( - mock_request_with_structured_output.structured_output_request.reasoning_ended - is True - ) - assert result is True - - def test_should_advance_reasoning_already_ended( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Test should_advance when reasoning has already ended.""" - # Set reasoning as already ended - ( - mock_request_with_structured_output.structured_output_request - ).reasoning_ended = True - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # Should return True since reasoning has ended - assert result is True - assert ( - mock_request_with_structured_output.structured_output_request.reasoner - is None - ) - - def test_should_advance_uses_new_token_ids_when_provided( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Regression for #43388: when caller passes new_token_ids, the - reasoner sees the exact multi-token delta rather than the - placeholder-derived window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - - end_token_id = 248069 - - reasoner = MockReasoner(tokenizer=Mock()) - # Detection mirrors the real Qwen3 parser: end token in the delta. - reasoner.is_reasoning_end_streaming = Mock( - side_effect=lambda input_ids, delta_ids: end_token_id in list(delta_ids) - ) - structured_req.reasoner = reasoner - - # Scenario from #43388: async + spec decode K=4, 4 tokens accepted - # but only 1 placeholder remains (some drafts were rejected). - # The placeholder math would yield delta=[271] and miss . - # Passing new_token_ids must override that. - new_token_ids = [9, 198, end_token_id, 271] - mock_request_with_structured_output.all_token_ids = [ - 1, - 2, - 3, - 4, - 5, - ] + new_token_ids - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - # First call to is_reasoning_end_streaming was with the full - # new_token_ids (not the truncated placeholder window). - first_call = reasoner.is_reasoning_end_streaming.call_args_list[0] - _, called_delta = first_call.args - assert list(called_delta) == new_token_ids - - assert structured_req.reasoning_ended is True - assert result is True - - def test_should_advance_without_new_token_ids_falls_back( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Backward compat: callers that don't pass new_token_ids keep - the original placeholder-derived delta window. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - reasoner.is_reasoning_end_streaming.return_value = False - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - mock_request_with_structured_output.num_computed_tokens = 5 - mock_request_with_structured_output.num_output_placeholders = 2 - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output - ) - - # placeholder window: start = 5 - 2 = 3, delta = [4, 5] - _, called_delta = reasoner.is_reasoning_end_streaming.call_args[0] - assert list(called_delta) == [4, 5] - assert result is False - - def test_should_advance_whole_delta_fallback_pins_last_sequence_index( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """When only the whole-delta predicate fires, the boundary is the - last index of the sequence, even if the placeholder window - overshoots it and the delta is empty. - """ - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = MockReasoner(tokenizer=Mock()) - # Base-default shape: rescans the whole sequence, ignores the delta. - reasoner.is_reasoning_end_streaming.return_value = True - structured_req.reasoner = reasoner - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5] - # placeholder window: start = 9 - 1 = 8, past the end -> delta = [] - mock_request_with_structured_output.num_computed_tokens = 9 - mock_request_with_structured_output.num_output_placeholders = 1 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_trims_reasoning_prefix_for_json( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """JSON uses the common trim-then-advance path at the boundary.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.structured_output_key = ( - StructuredOutputOptions.JSON_OBJECT, - "{}", - ) - - marker = 248069 - - class MarkerReasoner: - def __init__(self, *_, **__): - pass - - def is_reasoning_end_streaming(self, input_ids, delta_ids): - return marker in list(delta_ids) - - structured_req.reasoner = MarkerReasoner() - - new_token_ids = [9, 198, marker, 271, 5005] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - result = manager_with_reasoner.should_advance( - mock_request_with_structured_output, - new_token_ids=new_token_ids, - ) - - structured_req.grammar.accept_tokens.assert_not_called() - assert structured_req.reasoning_ended is True - assert result is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271, 5005] - - def test_should_advance_engine_adapter_scans_only_the_delta( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Engine-based parsers locate the boundary from the delta alone.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - reasoner = make_qwen3_reasoner() - assert reasoner.reasoning_end_token_ids == {201, 202} - reasoner.is_reasoning_end_streaming = Mock( - side_effect=AssertionError("engine path must not rescan the prefix") - ) - structured_req.reasoner = reasoner - - new_token_ids = [9, 198, 201, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 - assert manager_with_reasoner.trim_reasoning_for_advance( - mock_request_with_structured_output, new_token_ids - ) == [271] - - def test_should_advance_engine_adapter_tool_start_ends_reasoning( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 202, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == 4 - - def test_should_advance_engine_adapter_no_boundary( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 198, 271] - mock_request_with_structured_output.all_token_ids = [1, 2, 3] + new_token_ids - - assert not manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_ended is False - - def test_should_advance_engine_adapter_never_reads_history( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Per-step cost must not depend on how much has been generated.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - new_token_ids = [9, 201] - total = (1 << 20) + len(new_token_ids) - mock_request_with_structured_output.all_token_ids = UntouchableTokenIds(total) - - assert manager_with_reasoner.should_advance( - mock_request_with_structured_output, new_token_ids=new_token_ids - ) - assert structured_req.reasoning_end_token_index == total - 1 - - def test_should_advance_engine_adapter_without_new_tokens_persists( - self, - manager_with_reasoner, - mock_request_with_structured_output, - ): - """Draft-validation callers persist the boundary like the legacy path.""" - structured_req = mock_request_with_structured_output.structured_output_request - structured_req.reasoning_ended = False - structured_req.reasoner = make_qwen3_reasoner() - - mock_request_with_structured_output.all_token_ids = [1, 2, 3, 4, 5, 201, 7] - # placeholder window: start = 7 - 2 = 5, delta = [201, 7] - mock_request_with_structured_output.num_computed_tokens = 7 - mock_request_with_structured_output.num_output_placeholders = 2 - - assert manager_with_reasoner.should_advance(mock_request_with_structured_output) - assert structured_req.reasoning_ended is True - assert structured_req.reasoning_end_token_index == 5 diff --git a/tests/v1/structured_output/test_structured_output_manager.py b/tests/v1/structured_output/test_structured_output_manager.py new file mode 100644 index 000000000000..17bda09b3efa --- /dev/null +++ b/tests/v1/structured_output/test_structured_output_manager.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from dataclasses import dataclass + +import pytest +from transformers import AutoTokenizer + +from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig +from vllm.config.model import ModelConfig +from vllm.config.speculative import SpeculativeConfig +from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.request import Request +from vllm.v1.structured_output import StructuredOutputManager + +TOKENIZER = "gpt2" +THINK_END = "\n" # reasoning-end marker (single GPT-2 token) +EOS = "<|eos|>" # resolved to tokenizer.eos_token_id +JSON_SCHEMA = '{"type": "object"}' +BACKENDS = ("xgrammar", "guidance") +MAX_WAIT_SECONDS = 5 +NUM_SPEC_TOKENS = 8 + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(TOKENIZER) + + +@dataclass(frozen=True) +class FlowCase: + """One validate -> bitmask -> accept -> post-accept bitmask scenario.""" + + raw_drafts: tuple[str, ...] + expected_row_pattern: str + expected_reasoning: bool | None + expect_terminated: bool + prefix: str = "" + expected_validated: tuple[str, ...] | None = None + reasoning_ended: bool | None = False + xfail_guidance: str | None = None + + +class MockReasoner: + def __init__(self, tokenizer, marker: int | None = None): + self.marker = marker + + def is_reasoning_end(self, input_ids): + if self.marker is None: + return True + return self.marker in list(input_ids) + + def is_reasoning_end_streaming(self, input_ids, delta_ids): + return self.is_reasoning_end(delta_ids) + + +def _single_token(tokenizer, text: str) -> int: + token_ids = tokenizer.encode(text) + assert len(token_ids) == 1, (text, token_ids) + return token_ids[0] + + +def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]: + """Convert literal text (or EOS sentinel) to token IDs.""" + return [ + tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t) + for t in texts + ] + + +def _wait_for_grammar(request: Request) -> None: + structured_req = request.structured_output_request + assert structured_req is not None + + deadline = time.time() + MAX_WAIT_SECONDS + while not structured_req._check_grammar_completion(): + if time.time() > deadline: + pytest.fail("Grammar compilation timed out") + time.sleep(0.01) + + +def _build_harness( + tokenizer, + backend: str, + prefix: str = "", + use_reasoner: bool = True, + reasoning_ended: bool | None = None, + enable_in_reasoning: bool = False, + reasoning_parser_kwargs: dict | None = None, +) -> tuple[StructuredOutputManager, Request]: + vllm_config = VllmConfig( + model_config=ModelConfig(tokenizer=TOKENIZER), + device_config=DeviceConfig(device="cpu"), + structured_outputs_config=StructuredOutputsConfig( + backend=backend, + enable_in_reasoning=enable_in_reasoning, + ), + speculative_config=SpeculativeConfig( + model="[ngram]", + num_speculative_tokens=NUM_SPEC_TOKENS, + ), + ) + manager = StructuredOutputManager(vllm_config) + if use_reasoner: + manager.reasoner_cls = MockReasoner + + sampling_params = SamplingParams( + structured_outputs=StructuredOutputsParams(json=JSON_SCHEMA) + ) + sampling_params.structured_outputs._backend = backend + sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) + + prompt_ids = tokenizer.encode(prefix) if prefix else [] + request = Request( + request_id=f"{backend}-flow", + prompt_token_ids=prompt_ids, + sampling_params=sampling_params, + pooling_params=None, + reasoning_ended=reasoning_ended, + reasoning_parser_kwargs=reasoning_parser_kwargs, + ) + manager.grammar_init(request) + _wait_for_grammar(request) + + structured_req = request.structured_output_request + assert request.prompt_token_ids is not None + assert structured_req is not None + assert structured_req.grammar is not None + if prompt_ids: + assert structured_req.grammar.accept_tokens( + request.request_id, + prompt_ids, + ) + return manager, request + + +def _row_pattern(bitmask) -> str: + return "".join("U" if (row == -1).all() else "C" for row in bitmask) + + +def _run_real_flow( + manager: StructuredOutputManager, + request: Request, + raw_drafts: list[int], + expected_validated: list[int], + expected_row_pattern: str, + expected_reasoning: bool | None, + expect_terminated: bool, +) -> None: + assert len(expected_validated) <= len(raw_drafts) + + validated = manager.validate_tokens(request, list(raw_drafts)) + assert validated == expected_validated + + padded = validated + [-1] * (len(raw_drafts) - len(validated)) + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={request.request_id: padded}, + ) + assert bitmask is not None + assert bitmask.shape[0] == len(raw_drafts) + 1 + assert _row_pattern(bitmask) == expected_row_pattern + + # grammar_bitmask() must rollback any speculative state it advanced. + assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated + + structured_req = request.structured_output_request + assert structured_req is not None + grammar = structured_req.grammar + assert grammar is not None + assert not grammar.is_terminated() + # `bitmask[i]` is the grammar state after the first `i` scheduled tokens. + # These tests commit the validated prefix, so the matching post-accept + # state is already present in the scheduled bitmask. + expected_post_accept_row = bitmask[len(expected_validated)].copy() + + # Mirror the scheduler flow: sampled tokens are appended before + # StructuredOutputManager.accept_tokens() trims them to the + # grammar-constrained suffix. + request.append_output_token_ids(list(expected_validated)) + assert manager.accept_tokens(request, list(expected_validated)) + + post_accept_bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert post_accept_bitmask is not None + assert post_accept_bitmask.shape[0] == 1 + assert (post_accept_bitmask[0] == expected_post_accept_row).all() + assert structured_req.reasoning_ended == expected_reasoning + assert grammar.is_terminated() is expect_terminated + + +FLOW_CASES = [ + pytest.param( + FlowCase( + raw_drafts=(" ", "z", " "), + expected_row_pattern="UUUU", + expected_reasoning=False, + expect_terminated=False, + ), + id="inactive", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}'), + expected_row_pattern="CCCCC", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_all_valid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=("z", ":", ' "', "b"), + expected_validated=(), + expected_row_pattern="CUUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_first_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", "z", ' "', "b"), + expected_validated=(":",), + expected_row_pattern="CCUUU", + expected_reasoning=True, + expect_terminated=False, + reasoning_ended=True, + ), + id="active_later_token_invalid", + ), + pytest.param( + FlowCase( + prefix='{"a": "b"}', + raw_drafts=(EOS, " ", " "), + expected_validated=(EOS,), + expected_row_pattern="CUUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + xfail_guidance=( + "guidance validate_tokens rejects EOS on an already-complete " + "object even though accept_tokens would accept it" + ), + ), + id="active_first_token_terminates", + ), + pytest.param( + FlowCase( + prefix='{"a"', + raw_drafts=(":", ' "', "b", '"}', EOS, " "), + expected_validated=(":", ' "', "b", '"}', EOS), + expected_row_pattern="CCCCCUU", + expected_reasoning=True, + expect_terminated=True, + reasoning_ended=True, + ), + id="active_later_token_terminates", + ), + pytest.param( + FlowCase( + raw_drafts=(THINK_END, "{", ' "', "b"), + expected_row_pattern="UCCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_first_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", ' "'), + expected_row_pattern="UUCCC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_middle_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", " ", THINK_END), + expected_row_pattern="UUUC", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_on_last_token", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "z", "{"), + expected_validated=(" ", THINK_END), + expected_row_pattern="UUCUU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_first_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "z"), + expected_validated=(" ", THINK_END, "{"), + expected_row_pattern="UUCCU", + expected_reasoning=True, + expect_terminated=False, + ), + id="becomes_active_later_token_invalid", + ), + pytest.param( + FlowCase( + raw_drafts=(" ", THINK_END, "{", "}", EOS, " "), + expected_validated=(" ", THINK_END, "{", "}", EOS), + expected_row_pattern="UUCCCUU", + expected_reasoning=True, + expect_terminated=True, + ), + id="becomes_active_terminates", + ), +] + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("case", FLOW_CASES) +def test_real_flow( + tokenizer, + backend: str, + case: FlowCase, +): + if backend == "guidance" and case.xfail_guidance: + pytest.xfail(case.xfail_guidance) + + reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} + manager, request = _build_harness( + tokenizer, + backend, + prefix=case.prefix, + reasoning_ended=case.reasoning_ended, + reasoning_parser_kwargs=reasoner_kwargs, + ) + + raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) + expected_texts = ( + case.raw_drafts if case.expected_validated is None else case.expected_validated + ) + expected_validated = _to_token_ids(tokenizer, expected_texts) + + _run_real_flow( + manager, + request, + raw_drafts=raw_drafts, + expected_validated=expected_validated, + expected_row_pattern=case.expected_row_pattern, + expected_reasoning=case.expected_reasoning, + expect_terminated=case.expect_terminated, + ) + + +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize( + ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), + [ + pytest.param(True, False, True, id="enable_in_reasoning"), + pytest.param(True, True, False, id="reasoning_ended"), + pytest.param(False, False, False, id="no_reasoner"), + pytest.param(True, None, False, id="inferred"), + ], +) +def test_initial_constraint_activation( + tokenizer, + backend: str, + use_reasoner: bool, + reasoning_ended: bool | None, + enable_in_reasoning: bool, +): + manager, request = _build_harness( + tokenizer, + backend, + use_reasoner=use_reasoner, + reasoning_ended=reasoning_ended, + enable_in_reasoning=enable_in_reasoning, + reasoning_parser_kwargs={ + "marker": _single_token(tokenizer, THINK_END) + if reasoning_ended is not None + else None + }, + ) + + # "{" is valid JSON start; "z" is not. Truncation proves grammar is active. + open_brace = _single_token(tokenizer, "{") + z = _single_token(tokenizer, "z") + assert manager.validate_tokens(request, [open_brace, z]) == [open_brace] + + bitmask = manager.grammar_bitmask( + requests={request.request_id: request}, + structured_output_request_ids=[request.request_id], + scheduled_spec_decode_tokens={}, + ) + assert bitmask is not None + request.append_output_token_ids([open_brace]) + assert manager.accept_tokens(request, [open_brace]) + + # Short-circuit paths must not construct a request-local reasoner (#52573). + if enable_in_reasoning or reasoning_ended is True: + structured_req = request.structured_output_request + assert structured_req is not None + assert structured_req.reasoner is None diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index 11b3e5c01b38..f5f80202205b 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string from vllm.v1.structured_output.backend_xgrammar import ( has_xgrammar_unsupported_json_features, @@ -10,13 +12,19 @@ pytestmark = pytest.mark.cpu_test +def grammar_accepts(schema: dict, text: str) -> bool: + return _is_grammar_accept_string(Grammar.from_json_schema(schema), text) + + @pytest.fixture def unsupported_string_schemas(): return [ {"type": "string", "format": "non_existing_format"}, + # TODO(arpera): # pattern/format is compiled but length bounds are silently dropped, # so the combination must be rejected instead of producing quietly # wrong output + # https://github.com/mlc-ai/xgrammar/issues/749 {"type": "string", "pattern": "^a+$", "maxLength": 2}, {"type": "string", "pattern": "^a+$", "minLength": 3}, {"type": "string", "format": "email", "maxLength": 10}, @@ -48,15 +56,93 @@ def unsupported_array_schemas(): @pytest.fixture -def unsupported_object_schemas(): +def unsupported_property_names_combinations(): + return [ + # TODO(arpera): + # this case is not covered by xgrammar, so we should report this bug to xgrammar + # xgrammar drops propertyNames whenever patternProperties is present, + # so "grade_12" matches the pattern and escapes the name constraint + { + "type": "object", + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, + "propertyNames": {"pattern": "^grade_[0-9]$"}, # does NOT match "grade_12" + }, + # propertyNames makes xgrammar discard the sibling additionalProperties + # value schema: https://github.com/mlc-ai/xgrammar/issues/826 + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "additionalProperties": {"type": "integer"}, + }, + # propertyNames is a string schema that conventionally omits "type", so + # it escapes the string check while xgrammar drops its length bound all + # the same: https://github.com/mlc-ai/xgrammar/issues/749 + { + "type": "object", + "propertyNames": {"pattern": "^a+$", "maxLength": 2}, + }, + # xgrammar emits named `properties` separately from `propertyNames` and + # only applies `propertyNames` in its additional-properties branch, so a + # key declared in `properties` escapes the name constraint regardless of + # what `propertyNames` contains. + { + "type": "object", + "properties": {"Bad": {"type": "integer"}}, + "required": ["Bad"], + "propertyNames": {"pattern": "^[a-z]+$"}, + }, + { + "type": "object", + "properties": {"Bad": {"type": "integer"}}, + "propertyNames": {"enum": ["good"]}, + }, + # xgrammar's propertyNames-only branch falls back to an unconstrained + # value type, so it drops unevaluatedProperties whether it restricts + # values (a schema) or forbids them (false). + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "unevaluatedProperties": {"type": "integer"}, + }, + { + "type": "object", + "propertyNames": {"pattern": "^[a-z]+$"}, + "unevaluatedProperties": False, + }, + ] + + +@pytest.fixture +def unsupported_pattern_properties_combinations(): return [ - {"type": "object", "propertyNames": {"pattern": "^[a-z]+$"}}, - {"type": "object", "patternProperties": {"^S": {"type": "string"}}}, + # JSON Schema requires a property matched by both `properties` and + # `patternProperties` to satisfy both (conjunction), but xgrammar + # compiles them as alternative branches, so satisfying either one is + # enough. + { + "type": "object", + "properties": {"x": {"type": "string"}}, + "patternProperties": {"^x$": {"type": "integer"}}, + }, + # The same alternative-branches problem applies to overlapping + # patternProperties patterns: a key matching both patterns only has to + # satisfy one of them. + { + "type": "object", + "patternProperties": { + "^a[a-z]*$": {"type": "string"}, + "^[a-z]*z$": {"type": "integer"}, + }, + }, ] @pytest.fixture -def supported_schema(): +def supported_frankenstein_schema(): + # IMPORTANT(arpera): + # Do NOT add more keywords here! This schema is overcrowded enough that a new + # keyword can end up having no effect on the compiled grammar while the + # test still passes. Give a new keyword its own fixture and named test. return { "type": "object", "properties": { @@ -88,6 +174,19 @@ def supported_schema(): } +@pytest.fixture +def property_names_schema(): + return {"type": "object", "propertyNames": {"pattern": "^[a-z_]+$"}} + + +@pytest.fixture +def pattern_properties_schema(): + return { + "type": "object", + "patternProperties": {"^grade_[0-9]+$": {"type": "integer"}}, + } + + @pytest.mark.parametrize( "schema_type", [ @@ -95,7 +194,8 @@ def supported_schema(): "unsupported_integer_schemas", "unsupported_number_schemas", "unsupported_array_schemas", - "unsupported_object_schemas", + "unsupported_property_names_combinations", + "unsupported_pattern_properties_combinations", ], ) def test_unsupported_json_features_by_type(schema_type, request): @@ -106,7 +206,29 @@ def test_unsupported_json_features_by_type(schema_type, request): ) -def test_supported_json_features(supported_schema): - assert not has_xgrammar_unsupported_json_features(supported_schema), ( - "Schema should be supported" +@pytest.mark.parametrize( + "schema_type", + [ + "supported_frankenstein_schema", + "pattern_properties_schema", + "property_names_schema", + ], +) +def test_supported_json_features(schema_type, request): + schema = request.getfixturevalue(schema_type) + assert not has_xgrammar_unsupported_json_features(schema), ( + f"Schema should be supported: {schema}" ) + + +def test_property_names_constrains_keys(property_names_schema): + assert grammar_accepts(property_names_schema, '{"score": 5}') + assert grammar_accepts(property_names_schema, '{"score": "seven"}') + assert not grammar_accepts(property_names_schema, '{"Score": 5}') + assert not grammar_accepts(property_names_schema, '{"score_1": 5}') + + +def test_pattern_properties_constrains_keys_and_values(pattern_properties_schema): + assert grammar_accepts(pattern_properties_schema, '{"grade_1": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"other": 5}') + assert not grammar_accepts(pattern_properties_schema, '{"grade_1": "five"}') diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index af417bc03e6f..a23189438475 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -137,7 +137,7 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): @pytest.mark.parametrize( "schema, expected_backend", [ - # multipleOf is unsupported by xgrammar, patternProperties also by guidance. + # multipleOf is unsupported by xgrammar. ( { "type": "object", @@ -145,8 +145,13 @@ def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): }, "guidance", ), + # patternProperties + properties is also unsupported by guidance. ( - {"type": "object", "patternProperties": {"^a": {"type": "string"}}}, + { + "type": "object", + "properties": {"a": {"type": "string"}}, + "patternProperties": {"^a$": {"type": "string"}}, + }, "outlines", ), ], diff --git a/tests/v1/worker/test_dspark_utils.py b/tests/v1/worker/test_dspark_utils.py new file mode 100644 index 000000000000..94655f192c89 --- /dev/null +++ b/tests/v1/worker/test_dspark_utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass, field + +from vllm.v1.worker.gpu.spec_decode.dspark.utils import _get_dspark_parallel_config + + +@dataclass +class _FakeEPLBConfig: + num_redundant_experts: int = 0 + + +@dataclass +class _FakeParallelConfig: + pipeline_parallel_size: int = 2 + tensor_parallel_size: int = 8 + enable_eplb: bool = True + eplb_config: _FakeEPLBConfig = field( + default_factory=lambda: _FakeEPLBConfig(num_redundant_experts=32) + ) + enable_elastic_ep: bool = True + + def __post_init__(self) -> None: + if not self.enable_eplb and self.eplb_config.num_redundant_experts: + raise ValueError("redundant experts require EPLB") + if self.enable_elastic_ep and not self.enable_eplb: + raise ValueError("elastic EP requires EPLB") + + +def test_dspark_parallel_config_disables_eplb_atomically(): + target_config = _FakeParallelConfig() + + draft_config = _get_dspark_parallel_config( + target_config, + tensor_parallel_size=4, + ) + + assert target_config.pipeline_parallel_size == 2 + assert target_config.tensor_parallel_size == 8 + assert target_config.enable_eplb + assert target_config.eplb_config.num_redundant_experts == 32 + assert target_config.enable_elastic_ep + + assert draft_config is not target_config + assert draft_config.pipeline_parallel_size == 1 + assert draft_config.tensor_parallel_size == 4 + assert not draft_config.enable_eplb + assert draft_config.eplb_config.num_redundant_experts == 0 + assert not draft_config.enable_elastic_ep + assert draft_config.eplb_config is not target_config.eplb_config diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index dc935929d4c4..2bd7be2cfa34 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -1,25 +1,81 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock import numpy as np import pytest import torch +from vllm.config import CUDAGraphMode from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context from vllm.v1.worker.gpu.pcp_manager import PCPManager -def _copy_to_cpu(value, out=None, device=None): +def _copy_to_cpu(value, device=None, dtype=None, out=None): tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value if out is not None: return out.copy_(tensor) return tensor +def _make_config(cudagraph_mode: CUDAGraphMode): + return SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=2, + decode_context_parallel_size=1, + pipeline_parallel_size=1, + dcp_comm_backend="ag_rs", + ), + model_config=SimpleNamespace( + use_mla=True, + is_encoder_decoder=False, + hf_text_config=SimpleNamespace(), + ), + lora_config=None, + speculative_config=None, + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + ) + + +def _make_capture_manager(block_table: torch.Tensor): + block_tables = SimpleNamespace( + input_block_tables=(block_table,), + num_kv_cache_groups=1, + kernel_block_sizes=(2,), + blocks_per_kv_block=(1,), + ) + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + max_num_reqs=4, + max_num_tokens=8, + block_tables=block_tables, + ) + return manager, block_tables + + +@pytest.mark.parametrize( + "cudagraph_mode", + [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], +) +def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode): + PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False) + + +def test_validate_config_rejects_full_graph_for_prefills(): + with pytest.raises(NotImplementedError, match="decode-only routines"): + PCPManager.validate_config( + _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False + ) + + def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager = PCPManager( pcp_world_size=2, @@ -27,7 +83,7 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): device=torch.device("cpu"), dcp_world_size=1, ) - monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) segments_by_rank, per_rank_num_tokens = manager._build_batch_layout( num_scheduled_tokens=np.ones(3, dtype=np.int32), @@ -102,7 +158,7 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): device=torch.device("cpu"), dcp_world_size=1, ) - monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + monkeypatch.setattr(pcp_manager_module, "async_tensor_h2d", _copy_to_cpu) with pytest.raises(ValueError, match="smaller than the largest rank-local batch"): manager._build_batch_layout( @@ -114,6 +170,89 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): ) +@pytest.mark.parametrize( + ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"), + [ + (CUDAGraphMode.NONE, 4, None, None), + (CUDAGraphMode.PIECEWISE, None, 8, None), + (CUDAGraphMode.FULL, 4, 8, 4), + ], +) +def test_partition_padding_is_derived_from_batch_descriptor( + cg_mode, num_reqs, expected_tokens, expected_reqs +): + manager = MagicMock() + input_batch = MagicMock() + manager.partition_batch.return_value = input_batch + batch_desc = BatchExecutionDescriptor( + cg_mode=cg_mode, + num_tokens=8, + num_reqs=num_reqs, + ) + + result = pcp_manager_module.maybe_partition_pcp_batch( + manager, + input_batch, + batch_desc, + ) + + assert result is input_batch + manager.partition_batch.assert_called_once_with( + input_batch, + padded_num_tokens=expected_tokens, + padded_num_reqs=expected_reqs, + ) + + +def test_capture_uses_pcp_persistent_inputs(): + manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32)) + + dummy_batch = InputBatch.make_dummy( + num_reqs=4, + num_tokens=4, + input_buffers=InputBuffers(4, 8, torch.device("cpu")), + max_query_len=1, + ) + dummy_batch.input_ids.copy_(torch.arange(4)) + dummy_batch.positions.fill_(7) + input_batch = manager.prepare_inputs_to_capture(dummy_batch) + + assert input_batch is not dummy_batch + assert input_batch.req_ids == dummy_batch.req_ids + for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"): + actual = getattr(input_batch, name) + torch.testing.assert_close(actual, getattr(dummy_batch, name)) + assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr() + + +def test_dummy_context_updates_pcp_local_block_tables(): + global_block_table = torch.full((4, 4), -1, dtype=torch.int32) + manager, block_tables = _make_capture_manager(global_block_table) + input_batch = InputBatch.make_dummy( + num_reqs=2, + num_tokens=2, + input_buffers=manager.input_buffers, + max_query_len=1, + ) + input_batch = manager.prepare_inputs_to_capture(input_batch) + local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs) + + set_dummy_context( + input_batch, + block_tables, + context_len=3, + num_kv_blocks=16, + max_model_len=16, + input_block_tables=local_block_tables, + ) + + torch.testing.assert_close( + local_block_tables[0][:2, :2], + torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + ) + assert torch.all(global_block_table == -1) + + def _rank_rows( pcp_rank: int, pcp_world_size: int, @@ -331,10 +470,22 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2] global_batch.dcp_local_seq_lens.fill_(-1) - local_batch = manager.partition_batch(global_batch, padded_num_tokens=2) + local_batch = manager.partition_batch( + global_batch, + padded_num_tokens=4, + padded_num_reqs=4, + ) assert local_batch.dcp_local_seq_lens is None - assert local_batch.seq_lens.tolist() == [17, 25] + assert local_batch.seq_lens.tolist() == [17, 25, 0, 0] + assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None + expected_dcp_upper_bound = get_dcp_local_seq_lens( + torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + ) + assert torch.equal( + local_batch.dcp_local_seq_lens_cpu_upper_bound, + torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]), + ) # What execute_model does next: derive DCP metadata from the final batch # on the PCP-owned buffers. @@ -348,7 +499,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): num_reqs_padded=local_batch.num_reqs_after_padding, ) expected = get_dcp_local_seq_lens( - torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1 ) assert local_batch.dcp_local_seq_lens is not None assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected) diff --git a/tests/v1/worker/test_worker_memory_snapshot.py b/tests/v1/worker/test_worker_memory_snapshot.py index fe8a5a21f8dc..6ed62cb771fa 100644 --- a/tests/v1/worker/test_worker_memory_snapshot.py +++ b/tests/v1/worker/test_worker_memory_snapshot.py @@ -12,6 +12,7 @@ from vllm.config import set_current_vllm_config from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment @@ -66,8 +67,22 @@ def worker_process( model="facebook/opt-125m", tensor_parallel_size=2, load_format="dummy" ).create_engine_config() + # Select worker class and patch target based on platform + if current_platform.is_xpu(): + from vllm.v1.worker.xpu_worker import XPUWorker + + worker_cls = XPUWorker + init_patch_target = ( + "vllm.v1.worker.xpu_worker.init_worker_distributed_environment" + ) + else: + worker_cls = Worker + init_patch_target = ( + "vllm.v1.worker.gpu_worker.init_worker_distributed_environment" + ) + # Create worker - worker = Worker( + worker = worker_cls( vllm_config=vllm_config, local_rank=rank, rank=rank, @@ -81,7 +96,7 @@ def worker_process( # Apply minimal patches to track operation order init_patch = patch( - "vllm.v1.worker.gpu_worker.init_worker_distributed_environment", + init_patch_target, side_effect=make_operation_tracker( "init_distributed", original_init_worker ), @@ -93,7 +108,7 @@ def worker_process( ) all_reduce_patch = patch( "torch.distributed.all_reduce", - side_effect=make_operation_tracker("nccl_all_reduce", original_all_reduce), + side_effect=make_operation_tracker("dist_all_reduce", original_all_reduce), ) with ( @@ -182,13 +197,13 @@ def test_init_distributed_is_called_before_memory_snapshot(): # Raises ValueError if the operation is not found init_distributed = rank_ops.index("init_distributed") - nccl_all_reduce = rank_ops.index("nccl_all_reduce") + dist_all_reduce = rank_ops.index("dist_all_reduce") memory_snapshot = rank_ops.index("memory_snapshot") # Verify order: init_distributed should happen before memory_snapshot - assert init_distributed < nccl_all_reduce < memory_snapshot, ( + assert init_distributed < dist_all_reduce < memory_snapshot, ( f"Rank {rank}: init_distributed (index {init_distributed}) " - f"must happen before nccl_all_reduce (index {nccl_all_reduce}) " + f"must happen before dist_all_reduce (index {dist_all_reduce}) " f"and memory_snapshot (index {memory_snapshot})" ) diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 67a527405e2d..f4438c67552f 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -50,7 +50,7 @@ str(src / "deep_gemm/include"), str(src / "third-party/cutlass/include"), str(src / "third-party/cutlass/tools/util/include"), - str(src / "third-party/fmt/include"), + str(src / "third-party/deep_jit/include"), *cpp_extension.include_paths(device_type="cuda"), ] diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index f6a64d8796c1..db81bf4be310 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,11 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" -# NOTE: This is currently targeting the nv_dev branch tip due to sm120 support -DEEPGEMM_GIT_REF="8b1392b978f5a03c828dd1711090d7fb50958b8a" +DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" +# NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 +# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE +# activations +DEEPGEMM_GIT_REF="ad1f1726aa540a76c1d26d6a120effb8de21eaa4" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 4cfb3b388486..8f2b2a24adac 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -583,10 +583,11 @@ def _rocm_aiter_mla_decode_fwd_impl( reduce_indptr: torch.Tensor | None = None, reduce_final_map: torch.Tensor | None = None, reduce_partial_map: torch.Tensor | None = None, + causal: bool = True, ) -> None: from aiter.mla import mla_decode_fwd - kwargs: dict[str, float | torch.Tensor | None] = { + kwargs: dict[str, float | torch.Tensor | None | bool] = { "sm_scale": sm_scale, "logit_cap": logit_cap, } @@ -596,6 +597,8 @@ def _rocm_aiter_mla_decode_fwd_impl( kwargs["q_scale"] = q_scale kwargs["kv_scale"] = kv_scale + kwargs["causal"] = causal + if work_meta_data is not None: assert work_indptr is not None, ( "work_indptr must be provided with work_meta_data" @@ -1749,6 +1752,30 @@ def _mhc_delayed_pre_tail( _OPS_REGISTERED = False +def _sync_aiter_situv2_moe_env() -> None: + """Mirror the SiTUv2 MoE toggle into AITER's a4w4 dispatch env. + + AITER selects afp8 vs afp4 activation kernels via AITER_SITUV2_A8W4 / + AITER_SITUV2_A4W4 (see ROCm/aiter fused_moe.py, A8W4 checked first). + When VLLM_ROCM_USE_AITER_MOE_SITUV2 is enabled we route to a4w4 + (afp4_wfp4_fp4 kernels) and clear any legacy AITER_SITUV2_A8W4 override. + + Requires AITER with ROCm/aiter#4463 (first tagged in v0.1.20): a4w4 + dispatch plus kimik3_a4w4_{un,}tuned_fmoe.csv. Older AITER still runs + afp4 FlyDSL but falls back to heuristic configs, which is not the + tuned a4w4 path this recipe is meant to use. + """ + import os + + import vllm.envs as envs + + if envs.VLLM_ROCM_USE_AITER_MOE_SITUV2: + os.environ["AITER_SITUV2_A4W4"] = "1" + os.environ.pop("AITER_SITUV2_A8W4", None) + else: + os.environ.pop("AITER_SITUV2_A4W4", None) + + class rocm_aiter_ops: """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM. @@ -1771,7 +1798,7 @@ class rocm_aiter_ops: VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply. VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings. VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion. - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: Controls a8w4 SiTU fused MoE variant. + VLLM_ROCM_USE_AITER_MOE_SITUV2: Controls SiTUv2 FlyDSL MoE (a4w4). VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM. Note: @@ -1838,7 +1865,7 @@ def get_moe_dispatch_policy(cls) -> int: # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - _MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 + _MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 # TODO: Consolidate under _LINEAR_ENABLED _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM # Lazily probed: whether aiter.topk_softmax supports the @@ -1867,7 +1894,8 @@ def refresh_env_variables(cls): cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS - cls._MOE_SITUV2_A8W4 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 + cls._MOE_SITUV2 = envs.VLLM_ROCM_USE_AITER_MOE_SITUV2 + _sync_aiter_situv2_moe_env() cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM @staticmethod @@ -1980,10 +2008,10 @@ def is_fusion_moe_shared_experts_enabled(cls) -> bool: @classmethod @if_aiter_supported - def is_fused_moe_situv2_a8w4_enabled(cls) -> bool: - # _MOE_SITUV2_A8W4 is a variant of aiter fused moe, so aiter + def is_fused_moe_situv2_enabled(cls) -> bool: + # _MOE_SITUV2 is a variant of aiter fused moe, so aiter # fused moe must be enabled as well. - return cls.is_fused_moe_enabled() and cls._MOE_SITUV2_A8W4 + return cls.is_fused_moe_enabled() and cls._MOE_SITUV2 @classmethod @if_aiter_supported @@ -2036,6 +2064,26 @@ def fuse_sigmoid_in_kernel(cls, aiter_topK_meta_data: object) -> bool: def is_mla_enabled(cls) -> bool: return cls._AITER_ENABLED and cls._MLA_ENABLED + @classmethod + @if_aiter_supported + def mla_decode_supports_non_causal(cls) -> bool: + """Whether installed aiter.mla.mla_decode_fwd accepts `causal`. + + Added in aiter v0.1.20. A missing argument is a hard error rather + than a boolean the caller can use to pick another MLA backend. + """ + import inspect + + from aiter.mla import mla_decode_fwd + + if "causal" not in inspect.signature(mla_decode_fwd).parameters: + raise RuntimeError( + "ROCM_AITER_MLA requires aiter.mla.mla_decode_fwd(..., causal=). " + "The installed aiter is causal-only; upgrade aiter rather than " + "falling back to another MLA backend." + ) + return True + @classmethod @if_aiter_supported def is_mha_enabled(cls) -> bool: @@ -2858,6 +2906,7 @@ def mla_decode_fwd( reduce_indptr: torch.Tensor | None = None, reduce_final_map: torch.Tensor | None = None, reduce_partial_map: torch.Tensor | None = None, + causal: bool = True, ): torch.ops.vllm.rocm_aiter_mla_decode_fwd( q, @@ -2878,6 +2927,7 @@ def mla_decode_fwd( reduce_indptr=reduce_indptr, reduce_final_map=reduce_final_map, reduce_partial_map=reduce_partial_map, + causal=causal, ) @staticmethod @@ -4004,3 +4054,4 @@ def mhc_fused_post_pre( rocm_aiter_ops.register_ops_once() +_sync_aiter_situv2_moe_env() diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index b55c5f4e164d..54c6ce63b0fa 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -606,6 +606,7 @@ async def _run_pooling_request( headers: dict[str, Any], pbar: tqdm | None = None, num_input_sequences: int = 1, + prompt_len: int = 0, ) -> RequestFuncOutput: output = RequestFuncOutput(num_input_sequences=num_input_sequences) st = time.perf_counter() @@ -613,11 +614,19 @@ async def _run_pooling_request( try: async with session.post(url=api_url, headers=headers, json=payload) as response: if response.status == 200: + encoding_format = payload.get("encoding_format", "float") + if encoding_format in ("bytes", "bytes_only"): + async for _ in response.content.iter_any(): + pass + else: + await response.read() output.ttft = output.latency = time.perf_counter() - st - if payload.get("encoding_format", "float") == "bytes": + if encoding_format == "bytes": metadata = json.loads(response.headers["metadata"]) usage = metadata.get("usage", {}) + elif encoding_format == "bytes_only": + usage = {"prompt_tokens": prompt_len} else: data = await response.json() usage = data.get("usage", {}) @@ -672,6 +681,7 @@ async def async_request_openai_embeddings( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) @@ -709,6 +719,7 @@ async def async_request_vllm_rerank( headers=headers, pbar=pbar, num_input_sequences=len(request_func_input.prompt) - 1, + prompt_len=request_func_input.prompt_len, ) @@ -743,6 +754,7 @@ async def async_request_openai_embeddings_chat( payload=payload, headers=headers, pbar=pbar, + prompt_len=request_func_input.prompt_len, ) @@ -847,6 +859,7 @@ async def async_request_infinity_embeddings( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) @@ -896,6 +909,7 @@ async def async_request_vllm_pooling( headers=headers, pbar=pbar, num_input_sequences=_get_num_input_sequences(request_func_input.prompt), + prompt_len=request_func_input.prompt_len, ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 05da048bdbd6..3646677728b5 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -86,9 +86,20 @@ class AttentionConfig: indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + indexer_sparse_logits: bool = False + """DeepSeek V4.1 two-level indexer: score only the candidate blocks with + DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over + the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`, + an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension + (the top-k runs on the kernels' bf16 logits). The sparse path costs + O(candidate blocks) per query regardless of context length, so it pays off + for long contexts (roughly 32K tokens and beyond) and is slower below.""" + hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental - Model Runner V2-only HiSparse sparse-MLA decode hot-buffering.""" + Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred + with defaults when HiSparseConnector is configured (directly or via + MultiConnector); set it explicitly only to tune its fields.""" use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 9bd898731eb1..ce3199c52a25 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -211,6 +211,8 @@ class CacheConfig: """The number of blocks to allocate for CPU memory.""" # Set after KV cache initialization. + effective_attention_block_size: int | None = field(default=None, init=False) + """Full-attention block size in tokens, including DCP, or None if unavailable.""" kv_cache_size_tokens: int | None = field(default=None, init=False) """Per-DP-engine KV cache capacity in tokens (group-aware). Uses group-aware capacity since num_gpu_blocks * block_size can be wrong @@ -277,6 +279,7 @@ def compute_hash(self) -> str: # Post-init/derived counters "num_gpu_blocks", "num_cpu_blocks", + "effective_attention_block_size", "kv_cache_size_tokens", "kv_cache_max_concurrency", # WIP feature toggle not impacting compiled graph shape diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 7d2baea26ce9..7167ca5a9063 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -27,7 +27,7 @@ from vllm.utils import random_uuid from vllm.utils.hashing import safe_hash -from .attention import AttentionConfig +from .attention import AttentionConfig, HiSparseConfig from .cache import CacheConfig from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode from .device import DeviceConfig @@ -650,6 +650,29 @@ def num_lookahead_tokens(self) -> int: return self.num_speculative_tokens return 0 + @property + def num_prefill_lookahead_tokens(self) -> int: + """Prefill tokens past the computed range that the drafter reads. + + Mid-prefill the drafter consumes tokens the target model has not been + scheduled for yet, so every component that has to keep them available + must apply this margin: the scheduler, which never ends a chunk within + it and shifts encoder scheduling by it, and the KV cache manager, which + treats the trailing `this - 1` tokens as re-prefillable rather than + finalized. Consumers must read this property rather than re-deriving + their own per-method lookahead, so those components cannot drift apart. + """ + speculative_config = self.speculative_config + if speculative_config is None or not speculative_config.use_eagle(): + return 0 + if speculative_config.use_multi_module_mtp(): + # Each MTP module reads one token further ahead than the one before + # it, so the chain needs num_speculative_tokens of runway at a + # chunked-prefill boundary. + return self.num_speculative_tokens + # Eagle-family drafters read only the immediate next token. + return 1 + @property def uniform_decode_query_len(self) -> int: """Query length of every request in a uniform decode batch. @@ -1630,6 +1653,13 @@ def has_blocked_weights(): self._maybe_disable_dynamic_sd_for_data_parallel() self._maybe_override_dynamic_sd_cudagraph_mode() + if ( + self.attention_config.hisparse_config is None + and self.kv_transfer_config is not None + and self.kv_transfer_config.has_connector("HiSparseConnector") + ): + self.attention_config.hisparse_config = HiSparseConfig() + if self.attention_config.hisparse_config is not None: if not current_platform.is_cuda(): raise ValueError("HiSparse currently requires NVIDIA CUDA.") @@ -2829,16 +2859,16 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: unsupported.append("pipeline parallelism with external_launcher") if speculative_config is not None: - # TODO: ngram / ngram_gpu are not supported by the v2 model runner yet - if speculative_config.method in ("ngram", "ngram_gpu"): - unsupported.append("ngram/ngram_gpu speculative decoding") - elif speculative_config.method not in ( - "eagle", - "eagle3", - "mtp", - "dflash", - "dspark", - "extract_hidden_states", + if speculative_config.method in ( + # https://github.com/vllm-project/vllm/pull/40704 + "ngram", + "ngram_gpu", + # https://github.com/vllm-project/vllm/pull/43091 + "draft_model", + "suffix", + "medusa", + "mlp_speculator", + "custom_class", ): unsupported.append(f"speculative method '{speculative_config.method}'") diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 5c01d996968e..54184aa33945 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -113,7 +113,7 @@ def __init__( # New requests are added by update_state_after_alloc in # the scheduler. Used to make metadata passed to Worker. self._reqs_need_recv: dict[ - ReqId, tuple[Request, BlockIds, tuple[int, ...]] + ReqId, tuple[Request, BlockIds, tuple[int, ...], bool] ] = {} self._reqs_need_save: dict[ReqId, Request] = {} # Reqs to send and their expiration time @@ -199,8 +199,8 @@ def shutdown(self): def on_new_request(self, request: "Request") -> None: """Track a request that may need heartbeats.""" params = request.kv_transfer_params - if params is not None and params.get("do_remote_decode") and self._has_mamba: - self._truncate_mamba_request_for_prefill(request) + if params is not None and params.get("do_remote_decode"): + self._truncate_request_for_prefill(request) # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are # effectively disabled for Bidirectional KV transfer. @@ -382,36 +382,56 @@ def _nixl_handshake_listener( (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts) ) + def _prefill_backoff(self) -> int: + """Trailing prompt tokens the prefiller must not compute; the decoder + recomputes them locally. + + Mamba needs h(N-1) so the decoder can derive h(N) itself. Multi-module + MTP needs to keep its whole lookahead window off of the prefiller, which + would otherwise embed the unverified drafts in the MTP layer's KV cache. The + decoder would never rebuild them, because the update is sized by the rejection + count, which is zero for the first decode. + """ + return max( + 1 if self._has_mamba else 0, + self.vllm_config.num_prefill_lookahead_tokens - 1, + ) + def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: - """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: - return num_prompt_tokens - 1 + """D-side only. The number of prompt tokens to load from the prefiller. + Stops short of the trailing ``_prefill_backoff()`` tokens that the decoder + will recompute locally.""" + backoff = self._prefill_backoff() + if backoff and num_prompt_tokens > backoff: + return num_prompt_tokens - backoff return num_prompt_tokens - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: - """P-side only: drop the last prompt token so the prefiller computes - h(N-1) instead of h(N). The decoder recomputes the last token to - derive h(N) correctly. + def _truncate_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop the trailing ``_prefill_backoff()`` prompt tokens + so the prefiller stops short of what the decoder recomputes locally. + For Mamba that is the single token needed to yield h(N-1); for + multi-module MTP it is the drafter's whole lookahead window. Guarded by ``_p_side_truncated`` to avoid repeated truncation if the request is preempted and rescheduled.""" + backoff = self._prefill_backoff() params = request.kv_transfer_params if ( - params is not None + backoff + and params is not None # Guard against repeated truncation after preemption/reschedule. and not params.get("_p_side_truncated") - and request.num_prompt_tokens > 1 + and request.num_prompt_tokens > backoff ): if request.prompt_token_ids is not None: - request.prompt_token_ids.pop() + del request.prompt_token_ids[-backoff:] elif request.prompt_embeds is not None: - request.prompt_embeds = request.prompt_embeds[:-1] + request.prompt_embeds = request.prompt_embeds[:-backoff] else: return - request._all_token_ids.pop() - request.num_prompt_tokens -= 1 + del request._all_token_ids[-backoff:] + request.num_prompt_tokens -= backoff request.max_tokens = 1 params["_p_side_truncated"] = True @@ -459,13 +479,19 @@ def build_connector_meta( meta = NixlConnectorMetadata() # Loop through scheduled reqs and convert to ReqMeta. - for req_id, (req, block_ids, cached) in self._reqs_need_recv.items(): + for req_id, ( + req, + block_ids, + cached, + awaiting_kvs, + ) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None meta.add_new_req_to_recv( request_id=req_id, local_block_ids=block_ids, kv_transfer_params=req.kv_transfer_params, local_num_computed_blocks=cached, + awaiting_kvs=awaiting_kvs, ) if self.use_host_buffer: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 58ab5850eaf7..2704de589f05 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -164,6 +164,10 @@ class NixlBaseConnectorWorker: # Overridden by NixlPushConnectorWorker. _TRANSFER_MODE: str = "pull" + # Layer-name routing is supported only by NixlPushConnector for HMA KV + # caches under PP. + _supports_pp_hma = False + def _compute_desc_ids( self, block_ids: BlockIds, @@ -191,28 +195,42 @@ def _compute_desc_ids( num_blocks = dst_num_blocks if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) - num_regions = self.num_regions or len(self.block_len_per_layer) + use_layer_name_routing = bool(self._transfer_layer_group_ids) + num_desc_regions = ( + len(self._transfer_layer_group_ids) + if use_layer_name_routing + else self.num_regions or len(self.block_len_per_layer) + ) if region_num_blocks is None: - region_num_blocks = [num_blocks] * num_regions + region_num_blocks = [num_blocks] * num_desc_regions elif block_size_ratio is not None: region_num_blocks = [ int(count * block_size_ratio) for count in region_num_blocks ] - if region_num_blocks: - num_regions = len(region_num_blocks) - assert len(region_num_blocks) == num_regions + if region_num_blocks and not use_layer_name_routing: + num_desc_regions = len(region_num_blocks) + assert len(region_num_blocks) == num_desc_regions region_offsets = np.cumsum([0, *region_num_blocks[:-1]]) num_fa_descs = sum(region_num_blocks) + if use_layer_name_routing: + # Descriptors in layer order, using each layer's KV-cache group. + return np.concatenate( + [ + np.asarray(block_ids[g], dtype=np.int32) + region_offsets[k] + for k, g in enumerate(self._transfer_layer_group_ids) + ] + ) + # All-attention fast path. if num_ssm_regions == 0: if region_group_ids is None: region_group_ids = self.region_group_ids if uses_region_group_mapping is None: uses_region_group_mapping = self._uses_region_group_mapping - if not region_group_ids and num_regions == 1: + if not region_group_ids and num_desc_regions == 1: region_group_ids = [0] - assert len(region_group_ids) == num_regions + assert len(region_group_ids) == num_desc_regions if uses_region_group_mapping is None: uses_region_group_mapping = len(set(region_group_ids)) > 1 region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32) @@ -392,6 +410,103 @@ def _is_region_replicated(self, region_idx: int) -> bool: """ return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] + def _set_region_layers(self, region_layers: list[list[str]]) -> None: + layer_names = [layer_name for region in region_layers for layer_name in region] + assert len(layer_names) == len(set(layer_names)), ( + "A KV cache layer spans multiple NIXL regions" + ) + self.region_members: list[list[str]] = region_layers + if not self._requires_layer_name_routing(): + return + group_by_layer = self.kv_cache_config.transfer_group_index_by_layer + ungrouped = [name for name in layer_names if name not in group_by_layer] + assert not ungrouped, f"KV cache layers outside any local group: {ungrouped}" + self._transfer_layer_names = tuple(layer_names) + self._transfer_layer_region_indices = tuple( + region_idx + for region_idx, region in enumerate(region_layers) + for _ in region + ) + self._transfer_layer_group_ids = tuple( + group_by_layer[name] for name in layer_names + ) + + def _requires_layer_name_routing(self) -> bool: + """Whether PP push must match HMA layers by name, not region index.""" + return ( + self._supports_pp_hma + and self.pp_size > 1 + and self._is_hma_required + and not self._has_mamba + ) + + def _align_remote_regions_by_layer( + self, nixl_agent_meta: NixlAgentMetadata + ) -> None: + """Select remote regions for this PP stage in local layer order.""" + remote_region_layers = nixl_agent_meta.region_members + # Region-index routing here would silently transfer stale KV. + assert remote_region_layers, "Remote advertised no region_members" + assert ( + len(nixl_agent_meta.kv_caches_base_addr) + == len(nixl_agent_meta.block_lens) + == len(nixl_agent_meta.block_strides) + == len(remote_region_layers) + ), "Remote region metadata lengths disagree" + assert all( + values is None or len(values) == len(remote_region_layers) + for values in ( + nixl_agent_meta.region_num_blocks, + nixl_agent_meta.region_group_ids, + nixl_agent_meta.region_names, + nixl_agent_meta.region_mem_types, + ) + ), "Remote region metadata lengths disagree" + + remote_region_by_layer: dict[str, int] = {} + for region_idx, layer_names in enumerate(remote_region_layers): + for layer_name in layer_names: + assert layer_name not in remote_region_by_layer, ( + f"Remote advertised layer {layer_name!r} in multiple regions" + ) + remote_region_by_layer[layer_name] = region_idx + + missing = [ + name + for name in self._transfer_layer_names + if name not in remote_region_by_layer + ] + assert not missing, f"Remote is missing locally owned layers: {missing}" + + remote_regions = [ + remote_region_by_layer[name] for name in self._transfer_layer_names + ] + nixl_agent_meta.kv_caches_base_addr = [ + nixl_agent_meta.kv_caches_base_addr[i] for i in remote_regions + ] + nixl_agent_meta.block_lens = [ + nixl_agent_meta.block_lens[i] for i in remote_regions + ] + nixl_agent_meta.block_strides = [ + nixl_agent_meta.block_strides[i] for i in remote_regions + ] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = [ + nixl_agent_meta.region_num_blocks[i] for i in remote_regions + ] + if nixl_agent_meta.region_group_ids is not None: + nixl_agent_meta.region_group_ids = [ + nixl_agent_meta.region_group_ids[i] for i in remote_regions + ] + if nixl_agent_meta.region_mem_types is not None: + nixl_agent_meta.region_mem_types = [ + nixl_agent_meta.region_mem_types[i] for i in remote_regions + ] + if nixl_agent_meta.region_names is not None: + nixl_agent_meta.region_names = list(self._transfer_layer_names) + # One layer per region keeps a second alignment pass a no-op. + nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names] + def __init__( self, vllm_config: "VllmConfig", @@ -638,18 +753,26 @@ def __init__( # transfers into the matching sub-range of a PP=1 remote's regions. self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self._remote_region_offset = 0 - # PP push slices regions per layer (uniform count); HMA breaks that. - if self.pp_size > 1 and self._is_hma_required: - raise NotImplementedError( - "NixlPushConnector does not support pipeline_parallel_size > 1 " - "with hybrid KV cache layouts (HMA) yet." - ) - # Decode-side PP is unsupported (completions counted per consumer rank). - if vllm_config.kv_transfer_config.kv_role == "kv_consumer" and self.pp_size > 1: - raise NotImplementedError( - "NixlPushConnector consumer (decode) does not support " - "pipeline_parallel_size > 1." - ) + if self.pp_size > 1: + if self._is_hma_required and not self._supports_pp_hma: + raise NotImplementedError( + "NixlConnector (pull) does not support pipeline_parallel_size " + "> 1 with hybrid KV cache layouts (HMA); use NixlPushConnector " + "for PP + HMA." + ) + # PP push routes HMA (hybrid) attention layouts by layer name; + # Mamba/SSM hybrids are not yet supported under PP. + if self._has_mamba: + raise NotImplementedError( + "NixlPushConnector does not support pipeline_parallel_size > 1 " + "with Mamba/SSM hybrid KV cache layouts yet." + ) + # Decode-side PP is unsupported (completions counted per consumer rank). + if vllm_config.kv_transfer_config.kv_role == "kv_consumer": + raise NotImplementedError( + "NixlPushConnector consumer (decode) does not support " + "pipeline_parallel_size > 1." + ) # Keep heartbeat handshakes to a PP-sharded producer notif-only. self._hb_handshake_notif_only = False @@ -692,7 +815,7 @@ def __init__( # Uses Queue for thread-safe cross-thread coordination with the # background handshake thread, matching the _ready_requests pattern. self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue() - self._failed_inflight_recvs: set[ReqId] = set() + self._recv_failures: set[ReqId] = set() self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {} # Handshake metadata of this worker for NIXL transfers. @@ -773,6 +896,13 @@ def __init__( # within a block (BLHNC/BHLNC), where stride > block_len. self.block_stride_per_layer = list[int]() + # Region layer membership is populated for HMA layouts. + self.region_members = [] + # Local layer order, set by ``_set_region_layers`` only when + # routing by layer name. Empty means region-index routing. + self._transfer_layer_names = tuple[str, ...]() + self._transfer_layer_region_indices = tuple[int, ...]() + self._transfer_layer_group_ids = tuple[int, ...]() # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} @@ -1212,7 +1342,7 @@ def request_ready(f: Future[Any], entry=(req_id, meta)): error=e, meta=meta, ) - self._handle_failed_transfer(req_id, None) + self._failed_recv_reqs.put(req_id) fut.add_done_callback(request_ready) @@ -1306,6 +1436,11 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ): compressed_region_owners.setdefault(cache.data_ptr(), cache) + track_region_layers = ( + self._supports_pp_hma and self._is_hma_required and not self._has_mamba + ) + region_layers: list[list[str]] = [] + # K and V are packed into the content dim, so each attention layer is a # single NIXL region whose block transfers as one unit. Mamba layers instead # register separate conv/ssm sub-regions (see `_build_mamba_local`). @@ -1508,6 +1643,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self._region_is_mla.append(is_mla_region) region_mem_types.append(mem_type) + if track_region_layers: + if region_index == len(region_layers): + region_layers.append([layer_name]) + else: + region_layers[region_index].append(layer_name) + if isinstance(layer_spec, MambaSpec): if layer_spec.tp_replicated: assert self._is_csa_linear @@ -1563,7 +1704,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): if self._has_mamba: self.region_num_blocks = [self.num_blocks] * self.num_regions - if self.pp_size > 1: + self._set_region_layers(region_layers) + + if self.pp_size > 1 and not self._is_hma_required: start_layer, end_layer = self.model_config.get_layers_start_end_indices( self.vllm_config.parallel_config ) @@ -1573,7 +1716,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self._remote_region_offset = regions_per_layer * start_layer # Total local FA descriptors (boundary between FA and mamba descs). - self.num_descs = sum(self.region_num_blocks) + xfer_region_num_blocks = ( + [self.region_num_blocks[i] for i in self._transfer_layer_region_indices] + if self._transfer_layer_region_indices + else self.region_num_blocks + ) + self.num_descs = sum(xfer_region_num_blocks) self._mixed_mem_types = len(set(region_mem_types)) > 1 if self._mixed_mem_types: @@ -1597,7 +1745,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.device_kv_caches = kv_caches self.dst_num_blocks[self.engine_id] = self.num_blocks - self.dst_region_num_blocks[self.engine_id] = self.region_num_blocks + self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks self.dst_region_group_ids[self.engine_id] = self.region_group_ids self.dst_uses_region_group_mapping[self.engine_id] = ( self._uses_region_group_mapping @@ -1647,6 +1795,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_mem_types=self.region_mem_types, dcp_size=self.dcp_size, pcp_size=self.pcp_size, + region_members=self.region_members, ) # Wrap metadata in payload with hash for defensive decoding assert self.compat_hash is not None @@ -1805,12 +1954,16 @@ def _build_fa_local( base_addresses: list[int], block_size_ratio: int, ) -> np.ndarray: - """Build local FA descriptors for all layers as an Nx3 uint64 array.""" + """Build local FA descriptors as an Nx3 uint64 array.""" assert self.transfer_topo is not None assert base_addresses, "Local KV cache base addresses must not be empty." device_id = self.device_id + region_indices = self._transfer_layer_region_indices or range( + len(base_addresses) + ) parts: list[np.ndarray] = [] - for i, base_addr in enumerate(base_addresses): + for i in region_indices: + base_addr = base_addresses[i] block_len = self.block_len_per_layer[i] // block_size_ratio logical_blocks = np.repeat( np.arange(self.region_num_blocks[i], dtype=np.uint64), @@ -1834,7 +1987,7 @@ def _build_fa_remote( nixl_agent_meta: NixlAgentMetadata, block_size_ratio: int, ) -> np.ndarray: - """Build remote FA descriptors for all layers as an Nx3 uint64 array.""" + """Build remote FA descriptors as an Nx3 uint64 array.""" assert self.transfer_topo is not None assert nixl_agent_meta.kv_caches_base_addr, ( "Remote KV cache base addresses must not be empty." @@ -1851,10 +2004,14 @@ def _build_fa_remote( device_id = nixl_agent_meta.device_id block_strides = nixl_agent_meta.block_strides parts: list[np.ndarray] = [] + use_layer_name_routing = bool(self._transfer_layer_region_indices) for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - replicated = self._is_region_replicated(i) + local_region = ( + self._transfer_layer_region_indices[i] if use_layer_name_routing else i + ) + replicated = self._is_region_replicated(local_region) # Read our whole local region size from remote.. - local_block_len = self.block_len_per_layer[i] + local_block_len = self.block_len_per_layer[local_region] remote_kv_block_len = local_block_len // block_size_ratio if block_size_ratio > 1: # ..using remote kv_block_len as transfer unit @@ -2021,9 +2178,22 @@ def add_remote_agent( ) return self._remote_agents[engine_id][(0, remote_tp_rank)] + assert self.transfer_topo is not None + transfer_topo = self.transfer_topo # Number of physical regions registered locally (one per layer/tensor). num_local_regions = len(self.block_len_per_layer) - if ( + if self._transfer_layer_region_indices: + if self.block_size != nixl_agent_meta.block_size: + raise NotImplementedError( + "Attention-HMA push requires identical P/D block sizes." + ) + if remote_tp_size > transfer_topo.tp_size and not self.use_mla: + raise NotImplementedError( + "Attention-HMA push does not support decode TP greater than " + "prefill TP yet" + ) + self._align_remote_regions_by_layer(nixl_agent_meta) + elif ( self.pp_size > 1 and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions ): @@ -2038,6 +2208,10 @@ def add_remote_agent( ] nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end] nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end] + if nixl_agent_meta.region_num_blocks is not None: + nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ + start:end + ] if nixl_agent_meta.region_group_ids is not None: nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[ start:end @@ -2048,14 +2222,8 @@ def add_remote_agent( nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[ start:end ] - if nixl_agent_meta.region_num_blocks is not None: - nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[ - start:end - ] ### Register remote engine in TransferTopology (idempotent). - assert self.transfer_topo is not None - transfer_topo = self.transfer_topo physical_blocks_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block ) @@ -2349,6 +2517,9 @@ def _validate_remote_agent_handshake( # the per-rank KV head ratio rather than the raw tp_ratio, because GQA # replication caps per-rank heads at 1 when tp > total_kv_heads # (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here. + local_regions = self._transfer_layer_region_indices or range( + len(self.block_len_per_layer) + ) if self._has_mamba and self.use_mla: # Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are # kernel-granularity views of the mamba-unified page. The MLA @@ -2365,7 +2536,7 @@ def _validate_remote_agent_handshake( f"remote={nixl_agent_meta.block_lens}." ) elif not self._has_mamba: - assert len(self.block_len_per_layer) == len(nixl_agent_meta.block_lens), ( + assert len(local_regions) == len(nixl_agent_meta.block_lens), ( "Number of KV layers must match between prefill and decode" ) model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated( @@ -2374,8 +2545,9 @@ def _validate_remote_agent_handshake( total_kv_heads = self.transfer_topo.total_num_kv_heads local_heads = self.transfer_topo.local_physical_heads remote_heads = max(1, total_kv_heads // remote_tp_size) - for i, local_len in enumerate(self.block_len_per_layer): - replicated = model_replicated or self._is_region_replicated(i) + for i, region_idx in enumerate(local_regions): + local_len = self.block_len_per_layer[region_idx] + replicated = model_replicated or self._is_region_replicated(region_idx) remote_len = nixl_agent_meta.block_lens[i] if replicated: assert local_len // block_size_ratio == remote_len, ( @@ -2406,8 +2578,7 @@ def _validate_remote_agent_handshake( # TP workers that handhshake with same remote have same #blocks. assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks - # Same number of regions/~layers. - assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer) + assert len(nixl_agent_meta.kv_caches_base_addr) == len(local_regions) num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr) if nixl_agent_meta.region_num_blocks is not None: assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions @@ -2423,7 +2594,12 @@ def _validate_remote_agent_handshake( if nixl_agent_meta.region_group_ids is not None: assert len(nixl_agent_meta.region_group_ids) == num_remote_regions if nixl_agent_meta.region_names is not None: - assert nixl_agent_meta.region_names == self.region_names, ( + expected_names = ( + list(self._transfer_layer_names) + if self._transfer_layer_names + else self.region_names + ) + assert nixl_agent_meta.region_names == expected_names, ( "NIXL transfer regions must name the same model caches in the " "same order" ) @@ -2618,22 +2794,32 @@ def get_transfer_results(self) -> KVConnectorTransferResults: """ assert self.transfer_topo is not None done_sending = self._get_new_notifs() - done_recving = self._pop_done_transfers(self._recving_transfers) + done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers) + if newly_failed: + self._recv_failures.update(newly_failed) done_sending.update(self._replicated_pcp_done_sending) self._replicated_pcp_done_sending.clear() - # Drain queue of requests where handshake or transfer setup failed. - failed_recv_reqs = set[ReqId]() + # Process receive failures reported by background threads. while not self._failed_recv_reqs.empty(): try: - failed_recv_reqs.add(self._failed_recv_reqs.get_nowait()) + req_id = self._failed_recv_reqs.get_nowait() except queue.Empty: break + self._handle_failed_transfer(req_id, None, self._recv_failures) + + failed_recv_reqs: set[ReqId] = set() + if self._recv_failures: + # Keep failure bookkeeping independent of healthy request count. + for req_id in tuple(self._recv_failures): + if req_id not in self._recving_metadata: + self._recv_failures.remove(req_id) + elif req_id not in self._recving_transfers: + self._recv_failures.remove(req_id) + failed_recv_reqs.add(req_id) + done_recving.update(failed_recv_reqs) - # Add failed requests to done_recving for scheduler tracking - # (blocks are already marked invalid, scheduler will handle recompute) - done_recving.update(failed_recv_reqs) if len(done_sending) > 0 or len(done_recving) > 0: logger.debug( "Rank %s, get_finished: %s requests done sending " @@ -2654,12 +2840,17 @@ def get_transfer_results(self) -> KVConnectorTransferResults: # Skip KV sync and post-processing for failed requests if req_id in failed_recv_reqs: + self._pending_recv_notifs.pop(req_id, None) + # TODO (NickLucche) handle failed transfer for HMA. + if not self._is_hma_required: + self._invalid_block_ids.put(set(meta.local_block_ids[0])) logger.warning( "Skipping KV post-processing for failed request %s", req_id, ) continue + self._send_pending_recv_notifs(req_id) assert meta.remote is not None if self.use_host_buffer: self.sync_recved_kv_to_device(req_id, meta) @@ -2783,22 +2974,26 @@ def _handle_heartbeat(self, payload: str) -> None: new_expiry, ) - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: - """ - Pop completed xfers by checking for DONE state. + def _pop_done_transfers( + self, transfers: dict[str, list[int]] + ) -> tuple[set[str], set[str]]: + """Poll transfers, retaining handles until they can be released. + Args: - transfers: dict of req_id -> list[running_xfer] + transfers: Outstanding handles by request, updated in place. + Returns: - set of req_ids that have all done xfers + Requests with no outstanding handles and requests with observed + failures. Failed requests may still have outstanding handles. """ done_req_ids: set[str] = set() + failed_req_ids: set[str] = set() for req_id, handles in list(transfers.items()): in_progress = [] for handle in handles: try: xfer_state = self.nixl_wrapper.check_xfer_state(handle) if xfer_state == "DONE": - # Get telemetry from NIXL res = self.nixl_wrapper.get_xfer_telemetry(handle) self.xfer_stats.record_transfer(res) self.nixl_wrapper.release_xfer_handle(handle) @@ -2807,58 +3002,52 @@ def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]: else: self._log_failure( failure_type="transfer_failed", - msg="Deferring request completion until its last " - "xfer is terminal", req_id=req_id, xfer_state=xfer_state, ) - self._handle_failed_transfer(req_id, handle) + if not self._handle_failed_transfer( + req_id, handle, failed_req_ids + ): + in_progress.append(handle) except Exception as e: self._log_failure( failure_type="transfer_exception", - msg="Handle is unpollable; treating it as terminal", req_id=req_id, error=e, ) - self._handle_failed_transfer(req_id, handle) + if not self._handle_failed_transfer(req_id, handle, failed_req_ids): + in_progress.append(handle) if not in_progress: - del transfers[req_id] done_req_ids.add(req_id) - if req_id in self._failed_inflight_recvs: - self._failed_inflight_recvs.remove(req_id) - self._report_failed_recv(req_id) - else: - self._send_pending_recv_notifs(req_id) + del transfers[req_id] else: transfers[req_id] = in_progress - return done_req_ids - - def _handle_failed_transfer(self, req_id: str, handle: int | None): - """ - Handle a failed transfer by marking all (logical) blocks as invalid and - recording the failure. + return done_req_ids, failed_req_ids - Args: - req_id: The request ID. - handle: The transfer handle. - """ - # A sibling READ may still be writing these blocks. Retain metadata - # and defer invalidation until every handle is terminal. - self._pending_recv_notifs.pop(req_id, None) - if req_id in self._recving_transfers: - self._failed_inflight_recvs.add(req_id) - else: - self._report_failed_recv(req_id) - if handle is not None: + def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool: + """Release a handle, returning False if the caller must retain it.""" + try: self.nixl_wrapper.release_xfer_handle(handle) + except Exception as e: + # A status error does not guarantee that the backend stopped DMA. + self._log_failure( + failure_type="transfer_release_failed", + msg="Retaining handle and blocks until release succeeds", + req_id=req_id, + error=e, + ) + return False + return True + + def _handle_failed_transfer( + self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None + ) -> bool: + """Record a failure and release its handle, returning False to retain it.""" self.xfer_stats.record_failed_transfer() - - def _report_failed_recv(self, req_id: str) -> None: - if (meta := self._recving_metadata.get(req_id)) is not None: - if not self._is_hma_required: - self._invalid_block_ids.put(set(meta.local_block_ids[0])) - self._failed_recv_reqs.put(req_id) + if failed_req_ids is not None: + failed_req_ids.add(req_id) + return handle is None or self._try_release_xfer_handle(req_id, handle) def _send_pending_recv_notifs(self, req_id: str) -> None: """Send notifications deferred by split DRAM/VRAM reads.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index f61fdfa83512..c2b4d9dce753 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Metadata dataclasses and helpers for the NIXL connector.""" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from vllm.config import VllmConfig @@ -47,8 +47,9 @@ # 9: Add block_strides # 10: Add dense virtual transfer pages for compressed MLA caches # 11: Add per-region transfer geometry and memory types to NixlAgentMetadata +# 12: Add per-region member names for PP push # -NIXL_CONNECTOR_VERSION: int = 11 +NIXL_CONNECTOR_VERSION: int = 12 @dataclass @@ -71,6 +72,8 @@ class NixlAgentMetadata: region_mem_types: list[str] | None = None dcp_size: int = 1 pcp_size: int = 1 + # Layer names sharing each advertised region, in region order. + region_members: list[list[str]] = field(default_factory=list) @dataclass @@ -246,6 +249,9 @@ class ReqMeta: remote_block_size: int | None = None # Remote producer pipeline-parallel size (push mode, D side). pp_size: int = 1 + # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS + # and expects it in finished_recving; notify-only recvs must not be reported. + awaiting_kvs: bool = False class NixlConnectorMetadata(KVConnectorMetadata): @@ -275,6 +281,7 @@ def _add_new_req( local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], local_num_computed_blocks: tuple[int, ...] = (), + awaiting_kvs: bool = False, ) -> ReqMeta: return ReqMeta( local_block_ids=local_block_ids, @@ -285,6 +292,7 @@ def _add_new_req( remote_block_size=kv_transfer_params.get("remote_block_size"), pp_size=kv_transfer_params.get("pp_size", 1), local_num_computed_blocks=local_num_computed_blocks, + awaiting_kvs=awaiting_kvs, ) def add_new_req_to_save( @@ -303,11 +311,13 @@ def add_new_req_to_recv( local_block_ids: BlockIds, kv_transfer_params: dict[str, Any], local_num_computed_blocks: tuple[int, ...] = (), + awaiting_kvs: bool = False, ): req = self._add_new_req( local_block_ids, kv_transfer_params, local_num_computed_blocks, + awaiting_kvs, ) req.remote = RemoteMeta( block_ids=kv_transfer_params["remote_block_ids"], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 446126c328ea..9b9012a532ea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -172,6 +172,9 @@ def update_state_after_alloc( request, local_block_ids, local_num_computed_blocks, + # Parked in WAITING_FOR_REMOTE_KVS only when there is + # something to pull; a full local hit stays RUNNING. + num_external_tokens > 0, ) else: @@ -224,7 +227,7 @@ def request_finished( # To avoid stranding the prefill blocks in the prefill instance, # we must add empty block_ids to _reqs_need_recv so that our # worker side will notify and free blocks in the prefill instance. - self._reqs_need_recv[request.request_id] = (request, [], ()) + self._reqs_need_recv[request.request_id] = (request, [], (), False) params["do_remote_prefill"] = False return False, None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 59205b943a5e..04caebbd1f21 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -149,7 +149,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): engine_id, ) self.xfer_stats.record_kv_expired_req() - self._handle_failed_transfer(req_id, None) + self._handle_failed_transfer(req_id, None, self._recv_failures) return if any(len(group) > 0 for group in meta.local_block_ids): @@ -303,6 +303,7 @@ def group_ids(block_ids: BlockIds, rank: int) -> list[list[int]]: local_dram_handle=local_dram_handle, remote_xfer_side_handle=remote_xfer_side_handle, expected_consumers=plan.local_consumers, + awaiting_kvs=meta.awaiting_kvs, ): return @@ -327,6 +328,7 @@ def _read_blocks( local_dram_handle: int | None, remote_xfer_side_handle: int, expected_consumers: int, + awaiting_kvs: bool, ) -> bool: """ Post a READ point-to-point xfer request from a single local worker to @@ -387,6 +389,11 @@ def _read_blocks( remote_agent_name=agent_name, ) self.xfer_stats.record_failed_notification() + # Report even on notification failure: the KV is already local, and + # an unreported parked request would hold its blocks forever. + # Notify-only recvs must stay unreported (scheduler asserts). + if awaiting_kvs: + self._recving_transfers.setdefault(request_id, []) return True if read_spec.block_ids_by_region: @@ -486,16 +493,19 @@ def _read_blocks( self._recving_transfers[request_id].append(handle) return True except Exception as e: - # mark all (logical) blocks for this request as invalid self._log_failure( failure_type="transfer_setup_failed", req_id=request_id, - msg="Marking blocks as invalid", + msg="Deferring failure reporting until outstanding transfers finish", error=e, dst_engine_id=dst_engine_id, remote_rank=remote_rank, ) - self._handle_failed_transfer(request_id, handle) + if not self._handle_failed_transfer( + request_id, handle, self._recv_failures + ): + assert handle is not None + self._recving_transfers[request_id].append(handle) return False def _read_blocks_mixed( @@ -537,7 +547,8 @@ def _read_blocks_mixed( ) except Exception: for handle in handles: - self.nixl_wrapper.release_xfer_handle(handle) + if not self._try_release_xfer_handle(request_id, handle): + self._recving_transfers[request_id].append(handle) raise self._pending_recv_notifs.setdefault(request_id, []).append( @@ -548,7 +559,8 @@ def _read_blocks_mixed( self.nixl_wrapper.transfer(handle) except Exception: for unstarted in handles[i:]: - self.nixl_wrapper.release_xfer_handle(unstarted) + if not self._try_release_xfer_handle(request_id, unstarted): + self._recving_transfers[request_id].append(unstarted) raise self._recving_transfers[request_id].append(handle) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py index c11a63bb3414..a007e0c5a3ae 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py @@ -198,7 +198,7 @@ def update_state_after_alloc( # ReqMeta without a KeyError — the actual remote block IDs are # learned by P over the NIXL handshake at WRITE time. params["remote_block_ids"] = () - self._reqs_need_recv[request.request_id] = (request, local_block_ids, ()) + self._reqs_need_recv[request.request_id] = (request, local_block_ids, (), False) # Mark as processed so a re-entry (e.g. preemption + reschedule) # doesn't re-stage the registration. @@ -240,7 +240,7 @@ def request_finished( # recv so the worker emits a notif that lets P free them. # Seed remote_block_ids so add_new_req_to_recv won't KeyError. params["remote_block_ids"] = () - self._reqs_need_recv[request.request_id] = (request, [], ()) + self._reqs_need_recv[request.request_id] = (request, [], (), False) params["do_remote_prefill"] = False return False, None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 603321260de2..01c956a5251d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -83,6 +83,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): # Distinguishes push from pull in the NIXL compatibility hash. _TRANSFER_MODE: str = "push" + _supports_pp_hma = True + def __init__( self, vllm_config: "VllmConfig", @@ -341,7 +343,7 @@ def _on_handshake( self._log_failure( failure_type="push_reg_handshake_failed", req_id=rid, error=e ) - self._handle_failed_transfer(rid, None) + self._failed_recv_reqs.put(rid) return # Re-queue for the writer to send now that the handshake is done. self._reg_send_inbox.put((rid, rd)) @@ -363,7 +365,7 @@ def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None: engine_id, req_id, ) - self._handle_failed_transfer(req_id, None) + self._failed_recv_reqs.put(req_id) return for rank, agent_name in agents.items(): try: @@ -530,7 +532,7 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): local_region_groups = self.region_group_ids remote_region_groups = self.dst_region_group_ids[engine_id] groups_differ = local_region_groups != remote_region_groups - if groups_differ: + if groups_differ and not self._transfer_layer_group_ids: raise NotImplementedError( "NixlPushConnector does not support different producer and " "consumer cache-group layouts" @@ -722,9 +724,8 @@ def _xfer_blocks( # don't have a ``_recving_metadata`` entry to invalidate, so # we just release the handle and let the engine reschedule # via the lease / watchdog. - if handle is not None: - self.nixl_wrapper.release_xfer_handle(handle) - self.xfer_stats.record_failed_transfer() + if not self._handle_failed_transfer(request_id, handle): + return handle return None # --- Notification handling on engine main thread ------------------ # @@ -803,7 +804,16 @@ def get_transfer_results(self) -> KVConnectorTransferResults: # ``_pop_done_transfers`` mutates ``_sending_transfers``; the # writer thread also appends to it, so guard the pop. with self._sending_transfers_lock: - done_pushing = self._pop_done_transfers(self._sending_transfers) + done_pushing, failed_pushing = self._pop_done_transfers( + self._sending_transfers + ) + # A failed send must never be reported as done: its blocks + # are freed via the lease / watchdog instead. + done_pushing = { + req_id + for req_id in done_pushing - failed_pushing + if req_id in self._recving_metadata + } for req_id in done_pushing: self._reqs_to_send.pop(req_id, None) self._reqs_to_process.discard(req_id) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 86e351093917..5cc66e4bf0ff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time from collections.abc import Iterable, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from itertools import chain, islice from typing import Any, NamedTuple @@ -114,7 +114,10 @@ def load_window_size_in_chunks(self, num_tokens: int) -> int | None: window = max( window, cdiv( - self.kv_cache_spec.sliding_window - 1 + right_padding, + self.kv_cache_spec.sliding_window + - 1 + + self.kv_cache_spec.extra_retained_tokens + + right_padding, self.tokens_per_chunk, ), ) @@ -126,7 +129,13 @@ def get_sliding_window_size_in_chunks( ) -> int | None: if isinstance(kv_cache_spec, SlidingWindowSpec): assert kv_cache_spec.sliding_window > 0 - return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk) + return max( + cdiv(kv_cache_spec.sliding_window, tokens_per_chunk), + cdiv( + kv_cache_spec.sliding_window - 1 + kv_cache_spec.extra_retained_tokens, + tokens_per_chunk, + ), + ) if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec): # Attention never reaches back past one chunk @@ -354,6 +363,8 @@ class RequestOffloadState: req_context: ReqContext offloading_context: RequestOffloadingContext group_states: tuple[RequestGroupState, ...] = field(init=False) + # upper bound on tokens to load for this request; None means no cap + max_load_tokens: int | None = None # upper bound on tokens to offload for this request; None means no cap max_offload_tokens: int | None = None # number of hits in the GPU cache @@ -376,6 +387,20 @@ def __post_init__(self) -> None: ) params = self.req.kv_transfer_params + # NOTE: This field is experimental and subject to change in the future. + raw = params.get("max_load_tokens") if params else None + if type(raw) is int and raw >= 0: + self.max_load_tokens = raw + logger.debug( + "Request %s: max_load_tokens set to %d", + self.req.request_id, + raw, + ) + elif raw is not None: + logger.warning( + "max_load_tokens must be a non-negative int, got %r; ignoring", raw + ) + # NOTE: This field is experimental and subject to change in the future. raw = params.get("max_offload_tokens") if params else None if type(raw) is int and raw >= 0: @@ -685,6 +710,7 @@ def _sliding_window_lookup( The first run may need a larger window for a partial rightmost chunk. Returns 0 on miss, None if the backend deferred a lookup.""" defer_lookup = False + pending_in_window = False consecutive_hits = 0 required_window = initial_window_size or sliding_window_size for idx in range(len(keys) - 1, -1, -1): @@ -695,7 +721,7 @@ def _sliding_window_lookup( # Block is in cache, just not readable yet — counts # as hit for the consecutive streak. Don't break: # keep scanning to let manager kick off async lookups. - defer_lookup = True + pending_in_window = True consecutive_hits += 1 case LookupResult.RETRY: # Block location uncertain — does not count as hit. @@ -703,13 +729,18 @@ def _sliding_window_lookup( # async lookups. defer_lookup = True consecutive_hits = 0 + pending_in_window = False required_window = sliding_window_size case LookupResult.MISS: consecutive_hits = 0 + # This gap rules out the incomplete window to its right. + pending_in_window = False required_window = sliding_window_size if consecutive_hits == required_window: - return idx + required_window if not defer_lookup else None - return consecutive_hits if not defer_lookup else None + return ( + None if defer_lookup or pending_in_window else idx + required_window + ) + return None if defer_lookup or pending_in_window else consecutive_hits def _touch(self, req_status: RequestOffloadState): for group_config, group_state in zip( @@ -761,6 +792,11 @@ def _lookup_complete_chunks( max_hit_size_tokens = min( max_hit_size_tokens, num_computed_tokens + max_num_new_tokens ) + if req_status.max_load_tokens is not None: + max_hit_size_tokens = min( + max_hit_size_tokens, + num_computed_tokens + req_status.max_load_tokens, + ) if self._sliding_window_groups: # the last prompt token has to be recomputed to get the logprobs # for sliding window attention, we must reduce by 1 to make sure @@ -804,6 +840,10 @@ def _lookup_complete_chunks( max_hit_size_tokens = min( max_hit_size_tokens, len(offload_keys) * tokens_per_chunk ) + if req_status.max_load_tokens is not None: + max_hit_size_tokens = round_down( + max_hit_size_tokens, tokens_per_chunk + ) if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk: # We can only load less than a chunk, so skip. return 0 @@ -959,6 +999,11 @@ def _lookup( max_boundary = min(req_status.req.num_prompt_tokens - 1, block_end - 1) if max_num_new_tokens is not None: max_boundary = min(max_boundary, local_tokens + max_num_new_tokens) + if req_status.max_load_tokens is not None: + max_boundary = min( + max_boundary, + local_tokens + req_status.max_load_tokens, + ) max_boundary = round_down(max_boundary, tokens_per_hash) if max_boundary <= complete_boundary: return complete_hit @@ -1425,11 +1470,24 @@ def _reachable_store_block_mask( ) -> list[bool] | None: """Build the block mask for a range of candidate offload chunks.""" blocks_per_chunk = self.config.blocks_per_chunk + kv_cache_spec = group_config.kv_cache_spec + if isinstance(kv_cache_spec, SlidingWindowSpec) and ( + kv_cache_spec.extra_retained_tokens + ): + # Offload restores all allocated history, including MTP re-prefill + # tokens. Widen only the store mask, not the model's attention window. + kv_cache_spec = replace( + kv_cache_spec, + sliding_window=( + kv_cache_spec.sliding_window + kv_cache_spec.extra_retained_tokens + ), + extra_retained_tokens=0, + ) return group_config.manager_cls.reachable_block_mask( start_block=start_chunk_idx * blocks_per_chunk, end_block=end_chunk_idx * blocks_per_chunk, alignment_tokens=self.config.alignment_tokens, - kv_cache_spec=group_config.kv_cache_spec, + kv_cache_spec=kv_cache_spec, use_eagle=group_config.is_eagle_group, retention_interval=self.config.retention_interval, reachable_boundaries=reachable_boundaries, diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py index 3879bee0ce29..c063b3ea7426 100644 --- a/vllm/distributed/nixl_utils.py +++ b/vllm/distributed/nixl_utils.py @@ -91,9 +91,26 @@ def is_nixl_available() -> bool: return pkg in sys.modules or importlib.util.find_spec(pkg) is not None +def alias_nixl_for_ray() -> None: + """Let Ray's ``nixl._api`` import resolve to the ROCm implementation.""" + if not current_platform.is_rocm() or "nixl._api" in sys.modules: + return + + package_name = _get_nixl_package_name() + if not is_nixl_available(): + raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}") + + _maybe_set_ucx_rcache_limit() + package = importlib.import_module(package_name) + api = importlib.import_module(f"{package_name}._api") + sys.modules.setdefault("nixl", package) + sys.modules.setdefault("nixl._api", api) + + __all__ = [ "NixlWrapper", "nixl_agent_config", "nixlXferTelemetry", "is_nixl_available", + "alias_nixl_for_ray", ] diff --git a/vllm/distributed/weight_transfer/sharded_rdt_common.py b/vllm/distributed/weight_transfer/sharded_rdt_common.py index a5a7ed923369..70110517d3fc 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_common.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_common.py @@ -77,6 +77,16 @@ def check_ray_rdt_version() -> None: ) +def register_nixl_memory(tensor: torch.Tensor) -> None: + """Register a tensor with Ray using the platform's NIXL package.""" + from vllm.distributed.nixl_utils import alias_nixl_for_ray + + alias_nixl_for_ray() + from ray.experimental import register_nixl_memory as ray_register_nixl_memory + + ray_register_nixl_memory(tensor) + + def assign_producer_indices( num_producers: int, num_consumers: int, consumer_idx: int ) -> list[int]: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_engine.py b/vllm/distributed/weight_transfer/sharded_rdt_engine.py index 2b8187007b8f..96dd53f0fdb9 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_engine.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_engine.py @@ -66,6 +66,7 @@ RdtRouter, buffer_alloc_bytes, check_ray_rdt_version, + register_nixl_memory, ) from vllm.distributed.weight_transfer.sharded_rdt_fake import ( BakeSink, @@ -588,8 +589,6 @@ def _preregister_at_init(self) -> None: plan = self._cached_plan if plan is None or not plan.chunks: return - from ray.experimental import register_nixl_memory - # (a) consumer receive buffers — one per ring slot, at the largest chunk. max_pack = max(c.pack_bytes for c in plan.chunks) alloc = buffer_alloc_bytes(max_pack, self._buffer_presize) @@ -740,7 +739,7 @@ def _issue_pull(self, chunk: "_Chunk", slot: int) -> "_PendingPull": any time after the metadata push. See the doc's "slot generation handshake". """ - from ray.experimental import register_nixl_memory, set_target_for_ref + from ray.experimental import set_target_for_ref if self._slot_read_done: with self._slot_cv: diff --git a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py index 60d0e60fb02d..1be17c4fa088 100644 --- a/vllm/distributed/weight_transfer/sharded_rdt_trainer.py +++ b/vllm/distributed/weight_transfer/sharded_rdt_trainer.py @@ -47,6 +47,7 @@ ALLOWED_OPS, buffer_alloc_bytes, check_ray_rdt_version, + register_nixl_memory, ) from vllm.logger import init_logger @@ -334,8 +335,6 @@ def warmup_nixl(self) -> None: doc). The warmup buffer stays registered so the agent's CUDA-HMEM path stays initialized. """ - from ray.experimental import register_nixl_memory - self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda") with self._reg_lock: register_nixl_memory(self._nixl_warmup_buf) @@ -510,8 +509,6 @@ def _new_serve_buffer(self, nbytes: int) -> torch.Tensor: """Allocate + NIXL-register one serve slot: the single allocation seam, so registration cannot be skipped on either of the two paths that make buffers (the init-time reservation and the serve-path backstop).""" - from ray.experimental import register_nixl_memory - t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device) with self._reg_lock: register_nixl_memory(t) diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index e43b273e5e96..6a8492f4c4b7 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -18,6 +18,7 @@ from vllm.entrypoints.launchers.dp_supervisor import run_dp_supervisor from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger +from vllm.reasoning import ReasoningParserManager from vllm.usage.usage_lib import UsageContext from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.network_utils import get_tcp_uri @@ -179,6 +180,9 @@ def run_headless(args: argparse.Namespace): if args.api_server_count > 1: raise ValueError("api_server_count can't be set in headless mode") + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) + # Create the EngineConfig. engine_args = vllm.AsyncEngineArgs.from_cli_args(args) usage_context = UsageContext.OPENAI_API_SERVER diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index b919435f65a1..73b9ebe06b8f 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -293,6 +293,18 @@ def _beam_search_step( allowed_sets[i] = set(entry[1]) for (start, end), instance in zip(instance_start_and_end, instances_batch): + instance_output = output[start:end] + if any( + result is not None and result.outputs[0].finish_reason == "abort" + for result in instance_output + ): + for beam, result in zip(all_beams[start:end], instance_output): + if result is not None: + beam.finish_reason = "abort" + instance.completed.append(beam) + instance.beams = [] + continue + instance_new_beams = [] for i in range(start, end): current_beam = all_beams[i] @@ -302,8 +314,6 @@ def _beam_search_step( continue if result.outputs[0].logprobs is not None: - # if logprobs is None, the sequence completed - # due to max-model-len or abortion. logprobs = result.outputs[0].logprobs[0] allowed = allowed_sets[i] for token_id, logprob_obj in logprobs.items(): diff --git a/vllm/entrypoints/generate/beam_search/online.py b/vllm/entrypoints/generate/beam_search/online.py index ac1a585e8f62..da02b974c0a5 100644 --- a/vllm/entrypoints/generate/beam_search/online.py +++ b/vllm/entrypoints/generate/beam_search/online.py @@ -99,11 +99,7 @@ async def beam_search( output = [x[0] for x in await asyncio.gather(*tasks)] - candidates = [] - # Iterate through all beam inference results - for i, result in enumerate(output): - current_beam = all_beams[i] - + for result in output: # check for error finish reason and abort beam search if result.outputs[0].finish_reason == "error": # yield error output and terminate beam search @@ -126,6 +122,15 @@ async def beam_search( ) return + if any(result.outputs[0].finish_reason == "abort" for result in output): + for beam in all_beams: + beam.finish_reason = "abort" + break + + candidates = [] + for i, result in enumerate(output): + current_beam = all_beams[i] + if result.outputs[0].logprobs is not None: logprobs = result.outputs[0].logprobs[0] for token_id, logprob_obj in logprobs.items(): diff --git a/vllm/entrypoints/generate/beam_search/utils.py b/vllm/entrypoints/generate/beam_search/utils.py index 47f69cbf7abe..61f97963b6f3 100644 --- a/vllm/entrypoints/generate/beam_search/utils.py +++ b/vllm/entrypoints/generate/beam_search/utils.py @@ -150,7 +150,8 @@ def get_beam_search_score( if tokens[-1] == eos_token_id: seq_len -= 1 - return cumulative_logprob / (seq_len**length_penalty) + # An aborted beam may contain only an EOS prompt token. + return cumulative_logprob / (max(seq_len, 1) ** length_penalty) def create_sort_beams_key_function(eos_token_id: int, length_penalty: float): diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index e30195a7b1f9..c1366734f7c1 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -171,6 +171,8 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin): + padding: Literal["max_length", "do_not_pad"] | None + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: default_max_total_tokens = model_config.max_model_len max_total_tokens: int | None = default_max_total_tokens @@ -184,7 +186,7 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: max_embed_len = pooler_config.max_embed_len or default_max_total_tokens max_output_tokens = default_max_total_tokens - max_embed_len - return self._build_pooling_tok_params( + tok_params = self._build_pooling_tok_params( model_config, add_special_tokens=self.add_special_tokens, max_total_tokens=max_total_tokens, @@ -192,6 +194,20 @@ def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: max_output_tokens_param="max_model_len - max_embed_len", ) + if ( + self.padding == "max_length" + and pooler_config is not None + and pooler_config.enable_chunked_processing + ): + max_padding_len = pooler_config.max_embed_len + if max_padding_len is None: + max_padding_len = default_max_total_tokens + return tok_params.with_kwargs( + pad_prompt_tokens=min(default_max_total_tokens, max_padding_len) + ) + + return tok_params + class CompletionRequestMixin(OpenAIBaseModel): # --8<-- [start:completion-params] diff --git a/vllm/entrypoints/pooling/scoring/api_router.py b/vllm/entrypoints/pooling/scoring/api_router.py index 90be32ba248a..8e99e9a1feca 100644 --- a/vllm/entrypoints/pooling/scoring/api_router.py +++ b/vllm/entrypoints/pooling/scoring/api_router.py @@ -57,8 +57,6 @@ async def create_score(request: ScoreRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation -@load_aware_call async def create_score_v1(request: ScoreRequest, raw_request: Request): logger.warning( "To indicate that Score API is not part of standard OpenAI API, we " @@ -91,7 +89,6 @@ async def do_rerank(request: RerankRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation async def do_rerank_v1(request: RerankRequest, raw_request: Request): logger.warning_once( "To indicate that the rerank API is not part of the standard OpenAI" @@ -110,6 +107,5 @@ async def do_rerank_v1(request: RerankRequest, raw_request: Request): HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, }, ) -@with_cancellation async def do_rerank_v2(request: RerankRequest, raw_request: Request): return await do_rerank(request, raw_request) diff --git a/vllm/envs.py b/vllm/envs.py index 6158f86f896b..bc4724e2e4a7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -137,7 +137,7 @@ VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True VLLM_ROCM_AITER_MOE_DISPATCH_POLICY: int = 0 - VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4: bool = False + VLLM_ROCM_USE_AITER_MOE_SITUV2: bool = False VLLM_ROCM_USE_AITER_RMSNORM: bool = True VLLM_ROCM_USE_AITER_MLA: bool = True VLLM_ROCM_AITER_MLA_ASM_PADDING: Literal["auto", "gluon", "asm"] = "auto" @@ -1251,12 +1251,18 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ROCM_USE_AITER_MOE": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_MOE", "True").lower() in ("true", "1") ), - # Route K3 SiTU MXFP4 MoE through the a8w4 (fp8 activation) gate/up- - # interleaved flydsl kernels instead of the default a16w4 separated path. - # This is the only flag users need: vLLM picks the kernels by passing - # gate_mode to AITER and sets the AITER-side workaround env at init. - "VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() in ("true", "1") + # Route K3 SiTU MXFP4 MoE through the FlyDSL SiTUv2 path (a4w4 fp4 + # activations, separated gate/up layout) instead of default a16w4. vLLM + # sets AITER_SITUV2_A4W4 at init when this flag is on and clears any + # legacy AITER_SITUV2_A8W4 override (AITER checks A8W4 first). + # VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 is a deprecated alias for existing + # recipes; it does not select a8w4 kernels. + # Needs AITER >= v0.1.20 (ROCm/aiter#4463) for the a4w4 dispatch flag + # and tuned kimik3_a4w4_*_fmoe.csv rows; otherwise FlyDSL uses heuristics. + "VLLM_ROCM_USE_AITER_MOE_SITUV2": lambda: ( + os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2", "0").lower() in ("true", "1") + or os.getenv("VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4", "0").lower() + in ("true", "1") ), # MoE sorting dispatch policy for AITER fused MoE kernels. # 0 = auto (default): single-pass for small batches, multi-pass @@ -2170,6 +2176,10 @@ def _resolve_rust_cli_path() -> str | None: # Each op additionally checks its own shape / dtype constraints and falls # back to the eager path when they do not hold. "VLLM_ENABLE_HPC_OPS": lambda: bool(int(os.getenv("VLLM_ENABLE_HPC_OPS", "0"))), + # Whether to skip version suffix when building package + "VLLM_SKIP_VERSION_SUFFIX": lambda: bool( + int(os.getenv("VLLM_SKIP_VERSION_SUFFIX", "0")) + ), } @@ -2338,6 +2348,7 @@ def compile_factors() -> dict[str, object]: "LOCAL_RANK", "CUDA_VISIBLE_DEVICES", "NO_COLOR", + "VLLM_SKIP_VERSION_SUFFIX", } from vllm.config.utils import normalize_value diff --git a/vllm/forward_context.py b/vllm/forward_context.py index f57fc2c950bd..8d9eb5d7db3b 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -209,6 +209,15 @@ def is_forward_context_available() -> bool: return _forward_context is not None +def in_piecewise_cudagraph() -> bool: + """Whether the current forward runs in piecewise cudagraph mode (graph + segments separated by eager breaks), at capture or replay time.""" + return ( + is_forward_context_available() + and get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE + ) + + def create_forward_context( attn_metadata: Any, vllm_config: VllmConfig, diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py new file mode 100644 index 000000000000..839c581583a7 --- /dev/null +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer. + +The candidate-source indexer publishes the top candidate blocks (in units of +``candidate_block_size`` compressed positions). Instead of scoring every KV +position and masking the dense logits down to those blocks, this path scores +only the candidate tokens with DeepGEMM's sparse MQA logits kernels +(``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs +DeepSelect's top-k directly on the bf16 logits they produce, and remaps the +selected columns back to request-local positions. + +Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed +(granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The +FP8 indexer cache stores one fp32 scale per quant block instead. + +The functions here take caller-owned output buffers (``sparse_indices``, +``end``, ``col_indices``) so the hot path allocates nothing; the indexer +metadata builder owns them. +""" + +import importlib.util + +import torch + +from vllm.model_executor.layers.indexer_topk import ( + deep_select_topk, + get_deep_select_stride_requirement, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import ( + fp8_fp4_paged_sparse_mqa_logits, + fp8_fp4_sparse_mqa_logits, + get_paged_sparse_mqa_logits_metadata, + get_sparse_mqa_logits_metadata, +) + +# Sparse-block sizes (in tokens) supported by the DeepGEMM kernels. +SPARSE_BLOCK_KV_CHOICES = (16, 8) + + +def has_deep_select() -> bool: + """Whether the DeepSelect top-k extension is built and this GPU runs it.""" + return ( + importlib.util.find_spec("vllm._deepselect_C") is not None + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ) + + +def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None: + """Raise unless DeepSelect accepts the sparse logits layout. + + The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]`` + logits and the top-k writes int32 ``[rows, topk_tokens]`` indices; + DeepSelect needs both row strides aligned (1024B / 32B on SM100). + """ + in_align, out_align = get_deep_select_stride_requirement() + logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize + if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0: + raise ValueError( + f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 " + f"columns, {logits_row_bytes} B) to be {in_align}B-aligned." + ) + if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0: + raise ValueError( + f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 " + f"output row, got topk={topk_tokens}." + ) + + +_INT32_BIG = tl.constexpr(2**31 - 1) + + +@triton.jit +def _expand_candidates_kernel( + cand_ptr, + cand_stride, + ks_ptr, + ks_stride, + ke_ptr, + ke_stride, + si_ptr, + si_stride, + end_ptr, + end_stride, + K: tl.constexpr, + S: tl.constexpr, + RATIO: tl.constexpr, + SBK: tl.constexpr, +): + """One program per row: expand + filter + sort candidate blocks. + + Each row's valid sparse blocks come out sorted ascending (no dedup pass: + production candidates are unique, and the kernel tolerates repeats), + padded by repeating the last valid block — DeepGEMM's padding convention. + ``end`` counts valid sparse-token columns; valid columns form a prefix + because block ids are sorted. + """ + row = tl.program_id(0) + ks = tl.load(ks_ptr + row * ks_stride) + ke = tl.load(ke_ptr + row * ke_stride) + cols = tl.arange(0, K) + cand = tl.load(cand_ptr + row * cand_stride + cols) # K is pow2, no mask + if RATIO == 2: + cand = tl.interleave(cand * RATIO, cand * RATIO + 1) + + # Keep only blocks inside the row's KV range [0, ke - ks); the lower + # bound is free since candidates are non-negative. + kv_len = tl.maximum(ke - ks, 0) + max_block = (kv_len + SBK - 1) // SBK + valid = (cand >= 0) & (cand < max_block) + vals = tl.where(valid, cand, _INT32_BIG) + vals = tl.sort(vals) + + # Anchor at the row's ks rounded down to a sparse-block boundary + # (identity for the paged path, where ks == 0). + base = ks // SBK + si = tl.where(vals != _INT32_BIG, vals, 0) + base + n_valid = tl.sum(valid.to(tl.int32)) + last_rel = tl.max(tl.where(valid, cand, -1)) + pad_val = tl.where(n_valid > 0, last_rel + base, base) + tl.store( + si_ptr + row * si_stride + tl.arange(0, S), + tl.where(tl.arange(0, S) < n_valid, si, pad_val), + ) + + # Valid sparse-token column count; only the last block can be partial. + # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK, + # which is exactly last_rel * SBK + ks. + last_start = last_rel * SBK + ks + fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK) + end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0) + tl.store(end_ptr + row * end_stride, end) + + +@triton.jit +def _sparse_topk_remap_kernel( + col_ptr, + col_stride, + si_ptr, + si_stride, + ks_ptr, + ks_stride, + out_ptr, + out_stride, + k, + SBK: tl.constexpr, + K_POW2: tl.constexpr, +): + """Remap sparse top-k columns to request-local positions. + + Column ``j * SBK + o`` scores the token at + ``(si[row, j] - ks // SBK) * SBK + o``; -1 columns stay -1. + """ + row = tl.program_id(0) + ks_base = tl.load(ks_ptr + row * ks_stride) // SBK + cols = tl.arange(0, K_POW2) + c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1) + valid = c >= 0 + cc = tl.where(valid, c, 0) + blocks = tl.load(si_ptr + row * si_stride + cc // SBK) + pos = (blocks - ks_base) * SBK + cc % SBK + tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k) + + +def pick_sparse_block_kv(candidate_block_size: int) -> int: + """Largest sparse-block size (tokens) dividing the candidate block size.""" + for sbk in SPARSE_BLOCK_KV_CHOICES: + if candidate_block_size % sbk == 0: + return sbk + raise ValueError( + f"candidate_block_size={candidate_block_size} is not a multiple of " + f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}." + ) + + +def candidate_blocks_to_sparse_indices( + candidate_blocks: torch.Tensor, + row_ks: torch.Tensor, + row_ke: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + out: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand request-local candidate blocks into DeepGEMM sparse block ids. + + Args: + candidate_blocks: [rows, K] int32 request-local candidate block ids + (-1 padded), in units of ``candidate_block_size`` positions. K + must be a power of two (the in-kernel sort); rows may be strided. + row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same + (packed-workspace) coordinates the sparse kernel iterates over. + Pass zeros for the paged path, whose blocks are context-relative. + candidate_block_size: Positions per candidate block. + sparse_block_kv: Positions per sparse block (8 or 16). + out: Optional ``(sparse_indices, end)`` buffers to write into, + shaped [rows, K * ratio] and [rows], both int32. + + Returns: + sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM, + sorted ascending per row with the valid prefix first (unique in + production, where candidates are unique) and the remaining slots + repeating the last valid block (the kernel's padding convention). + With the unaligned-ks variant, block ``i`` starts at + ``i * sparse_block_kv + ks % sparse_block_kv``. + end: [rows] int32 count of valid sparse-token columns; valid columns + always form a prefix because the block ids are sorted. + """ + rows, num_candidates = candidate_blocks.shape + ratio = candidate_block_size // sparse_block_kv + assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), ( + candidate_block_size, + sparse_block_kv, + ) + assert num_candidates & (num_candidates - 1) == 0, num_candidates + num_sparse = num_candidates * ratio + if out is None: + device = candidate_blocks.device + out = ( + torch.empty(rows, num_sparse, dtype=torch.int32, device=device), + torch.empty(rows, dtype=torch.int32, device=device), + ) + sparse_indices, end = out + assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,) + assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1 + for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end): + assert t.dtype == torch.int32, t.dtype + + _expand_candidates_kernel[(rows,)]( + candidate_blocks, + candidate_blocks.stride(0), + row_ks, + row_ks.stride(0), + row_ke, + row_ke.stride(0), + sparse_indices, + sparse_indices.stride(0), + end, + end.stride(0), + K=num_candidates, + S=num_sparse, + RATIO=ratio, + SBK=sparse_block_kv, + num_warps=16, + ) + return sparse_indices, end + + +def sparse_topk_remap( + logits: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + row_ks: torch.Tensor, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + col_indices: torch.Tensor, +) -> None: + """Row top-k over sparse logits, remapped to request-local positions. + + Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at + request-local position + ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``. + + DeepSelect consumes the bf16 logits as they come out of the sparse + kernels; the valid columns form a prefix (block ids are sorted), so the + per-row ``end`` bounds drive it directly. Slots beyond a row's valid + count come back as -1. + + Args: + col_indices: [rows, topk_tokens] int32 scratch for the sparse-column + top-k result (row stride 32B-aligned for DeepSelect). + """ + rows, width = logits.shape + assert logits.dtype == torch.bfloat16 and width >= topk_tokens + assert end.dtype == torch.int32 and end.shape == (rows,) + cols = col_indices[:rows] + deep_select_topk(logits, topk_tokens, end=end, output_idx=cols) + _sparse_topk_remap_kernel[(rows,)]( + cols, + cols.stride(0), + sparse_indices, + sparse_indices.stride(0), + row_ks, + row_ks.stride(0), + topk_indices, + topk_indices.stride(0), + topk_tokens, + SBK=sparse_block_kv, + K_POW2=triton.next_power_of_2(topk_tokens), + ) + + +def sparse_mqa_logits_prefill_chunk( + q: torch.Tensor, + q_scale: torch.Tensor, + k_quant: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace). + + Args: + q: [rows, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, H] int32 packed UE8M0 Q scales. + k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8). + k_scale: [total_kv] int32 packed UE8M0 K scales. + weights: [rows, H] bf16 per-head weights; the sparse kernels take + bf16 and do not fold the Q scale in. + cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the + packed workspace. + candidate_blocks: [rows, K] int32 request-local candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + sparse_indices/end/col_indices: Caller-owned scratch, see + `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`. + kernel_metadata: DeepGEMM schedule from a previous call with the same + candidates and bounds (i.e. another indexer layer in the same + step). When given, the candidate expansion is skipped and + ``sparse_indices``/``end`` are assumed to be up to date. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + assert weights.dtype == torch.bfloat16, weights.dtype + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + cu_seqlen_ks, + cu_seqlen_ke, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_sparse_mqa_logits_metadata( + cu_seqlen_ks, + cu_seqlen_ke, + k_quant.shape[0], + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_sparse_mqa_logits( + (q, q_scale), + (k_quant, k_scale), + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + cu_seqlen_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata + + +def sparse_mqa_logits_paged_decode( + q: torch.Tensor, + q_scale: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + row_indices: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + row_ks: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for paged decode rows. + + Every query is one row (spec-decode rows are flattened by the metadata + builder), so ``block_table`` is already per row. + + Args: + q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. + kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache + view; the page stride must be 512B-aligned. + weights: [rows, H] bf16 per-head weights. + context_lens: [rows] int32 per-row (compressed) context lengths. + block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. + row_indices: [rows] int32 row -> request map (pairing only affects + scheduling). + candidate_blocks: [rows, K] int32 candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + row_ks: [rows] int32 zeros (paged blocks are context-relative). + sparse_indices/end/col_indices: Caller-owned scratch. + kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + rows = q.shape[0] + page_kv = kv_cache.shape[1] + assert weights.dtype == torch.bfloat16, weights.dtype + assert kv_cache.stride(0) % 512 == 0, ( + "paged cache page stride must be 512B-aligned: " + f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" + ) + assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv) + assert block_table.shape[0] == rows and block_table.stride(-1) == 1 + assert context_lens.shape == (rows,) and context_lens.is_contiguous() + + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + row_ks, + context_lens, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_paged_sparse_mqa_logits_metadata( + context_lens, + block_table, + row_indices, + page_kv, + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_paged_sparse_mqa_logits( + (q, q_scale), + kv_cache, + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + row_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 9714f92dffd4..29739afae3f2 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -4,9 +4,9 @@ Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16); activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling -matrix cores. Falls back (via the kernel selector) to the BF16 -``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with -``K % 128 != 0``. +matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple +of that is dequantized to BF16 once at load and served by a plain linear +instead, since ``can_implement`` does not filter on K. """ import torch @@ -194,6 +194,9 @@ def _select_cfg(M, N, K): return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3) +_DOT_SCALED_K_ALIGN = 128 + + class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" @@ -219,6 +222,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: N, K = weight.shape scale_k = K // MXFP8_BLOCK_SIZE weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() + if K % _DOT_SCALED_K_ALIGN != 0: + weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale) layer.weight = Parameter(weight.contiguous(), requires_grad=False) layer.weight_scale = Parameter(weight_scale, requires_grad=False) @@ -235,10 +240,13 @@ def apply_weights( ) out_shape = (*x.shape[:-1], layer.weight.shape[0]) x2d = x.reshape(-1, x.shape[-1]) - if x2d.shape[-1] % 128 == 0: + if layer.weight.element_size() >= 2: + out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype)) + elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0: out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale) else: - # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise. + # process_weights_after_loading dequantizes these weights once, so + # this runs only if it did not; dot_scaled would be invalid here. w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale) out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype) out = out.reshape(out_shape) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 874c43f8cd44..1e673bd2da14 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -269,11 +269,11 @@ LayerNameType, _encode_layer_name, _resolve_layer_name, + async_tensor_h2d, direct_register_custom_op, get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, - np_to_pinned_tensor, ) from vllm.v1.attention.backend import ( AttentionBackend, @@ -1980,11 +1980,6 @@ def aligned_split_len(start: int, remaining: int, available: int) -> int: return plans -def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor: - """Pinned int32 CPU tensor backing one concatenated per-chunk field.""" - return np_to_pinned_tensor(np.asarray(values, dtype=np.int32)) - - def align_mla_chunked_context_workspace_size( vllm_config: VllmConfig, workspace_size: int, @@ -2161,24 +2156,28 @@ def padded_rows(rows: int) -> int: token_offset += num_tokens local_token_offset += num_local_tokens - seq_lens_cpu = _flat_int32(seq_lens_flat) - starts_and_context_lens = _flat_int32(starts_flat + context_lens).to( - device, non_blocking=True + seq_lens_cpu = torch.tensor( + seq_lens_flat, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY + ) + starts_and_context_lens = async_tensor_h2d( + starts_flat + context_lens, device=device, dtype=torch.int32 ) starts = starts_and_context_lens[: len(starts_flat)] context_lens_gpu = starts_and_context_lens[len(starts_flat) :] - cu_seq_lens = _flat_int32(cu_seq_lens_flat).to(device, non_blocking=True) - cu_seqlens_q = _flat_int32(cu_seqlens_q_flat).to(device, non_blocking=True) - token_to_seq = _flat_int32(np.concatenate(token_to_seq_parts)).to( - device, non_blocking=True + cu_seq_lens = async_tensor_h2d(cu_seq_lens_flat, device=device, dtype=torch.int32) + cu_seqlens_q = async_tensor_h2d(cu_seqlens_q_flat, device=device, dtype=torch.int32) + token_to_seq = async_tensor_h2d( + np.concatenate(token_to_seq_parts), device=device, dtype=torch.int32 ) if use_dcp: - padded_local_cu_seq_lens = _flat_int32(padded_local_cu_seq_lens_flat).to( - device, non_blocking=True + padded_local_cu_seq_lens = async_tensor_h2d( + padded_local_cu_seq_lens_flat, device=device, dtype=torch.int32 + ) + padded_local_token_to_seq = async_tensor_h2d( + np.concatenate(padded_local_token_to_seq_parts), + device=device, + dtype=torch.int32, ) - padded_local_token_to_seq = _flat_int32( - np.concatenate(padded_local_token_to_seq_parts) - ).to(device, non_blocking=True) chunks: list[MLACommonPrefillMetadata.ContextChunk] = [] for index, (plan, layout) in enumerate(zip(plans, layouts)): diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index cb4225c92a0f..62df395f692c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -385,14 +385,10 @@ def rocm_aiter_fused_experts( gate_mode = "" if activation == MoEActivation.SITU: - # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up- - # interleaved (_gui_) fp8 flydsl kernels; default a16w4 SiTU stays - # separated. - gate_mode = ( - GateMode.INTERLEAVE.value - if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() - else GateMode.SEPARATED.value - ) + # SiTUv2 flydsl (VLLM_ROCM_USE_AITER_MOE_SITUV2=1) uses a4w4 + # fp4 activations with separated gate/up weights (AITER #4463); + # default a16w4 SiTU also stays separated. + gate_mode = GateMode.SEPARATED.value elif quant_config.use_mxfp4_w4a16: gate_mode = GateMode.INTERLEAVE.value elif activation_interleave is not None: diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3b415acd815d..1c38e4378b26 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1572,12 +1572,6 @@ def convert_weight_to_mxfp4_moe_kernel_format( if w2_bias is not None: w2_bias = w2_bias.data.to(torch.float32) - import os - - # TODO: Remove this once AITER is fixed - # Necessary for AITER side from crashing - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" - if activation == MoEActivation.SITU: from aiter.utility.fp4_utils import e8m0_shuffle @@ -1585,9 +1579,8 @@ def convert_weight_to_mxfp4_moe_kernel_format( fp4_dtype = torch.float4_e2m1fn_x2 e8m0_dtype = torch.float8_e8m0fnu - # a8w4 uses gate/up-interleaved flydsl kernels; - # default a16w4 keeps the separated layout. - guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() + # SiTUv2 flydsl uses separated gate/up layout (a4w4). + guinterleave = False w13 = rocm_aiter_ops.shuffle_weight_a16w4( w13_weight.data.view(fp4_dtype), 16, guinterleave ) @@ -1606,6 +1599,14 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2.is_shuffled = True return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias) + import os + + # Interleaved a16w4 only (DeepSeekV4 etc.). AITER uses this bound to + # pick bf16 vs fp8 activations when gate_mode is INTERLEAVE. SiTUv2 + # a4w4 is separated and selects q_dtype_a independently, so the bound + # is unused on that path. + os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" + from aiter.ops.shuffle import shuffle_scale as _shuf_s from aiter.ops.shuffle import shuffle_weight as _shuf_w diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index d834c6eff73c..b3a146274d30 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -21,6 +21,7 @@ from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, + dbo_enabled, ) @@ -89,8 +90,10 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): use cudagraphs anyway - Provides expert_tokens_meta for efficient batched expert kernels - Both modes use async_with_compute_stream=False (synchronous from - caller's perspective). The ElasticBuffer handles comm internally. + Dispatch always uses async_with_compute_stream=False. finalize_async + issues the combine with async_with_compute_stream=True (except under + DBO) so the modular kernel can overlap the shared-expert FFN with the + combine a2a; the returned receiver joins via a device-side event wait. """ @staticmethod @@ -465,15 +468,31 @@ def _finalize( f"got {fused_expert_output.dtype}" ) + # DBO drives its own hook/receiver schedule; keep the combine + # synchronous there (the receiver then only performs the copy). + combine_async = do_async and not dbo_enabled() combined_x, _, event = self.buffer.combine( x=fused_expert_output, handle=handle, topk_weights=None, - async_with_compute_stream=False, + async_with_compute_stream=combine_async, + allocate_on_comm_stream=combine_async, ) - output.copy_(combined_x, non_blocking=True) - return None + if do_async: + # The combine ran on DeepEP's comm stream; the modular kernel + # issues the shared-expert FFN before calling the receiver, which + # joins via a device-side cudaStreamWaitEvent (no host sync, so + # this is safe inside a captured region). + def _receiver(): + if event.event is not None: + event.current_stream_wait() + output.copy_(combined_x, non_blocking=True) + + return _receiver + else: + output.copy_(combined_x, non_blocking=True) + return None def finalize_async( self, @@ -484,16 +503,17 @@ def finalize_async( apply_router_weight_on_input: bool, weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> Callable: - self._finalize( + receiver = self._finalize( output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input, weight_and_reduce_impl, - False, + True, ) - return lambda: None + assert receiver is not None + return receiver def finalize( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 2868690ebeb1..6e63563ab861 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -457,9 +457,18 @@ def _maybe_reduce_routed_output_before_transform( Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm. TP partial routed outputs must be summed in latent space before such transforms are applied. + + A transform that commutes with the TP sum is exempt: if + ``sum_r T(x_r) == T(sum_r x_r)``, applying the transform to the local + partial output and letting the existing late all-reduce sum the + combined result is equivalent, and costs one collective instead of two. + Such a transform opts out by setting ``reduce_commutative = True``. + The default is False, so transforms that do not declare themselves + keep being reduced early. """ if ( self.routed_output_transform is not None + and not getattr(self.routed_output_transform, "reduce_commutative", False) and not self.moe_config.is_sequence_parallel and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1) and not fused_output_is_reduced diff --git a/vllm/model_executor/layers/hpc/hpc_ihc.py b/vllm/model_executor/layers/hpc/hpc_ihc.py index 1a8681aa8ead..9a8065a052e3 100644 --- a/vllm/model_executor/layers/hpc/hpc_ihc.py +++ b/vllm/model_executor/layers/hpc/hpc_ihc.py @@ -20,9 +20,8 @@ - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint keeps them out of fp8 quantization via modules_to_not_convert) -The fused modules own no parameters: they hold references to the weights of the -eager layer they replace, bound in ``process_weights_after_loading`` (invoked -generically for every ``HpcModule`` by the model loader). +The fused modules own no parameters: their forwards read the weights of the +eager layer they replace directly via the captured owner reference. NOTE: The reference implementation also offers a cross-layer post+pre fusion (``HpcIHCPostPre``) that folds one segment's post into the next segment's pre. @@ -147,41 +146,26 @@ def __init__( object.__setattr__(self, "_fallback_op", fallback_op) # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module). object.__setattr__(self, "_norm_owner", norm_owner) - self.weight = None - self.hc_scale = None - self.hc_base = None - self.rms_weight = None - self.rms_eps = 0.0 @classmethod def support(cls, hc_mult: int, hidden_size: int) -> bool: return _ihc_supported(hc_mult, hidden_size) - def process_weights_after_loading(self, model=None): - # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d]. - # It stays float32 -- the HPC kernel consumes fp32 weights directly, so - # there is no conversion (and no precision loss) here. - owner = self._fallback_op - self.weight = owner.hc_fn.weight - self.hc_scale = owner.hc_scale - self.hc_base = owner.hc_base - if self._norm_owner is not None: - self.rms_weight = self._norm_owner.weight - self.rms_eps = self._norm_owner.variance_epsilon - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: import hpc + owner = self._fallback_op + norm_owner = self._norm_owner return hpc.fuse_ihc_pre( x, - self.weight, - self.hc_scale, - self.hc_base, + owner.hc_fn.weight, + owner.hc_scale, + owner.hc_base, self.norm_eps, self.hc_eps, self.magnitude, - self.rms_weight, - self.rms_eps, + norm_owner.weight if norm_owner is not None else None, + norm_owner.variance_epsilon if norm_owner is not None else 0.0, ) @@ -233,29 +217,20 @@ def __init__( self.norm_eps = norm_eps # See HpcIHCPre: keep the owner out of the module tree. object.__setattr__(self, "_fallback_op", fallback_op) - self.weight = None - self.hc_scale = None - self.hc_base = None @classmethod def support(cls, hc_mult: int, hidden_size: int) -> bool: return _ihc_supported(hc_mult, hidden_size) - def process_weights_after_loading(self, model=None): - # hc_head_fn weight is [hc_mult, hc_mult * d], float32. - owner = self._fallback_op - self.weight = owner.hc_head_fn.weight - self.hc_scale = owner.hc_head_scale - self.hc_base = owner.hc_head_base - def forward(self, x: torch.Tensor) -> torch.Tensor: import hpc + owner = self._fallback_op return hpc.fuse_ihc_head( x, - self.weight, - self.hc_scale, - self.hc_base, + owner.hc_head_fn.weight, + owner.hc_head_scale, + owner.hc_head_base, self.norm_eps, self.hc_eps, ) diff --git a/vllm/model_executor/layers/hpc/hpc_module.py b/vllm/model_executor/layers/hpc/hpc_module.py index e1efad832eb9..c3a0dab676a1 100644 --- a/vllm/model_executor/layers/hpc/hpc_module.py +++ b/vllm/model_executor/layers/hpc/hpc_module.py @@ -11,8 +11,5 @@ def __init__(self): def support(cls, *args, **kwargs): return True - def process_weights_after_loading(self, model): - pass - def forward(self, *args, **kwargs): pass diff --git a/vllm/model_executor/layers/hpc/rope_norm.py b/vllm/model_executor/layers/hpc/rope_norm.py index c1503c754c83..02f06ed89cbd 100644 --- a/vllm/model_executor/layers/hpc/rope_norm.py +++ b/vllm/model_executor/layers/hpc/rope_norm.py @@ -90,13 +90,13 @@ class HpcRopeNorm(CustomOp, HpcModule): """HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant). Registered as a sub-module in model layers (e.g. HunYuanAttention). - Norm weights are extracted from fallback norm modules via - process_weights_after_loading() after all weights are loaded. + The QK-Norm weights are read directly from the fallback norm modules + (built in float32 when HPC is active), eliminating the derived copy. forward() is dispatched by CustomOp framework: - In compiled mode: forward_cuda() calls torch.ops.vllm.hpc_rope_norm_forward as a splitting point — internal Python control flow is opaque - to torch.compile and not captured by CUDA Graph. + to torch.compile (its kernel launches can still be CUDA-graph captured). - In eager/native mode: forward_native() falls back to forward_cuda(). """ @@ -124,8 +124,6 @@ def __init__( self.num_kv_heads = num_kv_heads self.head_dim = head_dim - self.use_qk_norm = use_qk_norm - self.q_size = num_heads * head_dim self.kv_size = num_kv_heads * head_dim @@ -139,31 +137,18 @@ def __init__( self.head_per_group = num_heads // num_kv_heads - # Pre-allocate norm weight tensors as Parameters so they are tracked by - # CuMemAllocator (for sleep/wake_up) and have stable addresses for CUDA - # Graph replay. process_weights_after_loading() updates them inplace via - # copy_() so refit does not invalidate captured graph tensor pointers. - # Shape is [head_dim] to match the HPC kernel's q/k_norm_weight layout. - if use_qk_norm and fallback_qnorm is not None: - self.qnorm_weight: torch.nn.Parameter | None = torch.nn.Parameter( - torch.empty(head_dim, dtype=torch.float32), - requires_grad=False, - ) - else: - self.qnorm_weight = None - if use_qk_norm and fallback_knorm is not None: - self.knorm_weight: torch.nn.Parameter | None = torch.nn.Parameter( - torch.empty(head_dim, dtype=torch.float32), - requires_grad=False, - ) - else: - self.knorm_weight = None - self.use_fp8 = "fp8" in kv_cache_dtype # The RMSNorm/RoPE ordering is model dependent (e.g. HunYuan V3 applies # QK-Norm before RoPE -> NORM_THEN_ROPE), so it is supplied by the # caller. When QK-Norm is disabled the policy is forced to NONE. self.qk_norm_policy = qk_norm_policy if use_qk_norm else QkNormPolicy.NONE + if self.qk_norm_policy != QkNormPolicy.NONE and ( + fallback_qnorm is None or fallback_knorm is None + ): + raise ValueError( + "QK-Norm is enabled but fallback_qnorm/fallback_knorm were not " + "supplied; the fused kernel reads their weights directly." + ) # Register layer_name + add self to the global instance registry so the # module-level custom op (hpc_rope_norm_forward) can route back here. @@ -215,20 +200,6 @@ def support( logger.info_once("enable hpc rope_norm") return True - def process_weights_after_loading(self, model: torch.nn.Module = None) -> None: - """Copy norm weights (float32) from fallback norm modules inplace. - - Uses copy_() to preserve tensor addresses for CUDA Graph / refit - compatibility. Called by the model's load_weights() after all weights - are loaded (and generically from the model loader for DummyModelLoader - / sleep-wake_up reload paths). - """ - if self.use_qk_norm: - if self.fallback_qnorm is not None and self.qnorm_weight is not None: - self.qnorm_weight.data.copy_(self.fallback_qnorm.weight.data.float()) - if self.fallback_knorm is not None and self.knorm_weight is not None: - self.knorm_weight.data.copy_(self.fallback_knorm.weight.data.float()) - def register_layer_name(self, layer_name: str) -> None: """Register layer_name and add self to the global registry. @@ -430,12 +401,14 @@ def _forward_impl( k_scale = attn_layer._k_scale.reshape(1) v_scale = attn_layer._v_scale.reshape(1) - q_norm_weight = ( - self.qnorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None - ) - k_norm_weight = ( - self.knorm_weight if self.qk_norm_policy != QkNormPolicy.NONE else None - ) + if self.qk_norm_policy != QkNormPolicy.NONE: + assert self.fallback_qnorm is not None + assert self.fallback_knorm is not None + q_norm_weight = self.fallback_qnorm.weight + k_norm_weight = self.fallback_knorm.weight + else: + q_norm_weight = None + k_norm_weight = None # --- Prefill --- if num_prefill_reqs > 0: diff --git a/vllm/model_executor/layers/mamba/ops/replayssm_config.py b/vllm/model_executor/layers/mamba/ops/replayssm_config.py index aff4fc1c3f2d..3cee44ba5015 100644 --- a/vllm/model_executor/layers/mamba/ops/replayssm_config.py +++ b/vllm/model_executor/layers/mamba/ops/replayssm_config.py @@ -8,7 +8,6 @@ """ import functools -from contextlib import contextmanager from vllm.platforms import current_platform from vllm.triton_utils import triton @@ -26,20 +25,6 @@ def _is_blackwell() -> bool: _overrides: dict[str, tuple] = {} -@contextmanager -def override_replayssm_config(kernel: str, config: tuple): - """Pin ``kernel``'s launch config for the duration of the context.""" - prev = _overrides.get(kernel) - _overrides[kernel] = config - try: - yield - finally: - if prev is None: - _overrides.pop(kernel, None) - else: - _overrides[kernel] = prev - - def _dstate_tile(dstate: int, tile: int) -> int: return max(16, min(tile, triton.next_power_of_2(dstate))) diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py new file mode 100644 index 000000000000..78db1b3df5a4 --- /dev/null +++ b/vllm/model_executor/layers/sparse_mqa_indexer.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse Attention Indexer that scores only the candidate blocks. + +DeepSeek V4.1 two-level selection: the candidate-source indexer publishes +the top candidate blocks and later indexers pick their top-k inside them. +`SparseAttnIndexer` does that by computing dense logits over the whole +context and masking; this layer instead calls DeepGEMM's sparse MQA-logits +kernels on the candidate blocks only, so the work is O(candidate blocks) +instead of O(context). It requires the `DeepseekV41SparseIndexerBackend` +metadata (see ``AttentionConfig.indexer_sparse_logits``). +""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + sparse_mqa_logits_paged_decode, + sparse_mqa_logits_prefill_chunk, +) +from vllm.model_executor.layers.sparse_attn_indexer import ( + _gather_workspace_shapes, + kv_cache_as_quant_view, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, +) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerMetadata, +) +from vllm.v1.worker.workspace import current_workspace_manager + + +def _prefill_k_workspaces( + total_seq_lens: int, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with + the dense indexer layers of the same model.""" + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True + ) + k_quant, k_scale = current_workspace_manager().get_simultaneous( + values_spec, scales_spec + ) + return k_quant, k_scale + + +def _gather_prefill_chunk_k( + kv_cache: torch.Tensor, + k_quant_full: torch.Tensor, + k_scale_full: torch.Tensor, + chunk: DeepseekV32IndexerPrefillChunkMetadata, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather one prefill chunk's paged K into the packed workspace.""" + assert chunk.local_cu_seq_lens is not None + k_quant = k_quant_full[: chunk.max_local_total_seq_lens] + k_scale = k_scale_full[: chunk.max_local_total_seq_lens] + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.local_cu_seq_lens, + ) + return k_quant, k_scale + + +class SparseMQAIndexer(nn.Module): + """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels. + + Only valid for indexer layers that read candidate blocks with the MXFP4 + indexer cache on SM100. The K cache is written by the model before this + runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the + attention layer can call either. + """ + + weights_dtype = torch.bfloat16 + """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant + kernel writes it directly so no cast runs per step.""" + + def __init__( + self, + k_cache, + topk_tokens: int, + head_dim: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + ): + super().__init__() + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_deep_gemm_sparse_mqa() + ): + raise ValueError( + "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8." + ) + if not has_deep_select(): + raise ValueError( + "SparseMQAIndexer requires the DeepSelect top-k extension " + "(vllm._deepselect_C)." + ) + self.k_cache = k_cache + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + + def _reserve_workspaces(self, device: torch.device) -> None: + """Profiling run: claim the K-gather workspace and the peak sparse + logits allocation so the memory estimate covers them.""" + _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) + + def forward( + self, + hidden_states: torch.Tensor, + q_quant: tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ) -> torch.Tensor: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + self._reserve_workspaces(hidden_states.device) + return self.topk_indices_buffer + + metadata = attn_metadata[self.k_cache.prefix] + assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), ( + "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata" + ) + assert k is None, "the model writes the indexer K cache" + q_values, q_scale = q_quant + kv_cache = self.k_cache.kv_cache + topk_tokens = self.topk_tokens + topk_indices_buffer = self.topk_indices_buffer + topk_indices_buffer[: hidden_states.shape[0]] = -1 + sparse_block_kv = metadata.sparse_block_kv + + if metadata.num_prefills > 0: + prefill = metadata.prefill + sparse_chunks = metadata.sparse_prefill + assert prefill is not None and sparse_chunks is not None + k_quant_full, k_scale_full = _prefill_k_workspaces( + self.max_total_seq_len, self.head_dim + ) + for chunk, rows in zip(prefill.chunks, sparse_chunks): + k_quant, k_scale = _gather_prefill_chunk_k( + kv_cache, k_quant_full, k_scale_full, chunk + ) + if chunk.local_total_seq_lens == 0: + continue # buffer already holds -1 + start, end = chunk.token_start, chunk.token_end + rows.kernel_metadata = sparse_mqa_logits_prefill_chunk( + q_values[start:end].view(torch.int8), + q_scale[start:end], + k_quant.view(torch.int8), + k_scale.view(torch.int32).squeeze(-1), + weights[start:end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + self.candidate_blocks[start:end], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[start:end, :topk_tokens], + sparse_indices=rows.sparse_indices, + end=rows.end, + col_indices=rows.col_indices, + kernel_metadata=rows.kernel_metadata, + ) + + if metadata.num_decodes > 0: + decode_rows = metadata.sparse_decode + assert decode_rows is not None and decode_rows.block_table is not None + assert decode_rows.row_indices is not None + num_rows = decode_rows.row_ke.shape[0] + if num_rows > 0: + decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode( + q_values[:num_rows].view(torch.int8).unsqueeze(1), + q_scale[:num_rows].unsqueeze(1), + kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True), + weights[:num_rows], + decode_rows.row_ke, + decode_rows.block_table, + decode_rows.row_indices, + self.candidate_blocks[:num_rows], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[:num_rows, :topk_tokens], + row_ks=decode_rows.row_ks, + sparse_indices=decode_rows.sparse_indices, + end=decode_rows.end, + col_indices=decode_rows.col_indices, + kernel_metadata=decode_rows.kernel_metadata, + ) + return topk_indices_buffer diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index e12910865ec2..11e6be5cf8a1 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -15,7 +15,6 @@ from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention import is_deferred_attention_layer -from vllm.model_executor.layers.hpc import HpcModule from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -154,15 +153,6 @@ def process_weights_after_loading( with device_loading_context(module, target_device): module.process_weights_after_loading(model_config.dtype) - # Process HPC modules (HpcRopeNorm, etc.) that rely on - # process_weights_after_loading being called from the model's - # load_weights(). When using DummyModelLoader (e.g. profiling or - # sleep/wake_up reload), the model's load_weights() is not called, so we - # must handle HPC modules here generically. - for _, module in model.named_modules(): - if isinstance(module, HpcModule): - module.process_weights_after_loading(model) - # Model-level post-load hook, after the per-layer quant finalize. if hasattr(model, "process_weights_after_loading"): model.process_weights_after_loading() diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index b19fb9627237..f0af5a284a4d 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -93,11 +93,14 @@ def export_entries( ``named_parameters``/``named_buffers`` are iterated with ``remove_duplicate=False`` so tied weights (e.g. ``lm_head.weight`` sharing storage with ``embed_tokens.weight``) are not silently dropped. Each unique - tensor is exported once; every additional name that refers to the same + tensor is exported once per call; every additional name that refers to the same tensor object is recorded in the returned alias map so the client can re-establish the shared identity instead of allocating uninitialized memory for it. + CUDA reduction arguments must be exported separately for each consumer so + that PyTorch registers a reference for each IPC mapping's lifetime. + Returns: A ``(entries, aliases)`` pair where ``entries`` maps a canonical name to its ``TensorEntry`` and ``aliases`` maps each duplicate name to its @@ -167,8 +170,6 @@ def __init__( self.distributed_init_method = distributed_init_method self.socket_dir = socket_dir self.model: torch.nn.Module | None = None - self.entries: dict[str, TensorEntry] = {} - self.aliases: dict[str, str] = {} # Fingerprint before loading: process_weights_after_loading may # mutate hf_config.quantization_config. self.cache_config = WeightCacheKey.from_model_config( @@ -190,17 +191,11 @@ def load_model(self) -> None: with set_current_vllm_config(self.vllm_config): ensure_model_parallel_initialized(tp_size, 1) self.model = get_daemon_model(self.vllm_config) - self._export_entries() logger.info( - "Weight cache daemon rank %d cached %d tensors", + "Weight cache daemon rank %d loaded model", self.tp_rank, - len(self.entries), ) - def _export_entries(self) -> None: - assert self.model is not None - self.entries, self.aliases = export_entries(self.model) - def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None: """Serve requests until terminated. @@ -295,28 +290,28 @@ def _handle_get_state(self, conn: socket.socket, request: dict) -> None: logger.warning("WeightCacheKey mismatch on fields: %s", mismatched) send_msg(conn, {"status": "mismatch", "fields": mismatched}) return - if not self.entries: + if self.model is None: send_msg(conn, {"status": "error", "message": "Weights were released"}) return + gpu_uuid = get_current_device_uuid() + entries, aliases = export_entries(self.model) send_msg( conn, { "status": "ok", - "entries": self.entries, - "aliases": self.aliases, - "gpu_uuid": get_current_device_uuid(), + "entries": entries, + "aliases": aliases, + "gpu_uuid": gpu_uuid, }, ) logger.info_once( "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine", self.tp_rank, - len(self.entries), - len(self.aliases), + len(entries), + len(aliases), ) def _handle_release(self, conn: socket.socket) -> None: - self.entries.clear() - self.aliases.clear() self.model = None torch.accelerator.empty_cache() logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank) diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 80c41501e6e9..48d019cd357b 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -305,10 +305,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None def embed_input_ids( self, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 603d68410df6..c90d78bd6dd1 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -53,7 +53,7 @@ from vllm.v1.outputs import LogprobsTensors from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.attn_utils import build_attn_metadata -from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores @@ -1159,7 +1159,7 @@ def _finish_prefills( return # Move the slot indices across once, up front: indexing a device # tensor with a numpy array copies them over synchronously each time. - ps_gpu = async_copy_to_gpu( + ps_gpu = async_tensor_h2d( ps.astype(np.int64), device=states.is_encoder_phase.device ) states.init_canvas(ps_gpu) @@ -1266,7 +1266,7 @@ def __call__( # was truncated near max_model_len, in which case the scheduler gave us # fewer than CL logits for that request. valid_canvas_len_np = per_req_nlogits_np[per_req_nlogits_np > 0] - valid_canvas_len = async_copy_to_gpu( + valid_canvas_len = async_tensor_h2d( valid_canvas_len_np.astype(np.int64), device=device ) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 2dddc72b218d..2fe2752475a0 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -24,6 +24,7 @@ from torch import nn from transformers import AutoModel, BatchFeature from transformers.models.gemma4 import ( + Gemma4AudioFeatureExtractor, Gemma4Config, Gemma4Processor, Gemma4VisionConfig, @@ -456,6 +457,28 @@ def get_video_repl( # --------------------------------------------------------------------------- +def _dummy_audio_num_samples(processor: Any) -> int: + """Length in samples of a worst-case dummy audio item. + + ``fft_length`` is set in ``Gemma4AudioFeatureExtractor.__init__`` and is not + defined by ``Gemma4UnifiedAudioFeatureExtractor``, which the encoder-free + Gemma 4 variant uses. Reading it unconditionally therefore raises + ``AttributeError`` during memory profiling on that variant, which is reached + whenever audio is the modality with the largest per-item token count. + + For any other extractor, fall back to the audio budget the processor itself + declares -- the same pair of values used to warn about over-long audio + elsewhere in this file. + """ + feature_extractor = processor.feature_extractor + if isinstance(feature_extractor, Gemma4AudioFeatureExtractor): + return feature_extractor.fft_length + + sampling_rate = feature_extractor.sampling_rate + max_duration_s = processor.audio_seq_length * processor.audio_ms_per_token / 1000.0 + return int(max_duration_s * sampling_rate) + + class Gemma4DummyInputsBuilder(BaseDummyInputsBuilder[Gemma4ProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: num_images = mm_counts.get("image", 0) @@ -508,7 +531,7 @@ def get_dummy_mm_data( } if num_audios > 0: - audio_len = processor.feature_extractor.fft_length + audio_len = _dummy_audio_num_samples(processor) data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index 0b54c26f47a4..248ae04747ae 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -300,8 +300,11 @@ def __init__( prefix=f"{prefix}.attn", ) if self.use_qk_norm: - self.q_norm = RMSNorm(self.head_dim, rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, rms_norm_eps) + # The HPC fused kernel reads norm weights directly in float32; + # the fallback path keeps the model default dtype. + norm_dtype = torch.float32 if rope_support else None + self.q_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype) + self.k_norm = RMSNorm(self.head_dim, rms_norm_eps, dtype=norm_dtype) # HPC fused RoPE + QK-Norm + KV-Cache-Write (+ optional FP8 Q quant). # HunYuan V3 applies QK-Norm *before* RoPE, so NORM_THEN_ROPE. diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index c3d956a946d1..68e21c9dc284 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -335,10 +335,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None self.use_parallel_drafting = speculative_config.parallel_drafting diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 2edf56871dbf..6f5e073bd623 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -23,6 +23,7 @@ from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.logger import init_logger +from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.layers.activation import ReLUSquaredActivation from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -30,6 +31,7 @@ HasInnerState, IsHybrid, MultiModalEmbeddings, + SupportsLoRA, SupportsMultiModal, SupportsMultiModalPruning, ) @@ -897,11 +899,22 @@ def get_dummy_mm_data( dummy_inputs=NanoNemotronVLDummyInputsBuilder, ) class NemotronH_Nano_VL_V2( - nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning + nn.Module, + HasInnerState, + IsHybrid, + SupportsMultiModal, + SupportsMultiModalPruning, + SupportsLoRA, ): requires_sequential_video_encoding = True """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet""" + # LoRA covers the language model only + is_non_gated_moe = NemotronHForCausalLM.is_non_gated_moe + packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping + embedding_modules = NemotronHForCausalLM.embedding_modules + lora_skip_prefixes = NemotronHForCausalLM.lora_skip_prefixes + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "language_model.backbone": "language_model.model", @@ -1331,7 +1344,19 @@ def _create_final_video_embeddings( # Create final video embeddings, merging text embeddings for indicator # tokens with video embeddings - text_embeddings = self.get_language_model().embed_input_ids(repl_token_ids) + + # LoRA support - + # These replacement tokens are produced inside the encoder, so they are + # absent from the request-token batch that the LoRA adapter index + # mapping is built from, and there can be more + # of them than max_num_batched_tokens. + # Embed them with the base weights - + # the LoRA delta is undefined for tokens with no mapping entry. + embed_tokens = self.get_language_model().model.embed_tokens + if isinstance(embed_tokens, BaseLayerWithLoRA): + embed_tokens = embed_tokens.base_layer + text_embeddings = embed_tokens(repl_token_ids) + final_video_embeddings = _merge_multimodal_embeddings( inputs_embeds=text_embeddings, multimodal_embeddings=video_embeddings, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index c1bc4c6823f7..ea8a28ebd9d7 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -49,6 +49,7 @@ QKVParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 @@ -204,6 +205,17 @@ def __init__( disable_tp=self.is_sequence_parallel, prefix=f"{prefix}.fc2_latent_proj", ) + # A bias-free, unquantized linear commutes with the TP sum + # (sum_r W x_r == W sum_r x_r), so one reduce after the transform + # suffices. Test the layer, not the model-wide `quant_config`: ModelOpt + # excludes the latent projections, so quantized checkpoints still get + # an UnquantizedLinearMethod here. + self.fc2_latent_proj.reduce_commutative = ( + not config.mlp_bias + and isinstance( + self.fc2_latent_proj.quant_method, UnquantizedLinearMethod + ) + ) else: self.fc1_latent_proj = None self.fc2_latent_proj = None diff --git a/vllm/model_executor/models/qwen3_eagle3.py b/vllm/model_executor/models/qwen3_eagle3.py index 39d0173b3ce9..745c77a97730 100644 --- a/vllm/model_executor/models/qwen3_eagle3.py +++ b/vllm/model_executor/models/qwen3_eagle3.py @@ -309,10 +309,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor( self.config.draft_vocab_size, scale=logit_scale ) - self.draft_id_to_target_id = nn.Parameter( - torch.zeros(self.config.draft_vocab_size, dtype=torch.long), - requires_grad=False, - ) + target_vocab_size = vllm_config.model_config.get_vocab_size() + if self.config.draft_vocab_size != target_vocab_size: + self.draft_id_to_target_id = nn.Parameter( + torch.zeros(self.config.draft_vocab_size, dtype=torch.long), + requires_grad=False, + ) + else: + self.draft_id_to_target_id = None self.use_parallel_drafting = vllm_config.speculative_config.parallel_drafting diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 495650eb1989..7ea3ed131919 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -378,7 +378,7 @@ def __init__( self.use_fused_qk_norm_rope_gate = ( self.attn_output_gate and getattr(self.rotary_emb, "is_neox_style", False) - and current_platform.is_cuda() + and (current_platform.is_cuda() or current_platform.is_xpu()) and supports_dtype and (text_only or supports_mrope) ) diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index cf6e7235f9ef..669186edbb1b 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -309,12 +309,22 @@ def __call__( ) +def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]: + """Weights dtypes the model's indexer layers ask for: fp32 for the dense + scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer + (`SparseMQAIndexer`) is enabled.""" + if vllm_config.attention_config.indexer_sparse_logits: + return (torch.float32, torch.bfloat16) + return (torch.float32,) + + class FusedIndexerQRopeMxFp4TritonKernel( VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: dtype: torch.dtype + weights_out_dtype: torch.dtype num_heads: int index_q_half_rot_dim: int index_q_head_dim: int @@ -434,12 +444,14 @@ def dispatch( # type: ignore[override] self, *, dtype: torch.dtype, + weights_out_dtype: torch.dtype, num_heads: int, head_dim: int, rope_dim: int, ) -> CompileKey: return self.CompileKey( dtype=dtype, + weights_out_dtype=weights_out_dtype, num_heads=num_heads, index_q_half_rot_dim=rope_dim // 2, index_q_head_dim=head_dim, @@ -456,6 +468,7 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: return self._trace_dispatch(self.dispatch)( dtype=vllm_config.model_config.dtype, + weights_out_dtype=_indexer_weights_out_dtypes(vllm_config), num_heads=num_heads, head_dim=head_dim, rope_dim=rope_dim, @@ -491,7 +504,7 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: shape=(1, compile_key.num_heads, scale_head_dim), ), index_weights_out=TritonWarmupTensor( - torch.float32, + compile_key.weights_out_dtype, shape=(1, compile_key.num_heads), ), ) @@ -540,12 +553,17 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, + weights_out_dtype: torch.dtype = torch.float32, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, ]: """Fused RoPE + quantize Q for the sparse indexer. + ``weights_out_dtype`` is the dtype the downstream scoring kernel takes: + fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse + MQA-logits kernels (CUDA MXFP4 path only). + Weight-fold semantics (important — the two paths differ): FP8 path (use_fp4=False, default): @@ -577,7 +595,7 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype) if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( @@ -613,6 +631,7 @@ def fused_indexer_q_rope_quant( weights_out=index_weights_out, ) elif current_platform.is_xpu(): + assert weights_out_dtype == torch.float32, weights_out_dtype torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( index_q, positions, @@ -647,6 +666,7 @@ def fused_indexer_q_rope_quant( index_q_scale.view(torch.int32).squeeze(-1), ), index_weights_out + assert weights_out_dtype == torch.float32, weights_out_dtype fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 diff --git a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py index fcb12740f27f..5920ed290cad 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py @@ -23,10 +23,11 @@ class DequantGatherKCacheKernel( VllmCuTeDSLJitKernel["DequantGatherKCacheKernel.CompileKey"] ): - # Hard-coded for DSv4. + # DSv4 geometry. The record format (how much of the row is quantized and + # how wide a scale group is) comes from the compile key: V4 quantizes the + # 448 NoPE dims in groups of 64 and keeps a bf16 RoPE tail, V4.1 quantizes + # all 512 dims in groups of 32 (MXFP8) and has no tail. head_dim = 512 - fp8_dim = 448 - group_size = 64 # 1 scale per 64 elems num_warps = 4 tb_size = num_warps * 32 num_stages = 4 @@ -35,17 +36,20 @@ class DequantGatherKCacheKernel( class CompileKey: block_size: int has_gather_lens: bool + fp8_dim: int = 448 + group_size: int = 64 # dims per scale @staticmethod def kernel(compile_key: CompileKey) -> Any: head_dim = DequantGatherKCacheKernel.head_dim - fp8_dim = DequantGatherKCacheKernel.fp8_dim - group_size = DequantGatherKCacheKernel.group_size + fp8_dim = compile_key.fp8_dim + group_size = compile_key.group_size num_warps = DequantGatherKCacheKernel.num_warps tb_size = DequantGatherKCacheKernel.tb_size num_stages = DequantGatherKCacheKernel.num_stages bf16_dim = head_dim - fp8_dim data_dim = fp8_dim + bf16_dim * 2 + scale_bytes = head_dim // group_size block_size = compile_key.block_size @cute.jit @@ -62,8 +66,8 @@ def load_g2s( ): op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL) cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128) - cp8_atom = cute.make_copy_atom( - cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64 + cp_scale_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=scale_bytes * 8 ) page_id = block_table[req_id, pos // block_size] block_offset = pos % block_size @@ -88,7 +92,7 @@ def load_g2s( ) elif idx == cutlass.const_expr(data_dim // 16): cute.copy( - cp8_atom, + cp_scale_atom, k_scale[page_id, block_offset, None], s_kscale[None, stage_id], ) @@ -119,8 +123,8 @@ def device_kernel( )[None, warp_id, None] s_kscale = smem.allocate_tensor( Uint8, - cute.make_layout((8, num_warps, num_stages)), - byte_alignment=8, + cute.make_layout((scale_bytes, num_warps, num_stages)), + byte_alignment=scale_bytes, )[None, warp_id, None] # Prepare for 16B cp.async, also for BF16 smem loads later. @@ -239,14 +243,16 @@ def device_kernel( # Last 64 elems are BF16 tail, corresponds to dequant1 of last # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 + - # 8x 16B for BF16. - if lane_id + 32 >= fp8_dim // 8: - idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 - cute.copy( - cp16_atom, - s_kdata_16B_slice[(None, idx), compute_stage], - dequant1, - ) + # 8x 16B for BF16. The V4.1 record is fully quantized and has + # no tail. + if cutlass.const_expr(bf16_dim > 0): # noqa: SIM102 + if lane_id + 32 >= fp8_dim // 8: + idx = fp8_dim // 16 + (lane_id + 32) - fp8_dim // 8 + cute.copy( + cp16_atom, + s_kdata_16B_slice[(None, idx), compute_stage], + dequant1, + ) # Store two 16B BF16 chunks per lane: first half, then second half. dst = out_slice[req_id, offset + i, (None, lane_id)] @@ -269,7 +275,7 @@ def host_entrypoint( ): # Split k_cache into k_data and k_scale. Each [block_size, head_bytes] # block is actually a concat of - # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8]. + # [block_size, fp8_dim + bf16_dim * 2] and [block_size, scale_bytes]. k_data = cute.make_tensor( k_cache.iterator, layout=cute.make_layout( @@ -282,8 +288,8 @@ def host_entrypoint( k_scale = cute.make_tensor( k_cache.iterator + (block_size * data_dim), layout=cute.make_layout( - (k_cache.shape[0], block_size, 8), - stride=(k_cache.stride[0], 8, 1), + (k_cache.shape[0], block_size, scale_bytes), + stride=(k_cache.stride[0], scale_bytes, 1), ), ) @@ -311,10 +317,14 @@ def dispatch( # type: ignore[override] *, block_size: int, has_gather_lens: bool, + fp8_dim: int = 448, + group_size: int = 64, ) -> CompileKey: return self.CompileKey( block_size=block_size, has_gather_lens=has_gather_lens, + fp8_dim=fp8_dim, + group_size=group_size, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -345,9 +355,11 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: num_reqs = cute.sym_int() head_dim = self.head_dim - fp8_dim = self.fp8_dim + fp8_dim = compile_key.fp8_dim block_size = compile_key.block_size - head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8 + head_bytes = ( + fp8_dim + (head_dim - fp8_dim) * 2 + head_dim // compile_key.group_size + ) out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16) k_cache = cute.runtime.make_fake_tensor( @@ -384,10 +396,14 @@ def __call__( block_table: torch.Tensor, block_size: int, offset: int, + fp8_dim: int = 448, + group_size: int = 64, ) -> CuTeDSLLaunchSpec["DequantGatherKCacheKernel.CompileKey"]: compile_key = self.dispatch( block_size=block_size, has_gather_lens=gather_lens is not None, + fp8_dim=fp8_dim, + group_size=group_size, ) launch_args = ( out, diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 8052178a03d9..f67a77e2541f 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -165,6 +165,7 @@ class CompileKey: rope_dim: int num_heads: int cos_sin_dtype: type[cutlass.Numeric] + weights_out_dtype: type[cutlass.Numeric] coarsen: int @staticmethod @@ -177,6 +178,7 @@ def kernel(compile_key: CompileKey) -> Any: # process multiple heads at the same time to armotize RoPE load costs coarsen = compile_key.coarsen tb_size = IndexerQMxFp4Kernel.tb_size + weights_out_dtype = compile_key.weights_out_dtype assert num_heads % coarsen == 0 # later we will use 32B load = 16 BF16 elems @@ -287,7 +289,7 @@ def device_kernel( weight_head_id = global_tid % num_heads weights_out[weight_token_id, weight_head_id] = ( weights[weight_token_id, weight_head_id].to(Float32) * scale - ) + ).to(weights_out_dtype) @cute.jit def host_entrypoint( @@ -320,11 +322,13 @@ def dispatch( # type: ignore[override] self, *, cos_sin_dtype: type[cutlass.Numeric], + weights_out_dtype: type[cutlass.Numeric], **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, cos_sin_dtype=cos_sin_dtype, + weights_out_dtype=weights_out_dtype, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -340,11 +344,17 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: return [] + # The DeepSeek V4.1 sparse-logits indexer takes bf16 weights; the + # dense scoring kernels take fp32. + weights_out_dtypes: tuple[type[cutlass.Numeric], ...] = (Float32,) + if vllm_config.attention_config.indexer_sparse_logits: + weights_out_dtypes = (Float32, BFloat16) return self._trace_dispatch(self.dispatch)( head_dim=head_dim, rope_dim=rope_dim, num_heads=num_heads, cos_sin_dtype=Float32, + weights_out_dtype=weights_out_dtypes, coarsen=(1, 4), ) @@ -380,7 +390,9 @@ def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]: divisibility=4, ) weights_out = make_fake_tensor( - Float32, (num_tokens, compile_key.num_heads), divisibility=4 + compile_key.weights_out_dtype, + (num_tokens, compile_key.num_heads), + divisibility=4, ) return ( positions, @@ -413,6 +425,7 @@ def __call__( rope_dim=cos_sin_cache.shape[-1], num_heads=num_heads, cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + weights_out_dtype=torch_to_cute_dtype(weights_out.dtype), coarsen=1 if num_tokens < 512 else 4, ) launch_args = ( diff --git a/vllm/models/deepseek_v41/amd/dspark.py b/vllm/models/deepseek_v41/amd/dspark.py index 3a6cef8aa74d..78f07b6b11db 100644 --- a/vllm/models/deepseek_v41/amd/dspark.py +++ b/vllm/models/deepseek_v41/amd/dspark.py @@ -259,6 +259,8 @@ def _insert_context_kv( attn.padded_heads, attn.eps, block_size, + True, # apply_q_norm (unused: the query is a discarded dummy) + attn.kv_mxfp8, ) elif cache_dtype == torch.bfloat16: swa_3d = swa_cache.view(-1, block_size, attn.head_dim) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 131cf4d97e7e..134328dbe908 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -4,6 +4,7 @@ DeepseekV4 MLA Attention Layer """ +import math from abc import ABC, abstractmethod from collections.abc import Callable from functools import cached_property @@ -25,6 +26,7 @@ RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, @@ -51,6 +53,7 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.models.deepseek_v41.common.rope import build_deepseek_v4_rope from vllm.models.deepseek_v41.compressor import DeepseekCompressor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import ( execute_in_parallel, @@ -62,6 +65,9 @@ dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerBackend, +) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache from vllm.v1.kv_cache_interface import ( KVCacheSpec, @@ -109,6 +115,16 @@ def _fill_short_context_topk_indices( ) +# Which packed fp8_ds_mla record a V4.1 layer writes. FlashMLA decodes +# DeepSeek's V4.1 record -- all 512 dims (RoPE included) as fp8 e4m3 with one +# UE8M0 scale per 32 dims, 512 data bytes + 16 scale bytes per token, pages +# rounded to the kernel's 512 B TMA stride -- only in its SM100 sparse-decode +# kernels. Every other arch keeps the V4 record: 448 fp8 NoPE + 64 bf16 RoPE +# plus 7 UE8M0 scales of 64 dims and a pad byte (584 B, 576 B pages). +def _use_v41_mxfp8_kv_record() -> bool: + return current_platform.is_device_capability_family(100) + + def _resolve_dsv4_kv_cache_dtype( use_fp8_ds_mla_layout: bool, kv_cache_dtype: str, @@ -445,6 +461,9 @@ def __init__( self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config ) + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584 + self.kv_page_alignment = 512 if self.kv_mxfp8 else 576 self.swa_cache_layer = DeepseekV4SWACache( head_dim=self.head_dim, @@ -454,6 +473,8 @@ def __init__( cache_config=cache_config, backend_cls=self.swa_backend_cls, block_size=32, + packed_bytes_per_token=self.kv_bytes_per_token, + packed_page_alignment=self.kv_page_alignment, ) # The attention layer itself was already registered with the @@ -879,6 +900,7 @@ def _fused_qnorm_rope_kv_insert( self.eps, swa_metadata.block_size, False, + self.kv_mxfp8, ) # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV @@ -934,9 +956,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # DeepseekV4SWACache. if not self.is_kv_source: return None - # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B - # alignment; plain bf16 / per-tensor fp8 rows use natural element-size - # pages. + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout whose page rounds up + # to the decode kernel's TMA stride; plain bf16 / per-tensor fp8 rows + # use natural element-size pages. uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, @@ -945,12 +967,13 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, tokens_per_state=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else 512, + alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), - # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token; - # head_size stays semantic (512). - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, + # Packed record width; head_size stays semantic (512). + state_content_bytes=( + self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None + ), ) def _compressed_kv_cache(self) -> torch.Tensor: @@ -979,7 +1002,9 @@ def __init__( self.cache_config = cache_config self.dtype = dtype self.compress_ratio = compress_ratio - compilation_config = get_current_vllm_config().compilation_config + vllm_config = get_current_vllm_config() + self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self @@ -992,19 +1017,34 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + page_alignment = ( + 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512 + ) return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, tokens_per_state=self.compress_ratio, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + # The indexer page shares a physical block with the main KV cache + # (the block stride is the sum of every page in the group), so it + # has to carry the same padding: the packed record's TMA stride, + # else 512B for FlashInfer sparse (#44577). + alignment=page_alignment, + # DeepGEMM's paged sparse MQA logits address pages by the block + # stride, which packs every layer's page in this layout, and need + # it 512B-aligned; the main KV record read from the same blocks + # needs its own TMA stride, so keep both. + block_stride_alignment=( + math.lcm(512, page_alignment) if self.sparse_logits else None + ), ) def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: + if self.sparse_logits: + return DeepseekV41SparseIndexerBackend return DeepseekV41IndexerBackend @@ -1104,21 +1144,52 @@ def __init__( self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) self.k_cache = k_cache - self.indexer_op = SparseAttnIndexer( - self.k_cache, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - skip_k_cache_insert=True, - use_fp4_cache=self.use_fp4_kv, - compress_ratio=self.compress_ratio, - candidate_blocks=candidate_block_buffer, - candidate_block_size=candidate_block_size, - candidate_write=candidate_write, + # Candidate consumers can score only the candidate blocks (opt-in); + # the candidate source and pre-candidate indexers stay dense. + use_sparse_logits = ( + vllm_config.attention_config.indexer_sparse_logits + and candidate_block_buffer is not None + and not candidate_write + ) + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer + if use_sparse_logits: + if not self.use_fp4_kv: + raise ValueError( + "attention_config.indexer_sparse_logits requires the MXFP4 " + "indexer cache (attention_config.indexer_kv_dtype='mxfp4')." + ) + assert candidate_block_buffer is not None + assert topk_indices_buffer is not None + self.indexer_op = SparseMQAIndexer( + self.k_cache, + self.topk_tokens, + self.head_dim, + self.max_total_seq_len, + topk_indices_buffer, + candidate_block_buffer, + candidate_block_size, + ) + else: + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + compress_ratio=self.compress_ratio, + candidate_blocks=candidate_block_buffer, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, + ) + # The fused Q kernel writes the per-head weights in the dtype the + # scoring kernels take, so no cast runs per step. + self.indexer_weights_dtype = ( + SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32 ) def _produce_k( @@ -1206,6 +1277,7 @@ def forward( self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, + weights_out_dtype=self.indexer_weights_dtype, ) if isinstance(q_quant, tuple): q, q_scale = q_quant diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 64462eb73b4a..51743016bf07 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -36,6 +36,18 @@ from vllm.utils.import_utils import has_cutedsl from vllm.utils.math_utils import next_power_of_2 +# Per-token byte width of the two paged fp8 records. Both put a page's whole +# data region ahead of its whole scale region, so ``k_cache.shape[-1]`` -- the +# width the KV-cache spec publishes -- identifies the record: +# V4 (584 B): 448 fp8 e4m3 NoPE + 64 bf16 RoPE, then 7 UE8M0 scales of 64 +# dims each plus a pad byte. +# V4.1 (528 B): all 512 dims as fp8 e4m3 (RoPE included), then 16 UE8M0 +# scales of 32 dims each. FlashMLA's ``ModelType::V41``. +V4_BYTES_PER_TOKEN = 584 +V41_BYTES_PER_TOKEN = 528 +V41_QUANT_BLOCK = 32 +V41_NUM_SCALES = 512 // V41_QUANT_BLOCK # 16 + @triton.jit def quantize_and_insert_k_kernel( @@ -163,6 +175,67 @@ def quantize_and_insert_k_kernel( tl.store(bf16_out_ptr + chunk_offsets, bf16_vals) +@triton.jit +def _quantize_and_insert_k_mxfp8_kernel( + k_ptr, # [num_tokens, 512] bf16 + slot_mapping_ptr, # [num_tokens] int64 + k_cache_ptr, # [num_blocks, block_bytes] as uint8 (flattened view) + num_tokens, + head_dim: tl.constexpr, # 512, RoPE dims included + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 dims per UE8M0 scale + cache_block_size: tl.constexpr, # paged cache block size, in tokens + block_stride: tl.constexpr, # total bytes per block (padded) + fp8_max: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Quantize a bf16 K row to MXFP8 and insert it into a V4.1 paged cache. + + One program per token. Unlike the V4 record there is no bf16 tail, so every + dim is scaled. + """ + pid = tl.program_id(0) + if pid >= num_tokens: + return + + slot_idx = tl.load(slot_mapping_ptr + pid) + if slot_idx == -1: + return + + block_idx = slot_idx // cache_block_size + pos_in_block = slot_idx % cache_block_size + + # int64: block_idx * block_stride can exceed 2^31 with many KV-cache blocks. + cache_block_ptr = k_cache_ptr + block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + d = tl.arange(0, head_dim) + x = tl.load(k_ptr + pid.to(tl.int64) * head_dim + d).to(tl.float32) + + tiles = tl.reshape(x, (scale_dim, quant_block)) + # Floor the absmax like the CUDA encoder so an all-zero tile still lands on + # a finite scale. + amax = tl.maximum(tl.max(tl.abs(tiles), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / fp8_max))) + scaled = tl.clamp( + tiles * tl.reshape(tl.exp2(-exponent), (scale_dim, 1)), -fp8_max, fp8_max + ) + if use_fnuz: # noqa: SIM108 + fp8 = scaled.to(tl.float8e4b8) + else: + fp8 = scaled.to(tl.float8e4nv) + tl.store( + token_data_ptr + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (head_dim,)) + ) + + # UE8M0 encoding: stored_value = exponent + 127 (bias). + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), 255.0) + tl.store(token_scale_ptr + tl.arange(0, scale_dim), encoded.to(tl.uint8)) + + def quantize_and_insert_k_cache( k: torch.Tensor, # [num_tokens, 512] bf16 k_cache: torch.Tensor, # [num_blocks, block_bytes] uint8 @@ -170,16 +243,13 @@ def quantize_and_insert_k_cache( block_size: int = 64, is_ue8m0: bool = True, use_fnuz: bool = False, + bytes_per_token: int = V4_BYTES_PER_TOKEN, ): """ Quantize K tensor and insert into paged K cache. - K Cache block layout (block_size=64 tokens): - - First 64 * 576 = 36864 bytes: Token data - - Each token: 448 bytes (fp8) + 128 bytes (bf16) - - Next 64 * 8 = 512 bytes: Scales - - Each token: 8 bytes (uint8 scales, 7 real + 1 padding) - - Padded to multiple of 576 + ``bytes_per_token`` picks the record (see the module header): the V4 one, + or V4.1's all-dims MXFP8 one. ``use_fnuz=True`` selects FNUZ E4M3 cache encoding and is only valid on platforms whose FP8 format is FNUZ. ``use_fnuz=False`` selects OCP E4M3, @@ -196,6 +266,26 @@ def quantize_and_insert_k_cache( num_tokens = slot_mapping.shape[0] block_stride = k_cache.stride(0) # bytes per block + if bytes_per_token == V41_BYTES_PER_TOKEN: + if use_fnuz and not current_platform.is_fp8_fnuz(): + raise ValueError("use_fnuz=True requires a platform using FNUZ FP8") + _quantize_and_insert_k_mxfp8_kernel[(num_tokens,)]( + k, + slot_mapping, + k_cache, + num_tokens, + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=block_stride, + fp8_max=get_fp8_min_max()[1] + if use_fnuz + else torch.finfo(torch.float8_e4m3fn).max, + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -343,6 +433,69 @@ def _dequantize_and_gather_k_kernel( tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals) +@triton.jit +def _dequantize_and_gather_k_mxfp8_kernel( + out_ptr, + out_stride0, + out_stride1, + k_cache_ptr, + seq_lens_ptr, + block_table_ptr, + offset, + gather_lens_ptr, + max_blocks_per_seq: tl.constexpr, + head_dim: tl.constexpr, # 512 + scale_dim: tl.constexpr, # 16 + quant_block: tl.constexpr, # 32 + cache_block_size: tl.constexpr, + block_stride: tl.constexpr, + use_fnuz: tl.constexpr = False, +): + """Gather V4.1 MXFP8 rows into a bf16 workspace, dequantizing in place.""" + batch_idx = tl.program_id(0) + worker_id = tl.program_id(1) + num_workers = tl.num_programs(1) + + seq_len = tl.load(seq_lens_ptr + batch_idx) + if gather_lens_ptr is not None: # noqa: SIM108 + gather_len = tl.load(gather_lens_ptr + batch_idx) + else: + gather_len = seq_len + start_pos = seq_len - gather_len + + d = tl.arange(0, head_dim) + s = tl.arange(0, scale_dim) + for i in range(worker_id, gather_len, num_workers): + pos = start_pos + i + block_in_seq = pos // cache_block_size + pos_in_block = pos % cache_block_size + + block_table_row_ptr = block_table_ptr + batch_idx * max_blocks_per_seq + physical_block_idx = tl.load(block_table_row_ptr + block_in_seq) # int32 + + # int64: physical_block_idx * block_stride can exceed 2^31. + cache_block_ptr = k_cache_ptr + physical_block_idx.to(tl.int64) * block_stride + token_data_ptr = cache_block_ptr + pos_in_block * head_dim + token_scale_ptr = ( + cache_block_ptr + cache_block_size * head_dim + pos_in_block * scale_dim + ) + + x_uint8 = tl.load(token_data_ptr + d) + if use_fnuz: + x_fp8 = x_uint8.to(tl.float8e4b8, bitcast=True) + else: + x_fp8 = x_uint8.to(tl.float8e4nv, bitcast=True) + tiles = tl.reshape(x_fp8.to(tl.float32), (scale_dim, quant_block)) + + # UE8M0: scale = 2^(stored_value - 127). + encoded = tl.load(token_scale_ptr + s) + scale = tl.exp2(encoded.to(tl.float32) - 127.0) + dequant = tl.reshape(tiles * tl.reshape(scale, (scale_dim, 1)), (head_dim,)) + + output_row_ptr = out_ptr + batch_idx * out_stride0 + (offset + i) * out_stride1 + tl.store(output_row_ptr + d, dequant.to(tl.bfloat16)) + + def dequantize_and_gather_k_cache_triton( # [num_reqs, max_num_tokens, head_size] out: torch.Tensor, @@ -358,6 +511,28 @@ def dequantize_and_gather_k_cache_triton( offset: int, use_fnuz: bool = False, ) -> None: + num_reqs = seq_lens.shape[0] + NUM_WORKERS = 128 + if k_cache.shape[-1] == V41_BYTES_PER_TOKEN: + _dequantize_and_gather_k_mxfp8_kernel[(num_reqs, NUM_WORKERS)]( + out, + out.stride(0), + out.stride(1), + k_cache, + seq_lens, + block_table, + offset, + gather_lens, + max_blocks_per_seq=block_table.shape[-1], + head_dim=512, + scale_dim=V41_NUM_SCALES, + quant_block=V41_QUANT_BLOCK, + cache_block_size=block_size, + block_stride=k_cache.stride(0), + use_fnuz=use_fnuz, + ) + return + TOKEN_FP8_DIM = 448 TOKEN_BF16_DIM = 64 TOKEN_SCALE_DIM = 8 @@ -365,8 +540,6 @@ def dequantize_and_gather_k_cache_triton( FP8_MAX = 448.0 TOKEN_DATA_SIZE = TOKEN_FP8_DIM + TOKEN_BF16_DIM * 2 - num_reqs = seq_lens.shape[0] - NUM_WORKERS = 128 _dequantize_and_gather_k_kernel[(num_reqs, NUM_WORKERS)]( out, out.stride(0), @@ -408,6 +581,8 @@ def dequantize_and_gather_k_cache( ) -> None: """Dequantize and gather a paged DSv4 K cache. + The record is read off ``k_cache.shape[-1]``; see the module header. + ``use_fnuz`` MUST match the encoder of the specific cache being read: ``False`` for ``compressed_k_cache`` (Triton encoder is OCP everywhere), ``current_platform.is_fp8_fnuz()`` for ``swa_k_cache`` (C++ encoder @@ -419,6 +594,7 @@ def dequantize_and_gather_k_cache( _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL, ) + mxfp8 = k_cache.shape[-1] == V41_BYTES_PER_TOKEN _DEQUANT_GATHER_K_CACHE_CUTEDSL_KERNEL( out=out, k_cache=k_cache, @@ -427,6 +603,8 @@ def dequantize_and_gather_k_cache( block_table=block_table, block_size=block_size, offset=offset, + fp8_dim=512 if mxfp8 else 448, + group_size=V41_QUANT_BLOCK if mxfp8 else 64, ) return diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 5950c4b6ba5c..96f15849ddd3 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -232,10 +232,13 @@ def rope_quant_insert( The BF16 latent supplies both NoPE quantization and RoPE input. It is read only for valid slots at group boundaries. The cache dtype selects the - layout: ``uint8`` is the fp8_ds_mla paged layout (576 value bytes and eight - segregated UE8M0 scale bytes per token, including one zero padding scale); - ``bfloat16`` and ``float8_e4m3fn`` are the plain [448 NoPE | 64 RoPE] rows - read by FlashInfer, the latter scaled by the per-tensor ``fp8_scale``. + layout: ``uint8`` is the fp8_ds_mla paged layout, whose record the + per-token byte width names -- 584 B for V4 (576 value bytes and eight + segregated UE8M0 scale bytes, including one zero padding scale) or 528 B + for V4.1 (512 MXFP8 value bytes covering the RoPE dims too, then 16 UE8M0 + scales of 32 dims each). ``bfloat16`` and ``float8_e4m3fn`` are the plain + [448 NoPE | 64 RoPE] rows read by FlashInfer, the latter scaled by the + per-tensor ``fp8_scale``. """ assert compress_ratio in (1, 2) assert latent.shape[1] == 512 and latent.dtype == torch.bfloat16 @@ -246,8 +249,15 @@ def rope_quant_insert( return launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {} if kv_cache.dtype == torch.uint8: - assert kv_cache.shape[-1] == 584 - _rope_quant_insert_kernel[(num_tokens,)]( + assert kv_cache.shape[-1] in (584, 528), ( + f"unsupported paged KV record width {kv_cache.shape[-1]}" + ) + kernel = ( + _rope_quant_insert_mxfp8_kernel + if kv_cache.shape[-1] == 528 + else _rope_quant_insert_kernel + ) + kernel[(num_tokens,)]( latent, positions, cos_sin_cache, @@ -338,6 +348,60 @@ def _rope_quant_insert_kernel( tl.store(rope_dst + d - 448, rotated.to(tl.bfloat16), d >= 448) +@triton.jit +def _rope_quant_insert_mxfp8_kernel( + latent, + positions, + cos_sin, + cache, + cache_slots, + COS_STRIDE: tl.constexpr, + CACHE_STRIDE: tl.constexpr, + CACHE_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, +): + """V4.1 record: RoPE first, then MXFP8-quantize all 512 dims. + + The RoPE dims are quantized here too, so unlike the V4 kernel the rotation + has to happen before the scales are picked. + """ + t = tl.program_id(0) + slot = tl.load(cache_slots + t) + if slot < 0: + return + position = tl.load(positions + t) + if (position + 1) % COMPRESS_RATIO != 0: + return + d = tl.arange(0, 512) + normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32) + + # NoPE pairs load (cos, sin) = (1, 0), so the rotation is the identity there. + even, odd = tl.split(tl.reshape(normed, (256, 2))) + pair = tl.arange(0, 256) - 224 + cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE + c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32) + s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32) + rotated = tl.interleave(even * c - odd * s, odd * c + even * s) + if SANITIZE_CACHE_NANS: + rotated = tl.where(rotated == rotated, rotated, 0.0) + + page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE + values = page + (slot % CACHE_BLOCK) * 512 + scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16 + + quant = tl.reshape(rotated, (16, 32)) + amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4) + exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0))) + scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1)) + fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv) + tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,))) + + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded) + tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8)) + + @triton.jit def _rope_plain_insert_kernel( latent, diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 3c711681aa75..214681f3a932 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -217,15 +217,14 @@ def forward( residual, ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) - if self.use_sequence_parallel: - hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] - pre_mix = sp_all_gather(pre_mix)[:full_num_tokens] # Collapse the hc copies with the pre-mix from the last layer's FFN # mixes — the mix the reference forward_head applies (hc_pre with the # last block's ffn pre-mix). Return the PRE-norm head hidden; # compute_logits applies self.norm. assert pre_mix is not None hidden_states = hc_collapse_triton(hidden_states, pre_mix) + if self.use_sequence_parallel: + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] return hidden_states @@ -235,65 +234,24 @@ def _insert_context_kv( positions: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: - """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. - - Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy - query and discard it, since context tokens have no query). Mirrors - ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. - """ - swa_cache = attn.swa_cache_layer.kv_cache - block_size = attn.swa_cache_layer.block_size - cos_sin_cache = attn.rotary_emb.cos_sin_cache - cache_dtype = swa_cache.dtype - n_ctx = kv.shape[0] - dummy_q = torch.zeros( - (n_ctx, attn.n_local_heads, attn.head_dim), - dtype=kv.dtype, - device=kv.device, + """Insert normalized context KV without constructing an unused query.""" + # The bound cache already is [num_blocks, block_size, bytes_per_token], so + # the record width needs no per-layer attribute -- one whose name the + # attention layer is free to change (kv_bytes_per_token has to become + # swa_bytes_per_token once the sliding-window and compressed records can + # differ). block_size stays the layer's, so the op's size(1) check still + # cross-checks the two rather than restating the tensor. + cache = attn.swa_cache_layer.kv_cache + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + kv, + cache, + slot_mapping, + positions, + attn.rotary_emb.cos_sin_cache, + attn.swa_cache_layer.block_size, + attn._flashinfer_fp8_kv_scale if cache.dtype == torch.float8_e4m3fn else None, + attn.kv_mxfp8, ) - if cache_dtype == torch.uint8: - # fp8_ds_mla UE8M0 paged layout - swa_2d = swa_cache.view(swa_cache.shape[0], -1) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( - dummy_q, - kv, - swa_2d, - slot_mapping, - positions, - cos_sin_cache, - attn.padded_heads, - attn.eps, - block_size, - ) - elif cache_dtype == torch.bfloat16: - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( - dummy_q, - kv, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn.eps, - block_size, - ) - else: # per-tensor fp8 (torch.float8_e4m3fn) - # TODO(ben): double-check if this is being dispatched correctly for FI backend - swa_3d = swa_cache.view(-1, block_size, attn.head_dim) - dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( - dummy_q, - kv, - dummy_q_fp8, - swa_3d, - slot_mapping, - positions, - cos_sin_cache, - attn._flashinfer_fp8_kv_scale, - attn._flashinfer_fp8_q_scale_inv, - attn.eps, - block_size, - ) class DSparkDeepseekV4ForCausalLM(nn.Module): diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index f3530544da7b..9c9719e97729 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -20,7 +20,6 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -79,6 +78,7 @@ from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes +from .ops.mega_mhc import mhc_shifted_post_pre if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -403,25 +403,23 @@ def forward( else: # The collapse already reads the post-mapped streams, so the mean # aux consumers want comes out of the same kernel. - residual, post_mix, res_mix, x, attn_pre, aux = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=pre_mix, - norm_weight=self.attn_norm.weight, - norm_eps=self.attn_norm.variance_epsilon, - capture_aux=capture_previous_aux, - ) + residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + capture_aux=capture_previous_aux, ) if capture_previous_aux: previous_aux = aux @@ -433,24 +431,22 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual, post_mix, res_mix, x, ffn_pre, _ = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=attn_pre, - norm_weight=self.ffn_norm.weight, - norm_eps=self.ffn_norm.variance_epsilon, - ) + residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=attn_pre, + norm_weight=self.ffn_norm.weight, + norm_eps=self.ffn_norm.variance_epsilon, ) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre, previous_aux diff --git a/vllm/models/deepseek_v41/nvidia/ops/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py new file mode 100644 index 000000000000..f2c74902616f --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools + +import torch + +from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_delayed_tilelang, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import ( + _import_deep_gemm, + is_deep_gemm_supported, + mega_mhc, +) + + +@functools.cache +def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool: + """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel.""" + if not ( + is_deep_gemm_supported() + and current_platform.is_device_capability_family(100) + and hidden_size > 0 + and hidden_size % 1024 == 0 + and hc_mult == 4 + ): + return False + deep_gemm = _import_deep_gemm() + return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None)) + + +def mhc_shifted_post_pre_deep_gemm( + x: torch.Tensor, + residual: torch.Tensor, + shifted_prev_mix: torch.Tensor, + post_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + mix_scales: torch.Tensor, + mix_bases: torch.Tensor, + hc_norm_eps: float, + hc_pre_eps: float, + hc_post_scale: float, + sinkhorn_eps: float, + num_sinkhorn_iters: int, + rmsnorm_weight: torch.Tensor, + rmsnorm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC.""" + num_tokens, hidden_size = x.shape + hc_mult = residual.shape[1] + new_residual = torch.empty_like(residual) + new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32) + new_post_mix = torch.empty_like(post_mix) + new_comb_res_mix = torch.empty_like(comb_res_mix) + y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16) + mega_mhc( + x=x, + residual=residual, + shifted_prev_mix=shifted_prev_mix.unsqueeze(-1), + post_mix=post_mix, + comb_res_mix=comb_res_mix, + fn=fn, + mix_scales=mix_scales, + mix_bases=mix_bases, + hc_mult=hc_mult, + hc_norm_eps=hc_norm_eps, + hc_pre_eps=hc_pre_eps, + hc_post_scale=hc_post_scale, + sinkhorn_eps=sinkhorn_eps, + num_sinkhorn_iters=num_sinkhorn_iters, + rmsnorm_weight=rmsnorm_weight, + rmsnorm_eps=rmsnorm_eps, + rmsnorm_scale=1.0, + new_residual=new_residual, + new_prev_mix=new_prev_mix, + new_post_mix=new_post_mix, + new_comb_res_mix=new_comb_res_mix, + y_bf16=y_bf16, + ) + return ( + new_residual, + new_post_mix, + new_comb_res_mix, + y_bf16, + new_prev_mix.squeeze(-1), + ) + + +def mhc_shifted_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Use Mega-mHC for carried mixing, retaining TileLang for aux capture.""" + if ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ): + outputs = mhc_shifted_post_pre_deep_gemm( + x, + residual, + pre_mix, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_post_mult_value, + hc_sinkhorn_eps, + sinkhorn_repeat, + norm_weight, + norm_eps, + ) + return *outputs, x.new_empty(0, x.shape[1]) + + return mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + capture_aux=capture_aux, + ) diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py index cbd9ad7f4d78..27928b6e6f43 100644 --- a/vllm/models/dots3_note/nvidia/attention.py +++ b/vllm/models/dots3_note/nvidia/attention.py @@ -19,7 +19,11 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import ( + PIN_MEMORY, + async_tensor_h2d, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionCGSupport, AttentionLayer, @@ -231,8 +235,12 @@ def _build_sliding_window_metadata( query_lens = query_lens_cpu[req_start:req_end] kv_lens = kv_lens_cpu[req_start:req_end] num_reqs = req_end - req_start - cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) - cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32) + cu_seq_lens_q_cpu = torch.zeros( + num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY + ) + cu_seq_lens_k_cpu = torch.zeros( + num_reqs + 1, dtype=torch.int32, pin_memory=PIN_MEMORY + ) torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:]) torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:]) token_to_seq_cpu = torch.repeat_interleave( @@ -248,8 +256,8 @@ def _build_sliding_window_metadata( query_end=query_end, cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True), cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True), - starts=starts_cpu[req_start:req_end].to(device, non_blocking=True), - token_to_seq=token_to_seq_cpu.to(device, non_blocking=True), + starts=async_tensor_h2d(starts_cpu[req_start:req_end], device), + token_to_seq=async_tensor_h2d(token_to_seq_cpu, device), num_kv_tokens=num_kv_tokens, max_seq_len_q=int(query_lens.max().item()), max_seq_len_k=int(kv_lens.max().item()), @@ -396,8 +404,8 @@ def build( seq_lens = common_attn_metadata.seq_lens[ reqs_start + req_start : reqs_start + req_end ] - query_lens = query_lens_cpu[req_start:req_end].to( - self.device, non_blocking=True + query_lens = async_tensor_h2d( + query_lens_cpu[req_start:req_end], self.device ) kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1) chunk.cu_seq_lens_k.zero_() diff --git a/vllm/models/glm5next/__init__.py b/vllm/models/glm5next/__init__.py index a1146e2602eb..b38806260bd4 100644 --- a/vllm/models/glm5next/__init__.py +++ b/vllm/models/glm5next/__init__.py @@ -6,8 +6,8 @@ if current_platform.is_xpu(): raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.") -from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration -from .nvidia.mtp import Glm5NextMTP +from .common.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration +from .common.mtp import Glm5NextMTP __all__ = [ "Glm5NextForCausalLM", diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 64e6fd46d73f..09e7a40aa5e0 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -65,50 +65,6 @@ def _hadamard128(x): return x * 0.08838834764831845 # 1/sqrt(128) -# Map pool ids to physical cache slots. - - -def compute_pooled_write_locs( - page_table_64: torch.Tensor, - pool_ids: torch.Tensor, - pool_size: int, -) -> torch.Tensor: - """Map logical pooled-K ids to physical flat cache slots. - - ``pool_size`` consecutive tokens share one pool slot that lives at the - *first* token page of each page-group. ``page_table_64`` maps token pages - to physical block ids; we gather the block id of each pool's page-group - and add the in-block pool offset. - """ - assert page_table_64.ndim == 1 - pool_ids = pool_ids.to(torch.int64) - block_size = 64 - pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor") - token_page_row = pool_page_group * pool_size - packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64)) - return packed_page.to(torch.int64) * block_size + torch.remainder( - pool_ids, block_size - ) - - -def build_pooled_page_table( - page_table: torch.Tensor, - pool_size: int, -) -> torch.Tensor: - """Build a pool-granular page table by taking every ``pool_size``-th - token-page column (one pool maps to ``pool_size`` token pages). - - Uses gather (not strided slicing) so the result is always a fresh - row-major tensor — some downstream kernels require stride(-1) == 1. - """ - block_size = page_table.shape[-1] - assert block_size % pool_size == 0, ( - f"pool_size ({pool_size}) must divide page columns ({block_size})" - ) - idx = torch.arange(0, block_size, pool_size, device=page_table.device) - return page_table[..., idx].contiguous() - - # Fused pool compression and cache write. diff --git a/vllm/models/glm5next/common/__init__.py b/vllm/models/glm5next/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/glm5next/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/glm5next/nvidia/attention.py b/vllm/models/glm5next/common/attention.py similarity index 100% rename from vllm/models/glm5next/nvidia/attention.py rename to vllm/models/glm5next/common/attention.py diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/common/kda.py similarity index 100% rename from vllm/models/glm5next/nvidia/kda.py rename to vllm/models/glm5next/common/kda.py diff --git a/vllm/models/glm5next/nvidia/model.py b/vllm/models/glm5next/common/model.py similarity index 98% rename from vllm/models/glm5next/nvidia/model.py rename to vllm/models/glm5next/common/model.py index 21918bdf6bed..e4f669835e49 100644 --- a/vllm/models/glm5next/nvidia/model.py +++ b/vllm/models/glm5next/common/model.py @@ -1036,6 +1036,17 @@ class Glm5NextForConditionalGeneration( has_inner_state: ClassVar[Literal[True]] = True is_hybrid: ClassVar[Literal[True]] = True + # GLM-5.3-Flash stores the dense-MLP gate/up as separate tensors (like + # ``Glm4vMoeForConditionalGeneration``, ``glm4_moe`` and ``deepseek_v2``), + # so the fused ``gate_up_proj`` must expand to its real shard names for + # per-layer quant-scheme resolution. The identity ``gate_up_proj`` entry + # inherited from ``Glm4vForConditionalGeneration`` (pre-fused gate_up_proj) + # would otherwise route the module to ``global_quant_config`` and mismatch + # at load for mixed-precision Quark checkpoints. + packed_modules_mapping = { + "gate_up_proj": ["gate_proj", "up_proj"], + } + # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration # (``model.visual.`` -> ``visual.``, ``model.language_model.`` -> # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``), @@ -1259,7 +1270,9 @@ def _try_load_fp8_attn_proj( return False suffix, (key, target_base, shard_id, is_kva) = matched is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn - is_scale = "weight_scale_inv" in name + # Need to accept both the DeepSeek-native ``weight_scale_inv`` and the Quark + # ``weight_scale`` names before feeding the shared block dequant below. + is_scale = "weight_scale_inv" in name or name.endswith(".weight_scale") if not is_weight and not is_scale: return False diff --git a/vllm/models/glm5next/nvidia/mtp.py b/vllm/models/glm5next/common/mtp.py similarity index 99% rename from vllm/models/glm5next/nvidia/mtp.py rename to vllm/models/glm5next/common/mtp.py index 35b452f44f79..33a5402b8fd1 100644 --- a/vllm/models/glm5next/nvidia/mtp.py +++ b/vllm/models/glm5next/common/mtp.py @@ -22,6 +22,7 @@ from vllm.model_executor.models.deepseek_mtp import SharedHead from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts from vllm.model_executor.models.utils import maybe_prefix +from vllm.models.glm5next.nvidia.ops.fused_eh_norm import fused_eh_norm from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -33,7 +34,6 @@ _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) -from .ops.fused_eh_norm import fused_eh_norm class Glm5NextMultiTokenPredictorLayer(nn.Module): diff --git a/vllm/models/glm5next/nvidia/multimodal.py b/vllm/models/glm5next/common/multimodal.py similarity index 100% rename from vllm/models/glm5next/nvidia/multimodal.py rename to vllm/models/glm5next/common/multimodal.py diff --git a/vllm/models/inkling/amd/ops/lamport.py b/vllm/models/inkling/amd/ops/lamport.py deleted file mode 100644 index 9705a9a3df37..000000000000 --- a/vllm/models/inkling/amd/ops/lamport.py +++ /dev/null @@ -1,766 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. - -The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. - -Liveness --------- -Large grids are deliberately split at the two communication dependencies: - - 1. ``_publish_input_kernel`` only publishes rank partials. - 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. - 3. ``_sconv_publish_kernel`` only computes and publishes the local result. - 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. - -Without PDL, CUDA stream order completes a producer before its consumer. With -PDL, every producer CTA posts all peer stores before triggering its dependent, -and the consumer executes ``gdc_wait`` before polling. A consumer therefore -waits only for stores from a producer that is already running (or complete) on -another GPU. It never waits for another block in its own grid. Consequently -spinning consumers cannot occupy resources needed by any producer, and the -wait-for graph has no cycle. The proof is independent of grid size, block -dispatch order, and occupancy. - -For one token, the first three phases use eight independent channel slices to -expose enough CTA parallelism for decode latency. Each slice has exclusive -ownership of its cache and Lamport columns; the gather/RMSNorm phase retains -one CTA per token so no cross-CTA reduction or completion counter is needed. - -Immediate buffer reuse (including replay of a captured CUDA graph) is also -safe. Rank R cannot republish input for call n+1 until its gather for n has -finished; that gather waited for owner O's output, which O publishes only -after consuming R's input for n. Likewise, R cannot republish output for n+1 -until its reduction for n+1 has observed destination D's input; D publishes -that input only after its gather consumed R's output for n. Thus every prior -read happens-before a same-slot rewrite. Three generations reduce incidental -coupling for ordinary launches, but correctness does not depend on rotation. - -The payload itself is the Lamport flag. A 32-bit store publishes two bf16s -atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real -negative zeroes are changed to positive zero before publication. Consumers -use volatile 32-bit loads and restore the sentinel after consuming a slot. -""" - -from __future__ import annotations - -import os - -import torch - -from vllm.distributed import get_tp_group -from vllm.distributed.parallel_state import in_the_same_node_as -from vllm.logger import init_logger -from vllm.triton_utils import tl, triton - -logger = init_logger(__name__) - -_MAX_TOKENS = 16384 -_EMPTY_PAIR = tl.constexpr(0x80008000) - - -@triton.jit -def _pack_bf16_pairs(values): - """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" - lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) - lo = lo.to(tl.uint16, bitcast=True) - hi = hi.to(tl.uint16, bitcast=True) - lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) - hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) - return lo | (hi << 16) - - -@triton.jit -def _unpack_bf16_pairs(values): - lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) - hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) - return tl.interleave(lo, hi) - - -@triton.jit -def _wait_pairs(ptr, offsets, mask): - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) - while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: - values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) - return values - - -@triton.jit -def _publish_input_kernel( - stage_ptr, - peer_ptrs, - peer_offset_u32, - stride_stage_t, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Publish this rank's full partial row into every shard owner's slots.""" - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - elem = tl.arange(0, CS_P2) - elem_mask = elem < CSS - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) - # Address generation is independent of the preceding stream kernel. The - # acquire remains before stage/peer loads, which may consume its writes. - if USE_PDL: - tl.extra.cuda.gdc_wait() - - for owner in tl.static_range(WORLD): - values = tl.load( - stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, - mask=elem_mask, - other=0.0, - ) - packed = _pack_bf16_pairs(values) - base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) - dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair - tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _reduce_insert_kernel( - input_peer_ptrs, - input_peer_offset_u32, - cache_ptr, - slot_ptr, - stride_cache_block, - stride_cache_head, - stride_cache_token, - stride_cache_dim, - block_size, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - CACHE_OFFSET: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Consume all partials for this rank and insert the reduced cache row.""" - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - source = tl.arange(0, WORLD) - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - offsets = ( - (token * WORLD + source)[:, None] * (CS // 2) - + split * (CSS // 2) - + pair[None, :] - ) - mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] - input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) - input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) - input_u32 += input_peer_offset_u32 - # Slot metadata and the cache destination do not depend on input publish. - slot = tl.load(slot_ptr + token) - valid = slot >= 0 - channel = tl.arange(0, CS_P2) - channel_mask = channel < CSS - global_channel = split * CSS + channel - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) - dim = CACHE_OFFSET + global_channel % HEAD_SIZE - safe_slot = tl.maximum(slot, 0).to(tl.int64) - dst = ( - cache_ptr - + (safe_slot // block_size) * stride_cache_block - + head * stride_cache_head - + (safe_slot % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - if USE_PDL: - tl.extra.cuda.gdc_wait() - packed = _wait_pairs(input_u32, offsets, mask) - - lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) - hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) - reduced = tl.interleave( - tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), - tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), - ) - tl.store( - input_u32 + offsets, - tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), - mask=mask, - ) - - tl.store(dst, reduced, mask=valid & channel_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _sconv_publish_kernel( - peer_ptrs, - peer_offset_u32, - residual_ptr, - weight_ptr, - cache_ptr, - position_ptr, - sequence_ptr, - slot_ptr, - block_table_ptr, - stride_residual_t, - stride_cache_block, - stride_cache_head, - stride_cache_token, - stride_cache_dim, - stride_block_table_r, - max_blocks, - block_size, - C: tl.constexpr, - CS: tl.constexpr, - CS_P2: tl.constexpr, - SPLITS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - CACHE_OFFSET: tl.constexpr, - RANK: tl.constexpr, - WORLD: tl.constexpr, - WINDOW: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Compute this rank's shard, then publish it to every rank.""" - tl.static_assert(WINDOW == 4) - token = tl.program_id(0).to(tl.int64) - split = tl.program_id(1).to(tl.int64) - CSS: tl.constexpr = CS // SPLITS - channel = tl.arange(0, CS_P2) - channel_mask = channel < CSS - global_channel = split * CSS + channel - head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) - dim = CACHE_OFFSET + global_channel % HEAD_SIZE - slot = tl.load(slot_ptr + token) - valid = slot >= 0 - position = tl.load(position_ptr + token) - sequence = tl.load(sequence_ptr + token) - safe_slot = tl.maximum(slot, 0).to(tl.int64) - own_ptr = ( - cache_ptr - + (safe_slot // block_size) * stride_cache_block - + head * stride_cache_head - + (safe_slot % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - # These loads were made visible before reduce/insert could signal us and - # are independent of its cache store. Hoisting them gives PDL useful work - # to overlap while retaining the acquire before every cache read. - residual = tl.load( - residual_ptr + token * stride_residual_t + RANK * CS + global_channel, - mask=channel_mask, - other=0.0, - ) - weight0 = tl.load( - weight_ptr + global_channel * WINDOW, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight1 = tl.load( - weight_ptr + global_channel * WINDOW + 1, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight2 = tl.load( - weight_ptr + global_channel * WINDOW + 2, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - weight3 = tl.load( - weight_ptr + global_channel * WINDOW + 3, - mask=channel_mask, - other=0.0, - ).to(tl.float32) - ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) - if USE_PDL: - tl.extra.cuda.gdc_wait() - current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) - - conv = tl.zeros([CS_P2], tl.float32) - for tap_idx in tl.static_range(WINDOW): - source_position = position - (WINDOW - 1) + tap_idx - take = valid & (source_position >= 0) - if tap_idx == WINDOW - 1: - value = tl.where(take, current.to(tl.float32), 0.0) - else: - safe_position = tl.maximum(source_position, 0) - logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) - physical_block = tl.load( - block_table_ptr + sequence * stride_block_table_r + logical_block, - mask=take, - other=0, - ).to(tl.int64) - source_ptr = ( - cache_ptr - + physical_block * stride_cache_block - + head * stride_cache_head - + (safe_position % block_size) * stride_cache_token - + dim * stride_cache_dim - ) - cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) - value = tl.where(take, cached.to(tl.float32), 0.0) - if tap_idx == 0: - weight = weight0 - elif tap_idx == 1: - weight = weight1 - elif tap_idx == 2: - weight = weight2 - else: - weight = weight3 - conv += value * weight - - # Preserve both bf16 rounding points of the original sublayer. - short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) - output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( - tl.bfloat16 - ) - - pair = tl.arange(0, CS_P2 // 2) - pair_mask = pair < CSS // 2 - packed = _pack_bf16_pairs(output) - row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair - - for destination in tl.static_range(WORLD): - base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) - tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _gather_norm_kernel( - output_peer_ptrs, - output_peer_offset_u32, - norm_weight_ptr, - normed_ptr, - residual_out_ptr, - eps, - stride_output_t, - C: tl.constexpr, - C_P2: tl.constexpr, - RANK: tl.constexpr, - HAS_NORM: tl.constexpr, - USE_PDL: tl.constexpr, - launch_pdl: tl.constexpr, -): - """Consume a complete row; one CTA owns all outputs for one token.""" - token = tl.program_id(0).to(tl.int64) - pair = tl.arange(0, C_P2 // 2) - pair_mask = pair < C // 2 - output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) - output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) - output_u32 += output_peer_offset_u32 - offsets = token * (C // 2) + pair - channel = tl.arange(0, C_P2) - channel_mask = channel < C - # Gamma is independent of the preceding sconv publication. Keep the - # acquire immediately before polling the Lamport output slots. - if HAS_NORM: - weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) - if USE_PDL: - tl.extra.cuda.gdc_wait() - packed = _wait_pairs(output_u32, offsets, pair_mask) - row = _unpack_bf16_pairs(packed) - tl.store( - residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask - ) - if HAS_NORM: - row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) - inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) - tl.store( - normed_ptr + token * stride_output_t + channel, - (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), - mask=channel_mask, - ) - tl.store( - output_u32 + offsets, - tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), - mask=pair_mask, - ) - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _validate_lamport_init_kernel( - input_peer_ptrs, - output_peer_ptrs, - bad_ptr, - num_pairs, - RANK: tl.constexpr, - BLOCK: tl.constexpr, -): - """Validate both complete local allocations through their fabric pointers.""" - offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < num_pairs - ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) - ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) - local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) - local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) - value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) - value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) - bad = tl.max( - tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) - ) - if bad != 0: - tl.atomic_max(bad_ptr, 1) - - -class LamportRSConv: - """Persistent symmetric buffers for one TP group.""" - - def _initialize_lamport_buffers(self) -> None: - """Arm every slot and collectively verify the exact sentinel bits. - - Do not replace this with a zero-fill: Lamport readiness distinguishes - the bf16 bit pattern for -0.0 (0x8000) from every published payload. - The validation is deliberately collective so one rank cannot enter the - first polling kernel while another rank still has an unarmed buffer. - """ - if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): - raise RuntimeError("Lamport symmetric buffers must be contiguous") - if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: - raise RuntimeError("Lamport buffers must contain whole uint32 pairs") - - # Fill through int16 so each bf16 lane receives the exact -0.0 bits. - # This covers all three generations and all max-token slots, including - # slots that a smaller first invocation does not touch. - self.buf_in.view(torch.int16).fill_(-0x8000) - self.buf_out.view(torch.int16).fill_(-0x8000) - torch.accelerator.synchronize(self.device) - self.tp.barrier() - - # Scan the complete local allocations. Since every rank performs the - # scan and participates in MAX, success proves every symmetric backing - # allocation was armed before any rank is allowed to use generation 0. - bad = torch.logical_or( - self.buf_in.view(torch.int16).ne(-0x8000).any(), - self.buf_out.view(torch.int16).ne(-0x8000).any(), - ).to(dtype=torch.float32) - bad = self.tp.all_reduce(bad) - if int(bad.item()) != 0: - raise RuntimeError("Lamport sentinel initialization failed on a TP rank") - self.tp.barrier() - - def _initialize_mnnvl_buffers(self) -> None: - """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" - self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) - self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) - torch.accelerator.synchronize(self.device) - self.tp.barrier() - - bad = torch.zeros((), dtype=torch.int32, device=self.device) - num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 - _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( - self.input_peer_ptrs, - self.output_peer_ptrs, - bad, - num_pairs, - RANK=self.rank, - BLOCK=256, - num_warps=4, - ) - bad = self.tp.all_reduce(bad.to(torch.float32)) - if int(bad.item()) != 0: - raise RuntimeError("MNNVL Lamport sentinel initialization failed") - self.tp.barrier() - - def __init__( - self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS - ) -> None: - import torch.distributed._symmetric_memory as symm_mem - - tp = get_tp_group() - self.tp = tp - self.group = tp.device_group - self.world_size = tp.world_size - self.rank = tp.rank_in_group - self.device = torch.device(tp.device) - if self.world_size not in (2, 4, 8): - raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") - if hidden_size % (2 * self.world_size) != 0: - raise ValueError("hidden size must produce an even shard on every rank") - if window_size != 4: - raise ValueError(f"short-conv window size must be 4, got {window_size}") - if max_tokens < 1 or max_tokens > _MAX_TOKENS: - raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") - - is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) - self.hidden_size = hidden_size - self.window_size = window_size - self.max_tokens = max_tokens - self.shard_size = hidden_size // self.world_size - # Three generations follow FlashInfer's Lamport layout. A generation - # is reused only after two intervening collective calls. - self.num_buffers = 3 - self.input_generation_bytes = max_tokens * hidden_size * 2 - self.output_generation_bytes = max_tokens * hidden_size * 2 - self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 - if is_cross_node: - try: - from flashinfer.comm.mnnvl import ( - McastGPUBuffer, - TorchDistBackend, - is_mnnvl_fabric_supported, - ) - except ImportError as error: - raise RuntimeError( - "cross-node TP requires FlashInfer MNNVL support" - ) from error - - local_supported = int( - is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) - ) - unsupported = torch.tensor( - 1 - local_supported, dtype=torch.float32, device=self.device - ) - unsupported = tp.all_reduce(unsupported) - if int(unsupported.item()) != 0: - raise RuntimeError("cross-node TP is supported only on MNNVL fabric") - - comm_backend = TorchDistBackend(self.group) - allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 - self._mnnvl_input_handle = McastGPUBuffer( - allocation_bytes, - self.world_size, - self.rank, - self.device, - comm_backend, - ) - self._mnnvl_output_handle = McastGPUBuffer( - allocation_bytes, - self.world_size, - self.rank, - self.device, - comm_backend, - ) - self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() - self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() - self._initialize_mnnvl_buffers() - logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") - else: - self.buf_in = symm_mem.empty( - self.num_buffers, - max_tokens, - self.world_size, - self.shard_size, - dtype=torch.bfloat16, - device=self.device, - ) - self.buf_out = symm_mem.empty( - self.num_buffers, - max_tokens, - hidden_size, - dtype=torch.bfloat16, - device=self.device, - ) - group_name = self.group.group_name - input_handle = symm_mem.rendezvous(self.buf_in, group_name) - output_handle = symm_mem.rendezvous(self.buf_out, group_name) - self.input_peer_ptrs = input_handle.buffer_ptrs_dev - self.output_peer_ptrs = output_handle.buffer_ptrs_dev - self._input_handle = input_handle - self._output_handle = output_handle - self._initialize_lamport_buffers() - self.generation = 0 - - def usable(self, num_tokens: int) -> bool: - return 0 < num_tokens <= self.max_tokens - - def rs_sconv_ag_add_norm( - self, - input_tensor: torch.Tensor, - residual: torch.Tensor, - conv_weight: torch.Tensor, - norm_weight: torch.Tensor | None, - eps: float, - cache: torch.Tensor, - positions: torch.Tensor, - block_table: torch.Tensor, - seq_idx: torch.Tensor, - slot_mapping: torch.Tensor, - off_s: int, - ws: int, - block_size: int, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" - tokens, hidden_size = residual.shape - if not self.usable(tokens): - raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") - if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: - raise ValueError("residual must be bf16 [T, 6144]") - if ( - input_tensor.shape != residual.shape - or input_tensor.dtype != torch.bfloat16 - or input_tensor.stride(1) != 1 - ): - raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") - shard_size = hidden_size // self.world_size - if conv_weight.shape != (shard_size, self.window_size): - raise ValueError( - f"conv_weight must have shape [{shard_size}, {self.window_size}]" - ) - if ( - conv_weight.dtype != torch.bfloat16 - or conv_weight.stride(0) != self.window_size - ): - raise ValueError("conv_weight must be contiguous bf16") - if norm_weight is not None and ( - norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 - ): - raise ValueError("norm_weight must be bf16 [6144] or None") - if cache.dtype != torch.bfloat16 or cache.ndim != 4: - raise ValueError("cache must be a 4-D bf16 tensor") - if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: - raise ValueError("cache head layout is inconsistent with ws") - if off_s < 0 or off_s + ws > cache.shape[3]: - raise ValueError("cache channel offset is out of bounds") - - index = self.generation - input_offset = index * self.input_generation_bytes // 4 - output_offset = index * self.output_generation_bytes // 4 - normed = torch.empty_like(residual) if norm_weight is not None else None - residual_out = torch.empty_like(residual) - phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 - phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) - phase_grid = (tokens, phase_splits) - # Wide CTAs win before the grid saturates; smaller CTAs reduce register - # pressure once high-throughput batches provide enough parallelism. - if 128 <= tokens <= 2048: - phase_warps = 16 - elif tokens > 2048: - phase_warps = 8 - else: - phase_warps = 4 - gather_warps = 4 if tokens >= 256 else 8 - _publish_input_kernel[phase_grid]( - input_tensor, - self.input_peer_ptrs, - input_offset, - input_tensor.stride(0), - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - RANK=self.rank, - WORLD=self.world_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _reduce_insert_kernel[phase_grid]( - self.input_peer_ptrs, - input_offset, - cache, - slot_mapping, - cache.stride(0), - cache.stride(1), - cache.stride(2), - cache.stride(3), - block_size, - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - HEAD_SIZE=ws, - CACHE_OFFSET=off_s, - RANK=self.rank, - WORLD=self.world_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _sconv_publish_kernel[phase_grid]( - self.output_peer_ptrs, - output_offset, - residual, - conv_weight, - cache, - positions, - seq_idx, - slot_mapping, - block_table, - residual.stride(0), - cache.stride(0), - cache.stride(1), - cache.stride(2), - cache.stride(3), - block_table.stride(0), - block_table.shape[1], - block_size, - C=hidden_size, - CS=shard_size, - CS_P2=phase_tile_p2, - SPLITS=phase_splits, - HEAD_SIZE=ws, - CACHE_OFFSET=off_s, - RANK=self.rank, - WORLD=self.world_size, - WINDOW=self.window_size, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=phase_warps, - ) - _gather_norm_kernel[(tokens,)]( - self.output_peer_ptrs, - output_offset, - norm_weight if norm_weight is not None else residual, - normed if normed is not None else residual_out, - residual_out, - eps, - residual.stride(0), - C=hidden_size, - C_P2=triton.next_power_of_2(hidden_size), - RANK=self.rank, - HAS_NORM=norm_weight is not None, - USE_PDL=self.use_pdl, - launch_pdl=self.use_pdl, - num_warps=gather_warps, - ) - self.generation = (index + 1) % self.num_buffers - return normed, residual_out - - -_STATE: LamportRSConv | None = None -_STATE_FAILED = False - - -def initialize_lamport_rs_conv( - hidden_size: int, window_size: int, max_num_batched_tokens: int -) -> None: - """Collectively initialize the TP-group state during model construction.""" - global _STATE, _STATE_FAILED - if _STATE is not None: - if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: - raise RuntimeError("all Lamport users must share hidden and window sizes") - return - if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": - return - try: - max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) - _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) - except Exception: - _STATE_FAILED = True - logger.exception("fused collective unavailable; use the NCCL fallback") - - -def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: - """Return the state initialized with the model, or ``None`` for fallback.""" - if _STATE is not None and ( - _STATE.hidden_size != hidden_size or _STATE.window_size != window_size - ): - raise RuntimeError("all Lamport users must share hidden and window sizes") - return _STATE diff --git a/vllm/models/inkling/amd/ops/mm_towers.py b/vllm/models/inkling/amd/ops/mm_towers.py deleted file mode 100644 index 021cf5274bd2..000000000000 --- a/vllm/models/inkling/amd/ops/mm_towers.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Fused CUDA kernels for the Inkling vision/audio towers. - -Both kernels keep the reference paths' fp32 accumulation and per-op bf16 -rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently -bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order -(real-checkpoint-weight cosine vs reference > 0.9999998). -""" - -from __future__ import annotations - -import torch - -from vllm.triton_utils import tl, tldevice, triton - -from .norm import _get_num_warps_from_block_size - - -@triton.jit -def _dmel_embed_sum_norm_kernel( - idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) - w_ptr, # [NB * VOCAB, D] bf16 embedding table - norm_w_ptr, # [D] (unused if HAS_NORM=False) - out_ptr, # [T, D] bf16 - eps, - D, - stride_idx_t, - NB: tl.constexpr, - VOCAB: tl.constexpr, - D_P2: tl.constexpr, - HAS_NORM: tl.constexpr, -): - t = tl.program_id(0).to(tl.int64) - offs = tl.arange(0, D_P2) - mask = offs < D - # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), - # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). - acc = tl.zeros([D_P2], dtype=tl.float32) - for b in tl.static_range(NB): - v = tl.load(idx_ptr + t * stride_idx_t + b) - row = (b * VOCAB + v).to(tl.int64) - acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) - h = acc.to(tl.bfloat16) - if HAS_NORM: - # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded - # bf16 multiply with the bf16 weight. - x32 = h.to(tl.float32) - var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D - xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) - w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) - h = xn * w - tl.store(out_ptr + t * D + offs, h, mask=mask) - - -def dmel_embed_sum_norm( - dmel_idx: torch.Tensor, # [T, NB] int32 - weight: torch.Tensor, # [NB * VOCAB, D] bf16 - norm_weight: torch.Tensor | None, - eps: float, -) -> torch.Tensor: - """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no - [T, NB, D] intermediate).""" - T, nb = dmel_idx.shape - D = weight.shape[1] - assert weight.shape[0] % nb == 0 - vocab = weight.shape[0] // nb - out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) - if T == 0: - return out - d_p2 = triton.next_power_of_2(D) - _dmel_embed_sum_norm_kernel[(T,)]( - dmel_idx, - weight, - norm_weight if norm_weight is not None else weight, - out, - eps, - D, - dmel_idx.stride(0), - NB=nb, - VOCAB=vocab, - D_P2=d_p2, - HAS_NORM=norm_weight is not None, - # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% - # at large T (the 80-row gather chain is latency- not lane-bound). - num_warps=8, - ) - return out - - -@triton.jit -def _rmsnorm_gelu_kernel( - x_ptr, # [R, D] bf16 - w_ptr, # [D] - out_ptr, # [R, D] bf16 (or the folded layout when FOLD) - eps, - R, - D, - D_P2: tl.constexpr, - BLOCK_M: tl.constexpr, # rows per block (>1 for small D) - HAS_GELU: tl.constexpr, - FOLD: tl.constexpr, - # fold geometry: input rows index [N, T, H, W]; the store scatters each - # row to (out_row, slot) of fold_timespace_to_depth's output layout. - FT: tl.constexpr, - FH: tl.constexpr, - FW: tl.constexpr, - TF: tl.constexpr, - HF: tl.constexpr, -): - pid = tl.program_id(0).to(tl.int64) - rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) - rmask = rows < R - offs = tl.arange(0, D_P2) - mask = rmask[:, None] & (offs < D)[None, :] - x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( - tl.float32 - ) - var = tl.sum(x32 * x32, axis=1) / D - xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) - w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) - h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm - if HAS_GELU: - # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, - # matching F.gelu's opmath on a bf16 tensor. - g32 = h.to(tl.float32) - h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) - if FOLD: - # Store directly into the next layer's folded layout (a pure - # permutation — replaces the separate fold copy pass). - t = (rows // (FH * FW)) % FT - hh = (rows // FW) % FH - ww = rows % FW - n = rows // (FT * FH * FW) - slot = ((t % TF) * HF + hh % HF) * HF + ww % HF - out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( - FW // HF - ) + ww // HF - base = (out_row * (TF * HF * HF) + slot) * D - tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) - else: - tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) - - -def rmsnorm_gelu( - x: torch.Tensor, # [..., D] bf16 contiguous - weight: torch.Tensor, - eps: float, - gelu: bool = True, - fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold -) -> torch.Tensor: - """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block - when D is small. With ``fold``, x must be [N, T, H, W, D] and the output - comes back as ``fold_timespace_to_depth(result, *fold)``.""" - D = x.shape[-1] - flat = x.reshape(-1, D) - assert flat.stride(1) == 1 and flat.stride(0) == D - R = flat.shape[0] - if fold is None: - out = torch.empty_like(flat) - ft = fh = fw = tf = hf = 1 - out_shape = x.shape - else: - tf, hf = fold - N, ft, fh, fw, _ = x.shape - out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) - out = torch.empty(out_shape, dtype=x.dtype, device=x.device) - if R == 0: - return out.reshape(out_shape) - d_p2 = triton.next_power_of_2(D) - block_m = max(1, 4096 // d_p2) - _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( - flat, - weight, - out, - eps, - R, - D, - D_P2=d_p2, - BLOCK_M=block_m, - HAS_GELU=gelu, - FOLD=fold is not None, - FT=ft, - FH=fh, - FW=fw, - TF=tf, - HF=hf, - num_warps=_get_num_warps_from_block_size(d_p2 * block_m), - ) - return out.reshape(out_shape) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index c4f0092f8d9b..6a26ce78b4c4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -532,8 +532,10 @@ def forward( symm_buffer, activation_clamp=activation_clamp, activation=self.activation, - situ_beta=self.activation_beta, - situ_linear_beta=self.activation_linear_beta, + # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the + # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation_alpha=self.activation_beta or 1.0, + activation_beta=self.activation_linear_beta or 0.0, fast_math=fast_math, ) return y diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py index 82782fcb96d2..7baec0b42b44 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py @@ -587,34 +587,3 @@ def warp_sum_specialized( mask_and_clamp=31, ) return value - - -@cute.jit -def block_sum_specialized( - value: Float32, - warp_sums: cute.Tensor, - tidx: Int32, - warps: cutlass.Constexpr[int], - last_warp_lanes: cutlass.Constexpr[int], - last_warp_mask: cutlass.Constexpr[int], -) -> Float32: - """Upstream-equivalent FP32 block reduction.""" - - lane = cute.arch.lane_idx() - warp_idx = cute.arch.warp_idx() - value = warp_sum_specialized( - value, warp_idx, lane, warps, last_warp_lanes, last_warp_mask - ) - if lane == 0: - warp_sums[warp_idx] = value - cute.arch.barrier() - - block_sum = Float32(0.0) - if warp_idx == 0: - if lane < Int32(warps): - block_sum = warp_sums[lane] - block_sum = cute.arch.warp_reduction_sum(block_sum) - if lane == 0: - warp_sums[0] = block_sum - cute.arch.barrier() - return warp_sums[0] diff --git a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py index 836649b725b4..df7abc0760db 100644 --- a/vllm/models/minimax_m3/amd/ops/swiglu_oai.py +++ b/vllm/models/minimax_m3/amd/ops/swiglu_oai.py @@ -83,6 +83,8 @@ def _swiglu_oai_quant_kernel( limit, HAS_LIMIT: tl.constexpr, BLOCK_M: tl.constexpr, + FP8_MAX: tl.constexpr, + TINY: tl.constexpr, ): """SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0) quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one @@ -109,8 +111,8 @@ def _swiglu_oai_quant_kernel( gate = tl.minimum(gate, limit) up = tl.minimum(tl.maximum(up, -limit), limit) act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32 - amax = tl.maximum(tl.max(tl.abs(act), axis=1), 1e-30) # [BLOCK_M] - sb = tl.minimum(tl.maximum(tl.floor(tl.log2(amax)) + 127.0, 0.0), 254.0) + amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY) # [BLOCK_M] + sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0) descale = tl.exp2(sb - 127.0) aq = (act / descale[:, None]).to(aq_ptr.dtype.element_ty) tl.store( @@ -173,6 +175,8 @@ def swiglu_oai_quantize_mxfp8( 0.0 if limit is None else float(limit), HAS_LIMIT=limit is not None, BLOCK_M=block_m, + FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max), + TINY=float(torch.finfo(torch.float32).tiny), num_warps=4, ) return aq, asc diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index cf83bb996e63..f099d7752ae5 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -197,6 +197,8 @@ def is_reasoning_end(self, input_ids: list[int]) -> bool: Used by structured engines like `xgrammar` to check if the reasoning content ends in the model output. + Must be a pure function of the input_ids. + Args: input_ids: The token IDs of the model output. @@ -210,6 +212,8 @@ def is_reasoning_end_streaming( """ Check if the reasoning content ends during a decode step. + Must be a pure function of the input_ids and delta_ids. + Args: input_ids: The entire model output token IDs. delta_ids: The last few computed tokens at the current decode step. @@ -228,8 +232,8 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: Returns: The offset within ``token_ids`` of the last reasoning token, or - ``None`` when reasoning does not end in this window. Parsers that - cannot answer from a window alone always return ``None``; callers + ``len(token_ids)`` when reasoning does not end in this window. Parsers + that cannot answer from a window alone always return ``None``; callers fall back to :meth:`is_reasoning_end_streaming`. """ return None diff --git a/vllm/parser/deepseek_v32.py b/vllm/parser/deepseek_v32.py index f76c954a93d6..de6634cf7368 100644 --- a/vllm/parser/deepseek_v32.py +++ b/vllm/parser/deepseek_v32.py @@ -53,6 +53,7 @@ def deepseek_v32_config() -> ParserEngineConfig: return ParserEngineConfig( name="deepseek_v32", initial_state=ParserState.CONTENT, + wait_for_reasoning=False, terminals={ "TOOL_START": DSML_FUNC_START, "TOOL_END": DSML_FUNC_END, diff --git a/vllm/parser/deepseek_v4.py b/vllm/parser/deepseek_v4.py index 13eb83e6b804..79ddd93f6b2e 100644 --- a/vllm/parser/deepseek_v4.py +++ b/vllm/parser/deepseek_v4.py @@ -140,6 +140,7 @@ def deepseek_v4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( name="deepseek_v4", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ "THINK_START": DSML_THINK_START, "THINK_END": DSML_THINK_END, diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 4817dedb4663..fa306a5da4a1 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -645,27 +645,31 @@ def find_reasoning_end_offset(self, token_ids: Sequence[int]) -> int | None: for offset, token_id in enumerate(token_ids): if token_id in end_ids: return offset - return None + return len(token_ids) def is_reasoning_end(self, input_ids: list[int]) -> bool: - end_id = self._reasoning_end_token_id + config = self.parser_engine_config + wait_for_reasoning = config.wait_for_reasoning + if wait_for_reasoning is None: + wait_for_reasoning = config.initial_state is ParserState.REASONING + start_transition = config.transitions.get((config.initial_state, "THINK_START")) + start_opens_reasoning = ( + start_transition is not None + and start_transition.next_state is ParserState.REASONING + ) + end_ids = self._reasoning_end_token_ids start_id = self._reasoning_start_token_id - if end_id is not None: - if not input_ids: - return self.parser_engine_config.initial_state != ParserState.REASONING - boundary_ids = self._turn_boundary_token_ids - for i in range(len(input_ids) - 1, -1, -1): - token_id = input_ids[i] - if token_id == end_id: - return True - if start_id is not None and token_id == start_id: + boundary_ids = self._turn_boundary_token_ids + for token_id in reversed(input_ids): + if token_id in end_ids: + return True + if token_id == start_id: + if start_opens_reasoning: return False - if token_id in boundary_ids: - return ( - self.parser_engine_config.initial_state != ParserState.REASONING - ) - return False - return self._reasoning_ended + break + if token_id in boundary_ids: + break + return not wait_for_reasoning def extract_content_ids(self, input_ids: list[int]) -> list[int]: end_id = self._reasoning_end_token_id diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index 5c5cd2b467fd..b8d7dc6e465a 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -74,6 +74,8 @@ class ParserEngineConfig: initial_state: ParserState = ParserState.CONTENT + wait_for_reasoning: bool | None = None + arg_converter: Callable[[str, bool], str] | None = None stream_arg_deltas: bool = True diff --git a/vllm/parser/gemma4.py b/vllm/parser/gemma4.py index da1718332b21..e3baaf488a8d 100644 --- a/vllm/parser/gemma4.py +++ b/vllm/parser/gemma4.py @@ -293,16 +293,21 @@ def _gemma4_arg_converter(raw_args: str, partial: bool) -> str: @functools.cache -def gemma4_config() -> ParserEngineConfig: +def gemma4_config(thinking: bool = False) -> ParserEngineConfig: return ParserEngineConfig( name="gemma4", initial_state=ParserState.CONTENT, + # Decides a silent tail: the E-series HF templates write no marker + # with thinking off, and no template does after a tool response. + wait_for_reasoning=thinking, + turn_boundary_tokens=frozenset({"<|turn>", "<|tool_response>"}), terminals={ "THINK_START": CHANNEL_START, "THINK_END": CHANNEL_END, "TOOL_START": TOOL_CALL_START, "TOOL_END": TOOL_CALL_END, "CALL_PREFIX": "call:", + "COLON": ":", "OPEN_BRACE": "{", "OPEN_PAREN": "(", }, @@ -348,6 +353,13 @@ def gemma4_config() -> ParserEngineConfig: ParserState.TOOL_NAME, (), ), + # Bare opener: some checkpoints emit "<|tool_call>:name{...}" + # without the "call" prefix. Match the colon directly so the + # tool call is not silently dropped. + (ParserState.TOOL_PREAMBLE, "COLON"): Transition( + ParserState.TOOL_NAME, + (), + ), (ParserState.TOOL_NAME, "OPEN_BRACE"): Transition( ParserState.TOOL_ARGS, (), @@ -409,11 +421,11 @@ def __init__( **kwargs, ) -> None: chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {} - self._thinking_enabled = chat_kwargs.get("enable_thinking", False) + thinking = chat_kwargs.get("enable_thinking", False) super().__init__( tokenizer, tools, - parser_engine_config=gemma4_config(), + parser_engine_config=gemma4_config(thinking=thinking), **kwargs, ) vocab = self.vocab @@ -464,30 +476,6 @@ def _preprocess_feed( return delta_text, delta_token_ids - def is_reasoning_end(self, input_ids: list[int]) -> bool: - end_id = self._reasoning_end_token_id - start_id = self._reasoning_start_token_id - tool_call_id = self._tool_call_token_id - new_turn_id = self._new_turn_token_id - tool_response_id = self._tool_response_token_id - - if end_id is not None and not input_ids: - return self.parser_engine_config.initial_state != ParserState.REASONING - - for i in range(len(input_ids) - 1, -1, -1): - tid = input_ids[i] - if start_id is not None and tid == start_id: - return False - if tool_call_id is not None and tid == tool_call_id: - return True - if new_turn_id is not None and tid == new_turn_id: - return not self._thinking_enabled - if tool_response_id is not None and tid == tool_response_id: - return not self._thinking_enabled - if end_id is not None and tid == end_id: - return True - return True - def _prompt_ends_in_open_reasoning(self, prompt_token_ids: Sequence[int]) -> bool: """Whether the prompt tail is inside an open ``<|channel>`` block. diff --git a/vllm/parser/glm47_moe.py b/vllm/parser/glm47_moe.py index 8aa4feef2599..4a804a4c9c94 100644 --- a/vllm/parser/glm47_moe.py +++ b/vllm/parser/glm47_moe.py @@ -123,6 +123,7 @@ def glm47_moe_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="glm47_moe", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_START": TOOL_CALL_START, @@ -206,11 +207,6 @@ def _handle_tool_end(self, event, deltas) -> None: self._tool_slots[idx].name = self._tool_slots[idx].name.strip() super()._handle_tool_end(event, deltas) - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if not self.thinking_enabled: - return True - return super().is_reasoning_end(input_ids) - def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py index 16883b58aa7f..2081fb278464 100644 --- a/vllm/parser/inkling.py +++ b/vllm/parser/inkling.py @@ -282,6 +282,7 @@ def inkling_config() -> ParserEngineConfig: # `<|message_model|>`. Non-streaming parsing receives only the generated # suffix, so begin in the corresponding message-header state as well. initial_state=ParserState.MESSAGE_HEADER, + wait_for_reasoning=True, terminals=terminals, # Inkling content-kind markers are the grammar. When the engine is # used through DelegatingParser, the reasoning pass can hand the tool diff --git a/vllm/parser/kimi_k2.py b/vllm/parser/kimi_k2.py index 1e1067a8c2ee..7c5f30aafef4 100644 --- a/vllm/parser/kimi_k2.py +++ b/vllm/parser/kimi_k2.py @@ -81,6 +81,7 @@ def kimi_k2_config(thinking: bool = True) -> ParserEngineConfig: return ParserEngineConfig( name="kimi_k2", initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, terminals={ **reasoning_terminals, "TOOL_SECTION_START": TOOL_SECTION_START, @@ -236,24 +237,6 @@ def _handle_arg_chunk(self, event, deltas) -> None: def _extract_args_json(self, raw_args: str, func_name: str) -> str: return raw_args.strip() or "{}" - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if not self.thinking_enabled: - return True - - start_id = self._start_token_id - end_id = self._end_token_id - tool_section_id = self._tool_section_start_token_id - - for i in range(len(input_ids) - 1, -1, -1): - token_id = input_ids[i] - if start_id is not None and token_id == start_id: - return False - if end_id is not None and token_id == end_id: - return True - if tool_section_id is not None and token_id == tool_section_id: - return True - return False - def extract_content_ids(self, input_ids: list[int]) -> list[int]: if not self.thinking_enabled: return input_ids diff --git a/vllm/parser/minimax_m2.py b/vllm/parser/minimax_m2.py index ab1b7eb11427..c96f852b9df6 100644 --- a/vllm/parser/minimax_m2.py +++ b/vllm/parser/minimax_m2.py @@ -95,6 +95,7 @@ def minimax_m2_config() -> ParserEngineConfig: return ParserEngineConfig( name="minimax_m2", initial_state=ParserState.REASONING, + wait_for_reasoning=True, terminals={ "THINK_START": THINK_START, "THINK_END": THINK_END, diff --git a/vllm/parser/mistral.py b/vllm/parser/mistral.py index fc9b8fbefe37..ce516eea1419 100644 --- a/vllm/parser/mistral.py +++ b/vllm/parser/mistral.py @@ -193,6 +193,7 @@ def mistral_config( return ParserEngineConfig( name=name, initial_state=ParserState.CONTENT, + wait_for_reasoning=reasoning_encoding != "none", terminals={ **reasoning_terminals, "TOOL_CALLS": _TOOL_CALLS, @@ -891,24 +892,6 @@ def _split_delta( return (delta_to_be_parsed, "") - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if self._reasoning_encoding == "none": - return True - if super().is_reasoning_end(input_ids): - return True - # [TOOL_CALLS] acts as an implicit reasoning-end marker - if self.bot_token_id is not None: - reasoning_start_id = self._reasoning_start_token_id - for i in range(len(input_ids) - 1, -1, -1): - if ( - reasoning_start_id is not None - and input_ids[i] == reasoning_start_id - ): - return False - if input_ids[i] == self.bot_token_id: - return True - return False - def extract_reasoning( self, model_output: str, diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index 8edf42e26363..a11551a209a0 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -100,6 +100,7 @@ def qwen3_config( return ParserEngineConfig( name=name, initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, + wait_for_reasoning=thinking, turn_boundary_tokens=turn_boundary_tokens, terminals={ # Reasoning terminals @@ -205,8 +206,8 @@ class Qwen3Parser(ParserEngine): """Qwen3 parser: ````/```` reasoning + ```` XML tool calls in a single engine. - - ```` as implicit reasoning end - - Unpaired ```` token ID detection for ``is_reasoning_end`` + - ```` as implicit reasoning end (a grammar transition, so it + also feeds ``is_reasoning_end`` and the structured-output gate) Subclasses that share the grammar but differ only in the four wrapper token strings (reasoning + tool-call) override the class attributes @@ -245,9 +246,6 @@ def __init__( tools, **kwargs, ) - vocab = self.vocab - self._tool_call_token_id: int | None = vocab.get(self.TOOL_START) - self._tool_call_end_token_id: int | None = vocab.get(self.TOOL_END) def extract_reasoning( self, @@ -257,28 +255,3 @@ def extract_reasoning( if not self.thinking_enabled: return None, model_output return super().extract_reasoning(model_output, request) - - def is_reasoning_end(self, input_ids: list[int]) -> bool: - if super().is_reasoning_end(input_ids): - return True - tool_call_id = self._tool_call_token_id - tool_call_end_id = self._tool_call_end_token_id - reasoning_start_id = self._reasoning_start_token_id - boundary_ids = self._turn_boundary_token_ids - if tool_call_id is not None: - for i in range(len(input_ids) - 1, -1, -1): - if ( - reasoning_start_id is not None - and input_ids[i] == reasoning_start_id - ): - return False - if input_ids[i] in boundary_ids: - return False - if input_ids[i] == tool_call_id: - if tool_call_end_id is not None and any( - input_ids[j] == tool_call_end_id - for j in range(i + 1, len(input_ids)) - ): - continue - return True - return False diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index ecd2dc5fe623..89b68a401e26 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -605,26 +605,30 @@ def import_kernels(cls) -> None: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once( + "Failed to import from vllm._C: %s", repr(e) + ) else: try: import vllm._C_AVX512 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: logger.warning_once( - "Failed to import from vllm._C_AVX512: %r", e + "Failed to import from vllm._C_AVX512: %s", repr(e) ) else: try: import vllm._C_AVX2 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: - logger.warning_once("Failed to import from vllm._C_AVX2: %r", e) + logger.warning_once( + "Failed to import from vllm._C_AVX2: %s", repr(e) + ) else: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %s", repr(e)) @classmethod def pack_kv_cache( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 98bddbde296d..30eed82ccb16 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -239,7 +239,9 @@ def import_kernels(cls) -> None: try: import vllm._C_stable_libtorch # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e) + logger.warning_once( + "Failed to import from vllm._C_stable_libtorch: %s", repr(e) + ) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 with contextlib.suppress(ImportError): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 7ad0b42da774..bd5d98191201 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -365,7 +365,7 @@ def import_kernels(cls) -> None: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning_once("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %s", repr(e)) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 diff --git a/vllm/third_party/flash_linear_attention/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py index cc6f54d8af0b..880be7747dcc 100644 --- a/vllm/third_party/flash_linear_attention/ops/index.py +++ b/vllm/third_party/flash_linear_attention/ops/index.py @@ -10,6 +10,7 @@ import torch from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import async_tensor_h2d from vllm.triton_utils import triton from .utils import tensor_cache @@ -27,8 +28,8 @@ def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Te chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() indices = torch.cat([torch.arange(n) for n in chunk_counts]) chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1) - return chunk_indices.to( - device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True + return async_tensor_h2d( + chunk_indices, device=cu_seqlens.device, dtype=cu_seqlens.dtype ) diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index 34e57840bed7..989250e2ea5b 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -10,6 +10,7 @@ from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from openai.types.responses.tool_choice_function import ToolChoiceFunction from xgrammar import StructuralTag, normalize_tool_choice +from xgrammar import builtin_structural_tag as xgrammar_builtin_structural_tag from xgrammar import get_model_structural_tag as get_xgrammar_model_structural_tag from xgrammar.openai_tool_call_schema import ( BuiltinToolParam, @@ -242,9 +243,20 @@ def get_deepseek_v41_structural_tag( token_suffix: str = "", ) -> StructuralTag: # Serving enables this visible-text grammar after the reasoning boundary. + builder = getattr( + xgrammar_builtin_structural_tag, "get_deepseek_v4_1_structural_tag", None + ) + if builder is not None: + return builder( + tools=tools, + builtin_tools=builtin_tools, + tool_choice=tool_choice, + reasoning="disabled", + ) + del builtin_tools, reasoning, token_suffix - # TODO: lower parameter schemas into DSML constraints. This builder constrains + # Compatibility with xgrammar releases without the V4.1 builtin. This constrains # tool names and DSML/value syntax only. Parameter names, presence, uniqueness # and value schemas remain unconstrained, including for strict=true. The request # layer still uses strict to decide whether auto tool choice activates a grammar. diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 3fc70c0f5da3..91e315afb9e7 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -158,7 +158,12 @@ def _missing(*_: Any, **__: Any) -> NoReturn: _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None +_mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( @@ -231,7 +236,10 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl - global _tf32_hc_prenorm_gemm_impl + global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl + global _get_sparse_mqa_logits_metadata_impl + global _get_paged_sparse_mqa_logits_metadata_impl + global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl global _get_theoretical_mk_alignment_for_contiguous_layout_impl @@ -251,6 +259,7 @@ def _lazy_init() -> None: or _fp8_fp4_paged_mqa_logits_impl is not None or _get_paged_mqa_logits_metadata_impl is not None or _tf32_hc_prenorm_gemm_impl is not None + or _mega_mhc_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None or _pack_ue8m0_to_int_impl is not None @@ -280,7 +289,11 @@ def _lazy_init() -> None: _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None) _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None) _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None) - _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) + # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so + # an externally pinned older DeepGEMM still resolves. + _grouped_masked_impl = getattr( + _dg, "m_grouped_fp8_gemm_nt_masked", None + ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None) # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`. @@ -289,7 +302,19 @@ def _lazy_init() -> None: _get_paged_mqa_logits_metadata_impl = getattr( _dg, "get_paged_mqa_logits_metadata", None ) + # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only). + _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None) + _fp8_fp4_paged_sparse_mqa_logits_impl = getattr( + _dg, "fp8_fp4_paged_sparse_mqa_logits", None + ) + _get_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_sparse_mqa_logits_metadata", None + ) + _get_paged_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_paged_sparse_mqa_logits_metadata", None + ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) + _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( _dg, "get_mn_major_tma_aligned_tensor", None ) @@ -674,6 +699,171 @@ def fp8_fp4_paged_mqa_logits( ) +def has_deep_gemm_sparse_mqa() -> bool: + """Whether the installed DeepGEMM provides the sparse-indexer kernels + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" + _lazy_init() + return ( + _fp8_fp4_sparse_mqa_logits_impl is not None + and _fp8_fp4_paged_sparse_mqa_logits_impl is not None + and _get_sparse_mqa_logits_metadata_impl is not None + and _get_paged_sparse_mqa_logits_metadata_impl is not None + ) + + +def get_sparse_mqa_logits_metadata( + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + num_kv_tokens: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`. + + Always uses the unaligned-ks variant, which degenerates to the aligned + variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never + need a host-side sync to pick one. + + Args: + cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV + workspace, shape [num_q_tokens], dtype int32. + num_kv_tokens: Total KV tokens in the packed workspace. + sparse_kv_block_indices: Per-row candidate block ids, shape + [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's + valid prefix must be sorted ascending (repeats are tolerated but + wasteful); pad by repeating the last valid block. Block ``i`` + covers tokens + ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``. + qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4, + ``torch.float8_e4m3fn`` for FP8). + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_sparse_mqa_logits_metadata_impl( + cu_seqlen_ks, + cu_seqlen_ke, + num_kv_tokens, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def get_paged_sparse_mqa_logits_metadata( + context_lens: torch.Tensor, + block_table: torch.Tensor, + indices: torch.Tensor, + page_kv: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`. + + Args: + context_lens: Per-row context length, shape [num_q_tokens], int32. + block_table: Per-row page table, shape [num_q_tokens, max_pages], + int32; rows of queries paired with the same request must be + identical (expand it per query row before calling). + indices: Request index for each query row, shape [num_q_tokens]. + page_kv: Tokens per KV page; must be a multiple of sparse_block_kv. + sparse_kv_block_indices: Per-row logical (context-relative) candidate + block ids, [num_q_tokens, num_max_sparse_blocks], int32. + qk_dtype: dtype of the packed Q values. + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_paged_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_paged_sparse_mqa_logits_metadata_impl( + context_lens, + block_table, + indices, + page_kv, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + ) + + +def fp8_fp4_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (prefill path). + + Args: + q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed + as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H] + int32 (mandatory, unlike the dense kernels). + kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0 + int32 scales. + weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in). + metadata: From `get_sparse_mqa_logits_metadata`. + num_max_sparse_blocks: Width of each sparse-index row. + sparse_block_kv: Tokens per sparse block, 8 or 16. + + Returns: + bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv]; + column ``j * sparse_block_kv + o`` scores the token at + ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``. + """ + _lazy_init() + if _fp8_fp4_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_sparse_mqa_logits_impl( + q, + kv, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def fp8_fp4_paged_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (paged decode path). + + Args: + q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. + kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, + page stride 512B-aligned. + weights: [num_q_tokens, H] ``torch.bfloat16``. + metadata: From `get_paged_sparse_mqa_logits_metadata`. + + Returns: + bf16 logits of shape + [num_q_tokens, num_max_sparse_blocks * sparse_block_kv]. + """ + _lazy_init() + if _fp8_fp4_paged_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_paged_sparse_mqa_logits_impl( + q, + kv_cache, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + ) + + def tf32_hc_prenorm_gemm( x: torch.Tensor, fn: torch.Tensor, @@ -700,6 +890,15 @@ def tf32_hc_prenorm_gemm( ) +def mega_mhc(*args: Any, **kwargs: Any) -> None: + """Run DeepGEMM Mega mHC with caller-owned output tensors.""" + _lazy_init() + if _mega_mhc_impl is None: + _missing() + return + _mega_mhc_impl(*args, **kwargs) + + def _ceil_to_ue8m0(x: torch.Tensor): return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -793,6 +992,7 @@ def should_use_deepgemm_for_fp8_linear( "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", "is_deep_gemm_supported", + "mega_mhc", "get_num_sms", "set_num_sms", "should_use_deepgemm_for_fp8_linear", diff --git a/vllm/utils/gpu_sync_debug.py b/vllm/utils/gpu_sync_debug.py index 53fcdfe71004..8d0e7ae1871f 100644 --- a/vllm/utils/gpu_sync_debug.py +++ b/vllm/utils/gpu_sync_debug.py @@ -225,6 +225,16 @@ def _install_compile_time_sync_suppressors() -> None: _torch_cuda = torch.Tensor.cuda _torch_copy_ = torch.Tensor.copy_ + +# Once ``Tensor.to`` is replaced below, Dynamo sees the saved C++ descriptor +# through this module global instead of as a recognized tensor method. Put the +# descriptor behind a weak-referenceable Python callable so ``allow_in_graph`` +# can admit it and the downstream compiler can lower the original operation. +@torch.compiler.allow_in_graph +def _torch_to_in_graph(self, *args, **kwargs): + return _torch_to(self, *args, **kwargs) + + # torch's own parser for `Tensor.to`'s overloaded signature, returning # (device, dtype, non_blocking, memory_format) (also used by torch._dynamo). _parse_to = getattr(torch._C._nn, "_parse_to", None) @@ -232,7 +242,7 @@ def _install_compile_time_sync_suppressors() -> None: def _active_check_mode() -> str | None: """The mode the calling thread is being checked in, if any.""" - if not _sync_check_enabled or _allow_depth.get() or torch.compiler.is_compiling(): + if not _sync_check_enabled or torch.compiler.is_compiling() or _allow_depth.get(): return None return _checking.get() @@ -296,7 +306,7 @@ def _checked_to(self, *args, **kwargs): else: if non_blocking and (reason := _cpu_copy_stall_reason(self, device)): _report_implicit_sync("to", "H2D", reason, mode) - return _torch_to(self, *args, **kwargs) + return _torch_to_in_graph(self, *args, **kwargs) def _checked_cuda(self, *args, **kwargs): diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index b453487a01a8..0b330edf0038 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -705,17 +705,30 @@ def create_kv_caches_with_random( def async_tensor_h2d( data: list | np.ndarray | torch.Tensor, - device: str | torch.device, + device: str | torch.device | None = None, dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Copy list/numpy array/tensor async from host to device.""" + if dtype is None and out is not None: + dtype = out.dtype if isinstance(data, np.ndarray): data = torch.from_numpy(data) if isinstance(data, torch.Tensor): - t = data.pin_memory() if PIN_MEMORY else data + t = data + if PIN_MEMORY and not t.is_pinned(): + # Stage in pinned, contiguous buffer to ensure fully async copy. + t = torch.empty( + t.shape, dtype=dtype or t.dtype, device="cpu", pin_memory=True + ).copy_(t) else: t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu") assert t.is_cpu + + if out is not None: + assert out.dtype == dtype + return out.copy_(t, non_blocking=True) + assert device is not None, "must provide destination tensor or device" return t.to(device=device, dtype=dtype, non_blocking=True) @@ -816,8 +829,11 @@ def current_stream() -> torch.cuda.Stream: # https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77 # for more details. Therefore, we create a dedicated stream per process. if current_platform.is_rocm() or current_platform.is_cuda(): + # Ensure new stream is ordered w.r.t. replaced stream's work. + new_stream = torch.cuda.Stream() + new_stream.wait_stream(torch.cuda.current_stream()) # torch.cuda.set_stream here is the alias of _pathed_set_stream - torch.cuda.set_stream(torch.cuda.Stream()) + torch.cuda.set_stream(new_stream) elif current_platform.is_cpu(): _current_stream_tls.value = _StreamPlaceholder() else: diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 0da68e3e3d93..7ab2b2018e09 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1264,15 +1264,9 @@ def _compute_flashinfer_kv_metadata( """ # write self.paged_kv_indptr_cpu inplace (0-index is always 0) np.cumsum( - num_blocks_np, - dtype=np.int32, - out=self.paged_kv_indptr.np[1 : num_reqs + 1], + num_blocks_np, dtype=np.int32, out=self.paged_kv_indptr.np[1 : num_reqs + 1] ) - paged_kv_indptr = self.paged_kv_indptr.gpu[: num_reqs + 1] - paged_kv_indptr_cpu = self.paged_kv_indptr.cpu[: num_reqs + 1] - if PIN_MEMORY: - paged_kv_indptr_cpu = paged_kv_indptr_cpu.pin_memory() - paged_kv_indptr.copy_(paged_kv_indptr_cpu, non_blocking=True) + paged_kv_indptr = self.paged_kv_indptr.copy_to_gpu(num_reqs + 1) # write self.paged_kv_indices inplace num_actual_pages = self.paged_kv_indptr.np[num_reqs] @@ -1292,12 +1286,7 @@ def _compute_flashinfer_kv_metadata( page_size, paged_kv_last_page_len_np, ) - paged_kv_last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_reqs] - if PIN_MEMORY: - paged_kv_last_page_len_cpu = paged_kv_last_page_len_cpu.pin_memory() - self.paged_kv_last_page_len.gpu[:num_reqs].copy_( - paged_kv_last_page_len_cpu, non_blocking=True - ) + self.paged_kv_last_page_len.copy_to_gpu(num_reqs) return paged_kv_indices def build( diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index ae2e51284015..ba10a2e91cea 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -113,6 +113,16 @@ def supports_combination( from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 + and parallel_config.decode_context_parallel_size > 1 + ): + return ( + "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; " + "use FLASHMLA_SPARSE, which gathers each DCP KV shard before " + "running the rank-local PCP prefill queries" + ) if kv_cache_dtype == "fp8_ds_mla": return ( "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" @@ -265,6 +275,31 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS hisparse_supports_multi_token_decode: ClassVar[bool] = True + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + + # Under DCP the workspace must hold the trtllm-gen softmax-stats slab. + # The buffer address is baked into CUDA graphs, so it has to reach its + # final size here, before the first forward/capture (no lazy regrow). + if self.dcp_world_size > 1: + num_q_heads = vllm_config.model_config.get_num_attention_heads( + vllm_config.parallel_config + ) + _get_workspace_buffer( + device, + _required_workspace_bytes( + self.dcp_world_size, + num_q_heads, + vllm_config.scheduler_config.max_num_batched_tokens, + ), + ) + def _build_req_id_per_token( self, common_attn_metadata: "CommonAttentionMetadata", @@ -275,14 +310,98 @@ def _build_req_id_per_token( # Global workspace buffer (lazily initialized) _fi_sparse_workspace: torch.Tensor | None = None +# trtllm-gen carves a softmax-stats slab from the workspace whenever LSE is +# requested (FlashInfer csrc/trtllm_fmha_kernel_launcher.cu, unchanged from +# v0.6.14 through current main): +# sizeof(float2) * num_qo_heads * batch_size * round_up(max_q_len, 256) +# + 1 MiB guard +# forward_mqa always passes q_len == 1, so each (head, token) pair costs +# round_up(1, 256) == 256 slots. +_TRTLLM_GEN_SOFTMAX_STAT_BYTES = 8 # sizeof(float2) +_TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN = 256 # round_up(q_len=1, 256) +_TRTLLM_GEN_SOFTMAX_GUARD_BYTES = 1024 * 1024 + +# Keep in sync with the VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE default in +# vllm/envs.py. +_DEFAULT_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024 + + +def compute_trtllm_sparse_mla_workspace_bytes( + base_workspace_bytes: int, + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Workspace bytes needed by the trtllm-gen sparse MLA decode kernel. + + Under DCP the query is all-gathered across the DCP group in the head dim + and the kernel is asked for LSE, so trtllm-gen carves the softmax-stats + slab described above from the workspace before its (batch-independent) + counter and scratch regions. ``base_workspace_bytes`` must stay available + for those regions, so the slab is added on top of it (see #50781). + + Without DCP no LSE is requested, no slab is carved, and the base size is + returned unchanged. + """ + if dcp_world_size <= 1: + return base_workspace_bytes + softmax_bytes = ( + _TRTLLM_GEN_SOFTMAX_STAT_BYTES + * (num_heads_per_rank * dcp_world_size) + * max_num_batched_tokens + * _TRTLLM_GEN_SOFTMAX_SLOTS_PER_TOKEN + + _TRTLLM_GEN_SOFTMAX_GUARD_BYTES + ) + return base_workspace_bytes + softmax_bytes + + +def _required_workspace_bytes( + dcp_world_size: int, + num_heads_per_rank: int, + max_num_batched_tokens: int, +) -> int: + """Resolve the workspace size, honoring an explicit env override.""" + computed = compute_trtllm_sparse_mla_workspace_bytes( + _DEFAULT_WORKSPACE_BUFFER_SIZE, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + ) + if not envs.is_set("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE"): + return computed + env_bytes = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + if env_bytes < computed: + logger.warning_once( + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d is below the %d bytes " + "computed for the sparse MLA workspace with dcp_world_size=%d, " + "%d heads/rank and max_num_batched_tokens=%d. Respecting the " + "override, but the trtllm-gen kernel may crash with a workspace " + "overflow; set VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE=%d or unset " + "it to use the computed size.", + env_bytes, + computed, + dcp_world_size, + num_heads_per_rank, + max_num_batched_tokens, + computed, + ) + return env_bytes + -def _get_workspace_buffer(device: torch.device) -> torch.Tensor: +def _get_workspace_buffer( + device: torch.device, min_bytes: int | None = None +) -> torch.Tensor: global _fi_sparse_workspace - if _fi_sparse_workspace is None: + required = ( + min_bytes + if min_bytes is not None + else envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE + ) + if _fi_sparse_workspace is None or _fi_sparse_workspace.numel() < required: # FlashInfer's CuteDSL MLA-decode tactic requires an int8 workspace; # the trtllm-gen path views it as uint8, so int8 is safe for all backends. _fi_sparse_workspace = torch.zeros( - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, + required, dtype=torch.int8, device=device, ) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 90dc4cd94738..830bc52ab939 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py @@ -277,6 +277,7 @@ def __init__( attention_layer = vllm_config.compilation_config.static_forward_context[ layer_names[0] ] + self._attention_layer = attention_layer impl = attention_layer.impl if not isinstance(impl, FlashInferMLASparseSM90Impl): raise TypeError( @@ -376,6 +377,13 @@ def build( # Replan every step outside any CUDA graph capture with this step's # exact per-row lengths; captured runs read the refreshed buffers. num_rows, kv_lens = self._kv_lens_host(common_attn_metadata) + # MHA prefills are removed from q before forward_mqa. Plan only the + # decode prefix in that case, using the same routing decision as MLA. + if metadata.num_prefills > 0 and self._attention_layer._use_sparse_mha( + metadata + ): + num_rows = metadata.num_decode_tokens + kv_lens = kv_lens[:num_rows] self.state.plan(num_rows, kv_lens) metadata.state = self.state return metadata diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index d8c4368401c6..3a9ca2ffdb06 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -20,7 +20,11 @@ from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import current_stream, is_quantized_kv_cache +from vllm.utils.torch_utils import ( + async_tensor_h2d, + current_stream, + is_quantized_kv_cache, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -48,7 +52,6 @@ get_mla_metadata, ) from vllm.v1.kv_cache_interface import AttentionSpec -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -91,6 +94,17 @@ The first `ue8m0` is the scale for the first 64 `float8_e4m3` values, the second for the next 64, and so on. +For DeepSeek V4.1, each token's KV cache is 528 Bytes, structured as: +- **First 512 bytes:** all 512 dims as `float8_e4m3`. Unlike V4 the 64 RoPE + dims are quantized too, so there is no `bfloat16` part. +- **Last 16 bytes:** Scale factors, containing 16 `ue8m0` values, one per 32 + consecutive `float8_e4m3` values (i.e. MXFP8). + +The V4 and V4.1 records are not laid out token-by-token within a page: a page +holds all its data rows first and all its scale rows after, so a page is +`block_size * bytes_per_token` bytes rounded up to the decode kernel's TMA +stride (576 B for V4, 512 B for V4.1). + In the "nvfp4_ds_mla" format (SM100 only, DeepSeek V3.2 geometry), each token's KV cache is 352 Bytes, structured as: - **First 256 bytes:** 512 `e2m1` NoPE values packed 2/byte (low nibble = @@ -508,7 +522,7 @@ def _build_fp8_separate_prefill_decode( ) workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32)) max_prefill_buffer_size //= self.dcp_world_size - entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device) + entry_rows = async_tensor_h2d(row_bounds[:-1], device=self.device) prefill_block_table = prefill_block_table.index_select(0, entry_rows) prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows) num_entries = len(workspace_rows) diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 8ba0171835fd..1b8455a2328a 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -10,6 +10,7 @@ from vllm import _custom_ops as ops from vllm.config import VllmConfig +from vllm.forward_context import in_piecewise_cudagraph from vllm.utils.torch_utils import current_stream from vllm.v1.attention.backends.mla.sparse_utils import ( triton_convert_req_index_to_global_index, @@ -56,8 +57,16 @@ def register_layer( return layer_index def set_logical_topk_ready(self, layer_index: int) -> None: - if layer_index == 0 and self.has_indexer: - self.logical_topk_ready.record(current_stream()) + if layer_index != 0 or not self.has_indexer: + return + # In piecewise cudagraph mode the consumer converts on the side stream + # from an eager break; an event recorded inside a captured segment is + # graph-local and any eager wait on it raises cudaErrorInvalidValue. + # The convert side falls back to wait_stream there instead (and never + # touches this event), so skip the record as well. + if in_piecewise_cudagraph(): + return + self.logical_topk_ready.record(current_stream()) def prepare_for_batch(self, layer_index: int, attn_metadata: Any | None) -> None: pass @@ -129,7 +138,7 @@ def _convert_once( valid_topk_counts = self.valid_topk_counts[:num_tokens] compute_stream = current_stream() if layer_index == 0: - if self.has_indexer: + if self.has_indexer and not in_piecewise_cudagraph(): self.side_stream.wait_event(self.logical_topk_ready) else: self.side_stream.wait_stream(compute_stream) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2321cf8d13d4..53206afd9337 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -10,7 +10,7 @@ from vllm.config import VllmConfig from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup import kernel_launcher +from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonPointerInputVariant, @@ -27,7 +27,7 @@ ) from vllm.utils.math_utils import round_down from vllm.utils.platform_utils import num_compute_units -from vllm.utils.torch_utils import PIN_MEMORY +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -47,7 +47,6 @@ KVCacheSpec, MLAAttentionSpec, ) -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) @@ -338,14 +337,14 @@ def build_pcp_global_chunk_plan( + local ) - cu = async_copy_to_gpu(cu, device=device) + cu = async_tensor_h2d(cu, device=device) return PCPGlobalChunkPlan( row_start_cu=cu[0], global_cu=cu[1], padded_local_cu=cu[2], padded_local_total=padded_total, total=total, - deinterleave_idx=async_copy_to_gpu(idx, device=device), + deinterleave_idx=async_tensor_h2d(idx, device=device), ) @@ -493,6 +492,14 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: dcp_world = parallel_config.decode_context_parallel_size dcp_interleave = parallel_config.cp_kv_cache_interleave_size dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0 + dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)] + if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1: + # PCP+DCP prefill constructs a global PCP chunk plan and therefore + # dispatches this kernel with already-localized row starts. The + # runtime key is normalized to rank 0/world 1 in that path; warm it + # explicitly so one PCP rank cannot enter a collective while + # another rank is still compiling the first real inference batch. + dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1)) compress_ratios = tuple( dict.fromkeys( max(1, int(ratio)) @@ -506,17 +513,33 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios: compress_ratios = compress_ratios + (index_kpool,) return self._trace_dispatch(self.dispatch)( + zip_inputs(*dcp_cases), # Cover Triton's divisible, exact-one, and generic i32 classes. query_slice_start=(0, 1, 2), query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), - DCP_RANK=dcp_rank, - DCP_WORLD=dcp_world, DCP_INTERLEAVE=dcp_interleave, BLOCK_SIZE=self.BLOCK_SIZE, COMPRESS_RATIO=list(compress_ratios), + # PCP's global cumulative lengths are the second row of one packed + # allocation, so their pointer is not always 16-byte aligned. + # Prefill chunking can independently unalign uncompressed lengths. input_variant=( - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True), - TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=True, + cu_compressed_seq_lens=True, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=True, + cu_compressed_seq_lens=False, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=False, + cu_compressed_seq_lens=True, + ), + TritonPointerInputVariant.from_alignment( + uncompressed_seq_lens=False, + cu_compressed_seq_lens=False, + ), ), ) @@ -527,7 +550,9 @@ def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: uncompressed_seq_lens=compile_key.input_variant.pointer( "uncompressed_seq_lens", torch.int32 ), - cu_compressed_seq_lens=int32_ptr, + cu_compressed_seq_lens=compile_key.input_variant.pointer( + "cu_compressed_seq_lens", torch.int32 + ), row_start_cu_compressed_seq_lens=int32_ptr, token_to_seq=int32_ptr, cu_compressed_seq_len_ks=int32_ptr, @@ -836,6 +861,12 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: dtype=torch.int32, device=self.device, ) + # Materialize the rank on device during builder initialization. Creating + # this scalar in build() would introduce a GPU<->CPU sync in the decode + # hot path. + self.dcp_rank_tensor = torch.tensor( + self.dcp_rank, dtype=torch.int32, device=self.device + ) self.expanded_block_table_buffer = torch.zeros( (scheduler_config.max_num_batched_tokens, block_table_width), dtype=torch.int32, @@ -889,7 +920,7 @@ def _dcp_localize_decode_seq_lens( local_seq_lens = get_dcp_local_seq_lens( seq_lens, self.dcp_world_size, - self.dcp_rank, + self.dcp_rank_tensor, self.cp_kv_cache_interleave_size, ) if seq_lens_is_buffer_view: @@ -1085,6 +1116,11 @@ def _split_pcp_dcp_prefill_chunks( for request_slice, query_slice in chunk_specs ] + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + """Per-request KV lengths the prefill chunker budgets logits with; + subclasses whose logits rows are wider than the context override.""" + return seq_lens_cpu + @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1247,7 +1283,7 @@ def build( ) else: chunk_specs = self._split_indexer_prefill_chunks( - compressed_seq_lens_cpu[num_decodes:], + self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), prefill_query_lens_cpu, self.max_prefill_buffer_size, max_logits_bytes, diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index b850deb07091..c9ba2b0a566e 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -4,7 +4,7 @@ import functools from dataclasses import dataclass from pathlib import Path -from typing import ClassVar, Final +from typing import TYPE_CHECKING, ClassVar, Final import torch @@ -38,6 +38,9 @@ ) from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache +if TYPE_CHECKING: + from vllm.platforms.interface import DeviceCapability + logger = init_logger(__name__) @@ -144,6 +147,20 @@ def _aiter_mla_native_h24_supported() -> bool: ) +def _aiter_mla_non_causal_asm_kernels() -> bool: + """Whether this arch ships non-causal MLA decode ASM kernels. + + The Python `causal=` probe cannot see the per-arch manifest. gfx950 has + the kernels; gfx942 and everything else must fall through to another + backend. + """ + try: + from vllm.platforms.rocm import on_gfx950 + except Exception: # noqa: BLE001 + return False + return bool(on_gfx950()) + + @functools.lru_cache(maxsize=1) def _gluon_mla_decode_supported() -> bool: """The small-head Gluon MLA decode kernel only has a gfx950 (CDNA4) build. @@ -252,6 +269,72 @@ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: def get_name() -> str: return "ROCM_AITER_MLA" + @classmethod + def supports_non_causal(cls) -> bool: + # causal= on mla_decode_fwd is necessary but not sufficient: pinned + # AITER v0.1.21.post2 still exports that argument on gfx942, whose ASM + # manifest has no non-causal decode kernels. Selecting this backend + # there reaches `cannot get heuristic kernel!` instead of TRITON_MLA. + return bool( + rocm_aiter_ops.mla_decode_supports_non_causal() + and _aiter_mla_non_causal_asm_kernels() + ) + + @classmethod + def validate_configuration( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + use_per_head_quant_scales: bool, + device_capability: "DeviceCapability", + attn_type: str, + has_sliding_window: bool = False, + use_non_causal: bool = False, + use_batch_invariant: bool = False, + use_kv_connector: bool = False, + use_pcp: bool = False, + use_adaptive_verification: bool = False, + use_dcp: bool = False, + use_rswa: bool = False, + ) -> list[str]: + invalid_reasons = super().validate_configuration( + head_size=head_size, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + block_size=block_size, + use_mla=use_mla, + has_sink=has_sink, + use_sparse=use_sparse, + use_mm_prefix=use_mm_prefix, + use_per_head_quant_scales=use_per_head_quant_scales, + device_capability=device_capability, + attn_type=attn_type, + has_sliding_window=has_sliding_window, + use_non_causal=use_non_causal, + use_batch_invariant=use_batch_invariant, + use_kv_connector=use_kv_connector, + use_pcp=use_pcp, + use_adaptive_verification=use_adaptive_verification, + use_dcp=use_dcp, + use_rswa=use_rswa, + ) + # Arch/signature cannot express this: the backend still advertises + # fp16, so a DSpark draft with kv_cache_dtype="auto" would select + # AITER and abort on fp16 query. Causal fp16 is fine. + if ( + use_non_causal + and dtype == torch.float16 + and kv_cache_dtype in (None, "auto", "float16") + ): + invalid_reasons.append("non-causal fp16 MLA decode not supported") + return invalid_reasons + @staticmethod def get_impl_cls() -> type["AiterMLAImpl"]: return AiterMLAImpl @@ -347,6 +430,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # https://github.com/vllm-project/vllm/issues/22945 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM + # Served by passing the mask to the kernel; _build_decode turns away the + # shapes AITER has no non-causal kernel for. + supports_non_causal_multi_token_decode: ClassVar[bool] = True + # Set from the common metadata every build; a batch is causal unless the + # drafter says otherwise. + _decode_causal: bool = True @staticmethod def _uniform_padded_mtp_qo_len( @@ -919,6 +1008,7 @@ def _build_decode( num_decode_tokens: int, dcp_tot_seq_lens_device: torch.Tensor | None, ) -> AiterMLADecodeMetadata: + causal = self._decode_causal device = self.device num_reqs = seq_lens_device.size(0) qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] @@ -970,6 +1060,7 @@ def _build_decode( int(max_qo_len), self._kv_cache_dtype_str, self.dcp_world_size, + causal, ) use_segmented_dcp_verify = ( self._supports_segmented_dcp_verify and max_qo_len > 1 @@ -1047,14 +1138,36 @@ def _build_decode( # A padded rank has no bf16 persistent kernel past qlen 4 where the # gfx950 fold is absent; the non-persistent entry covers it. fp8 # keeps the schedule -- its fold rejects non-persistent outright. + # A non-causal block keeps it too: what the fold drops past qlen 4 + # is the block's causal staircase, which a non-causal block does + # not have, and the schedule is the only thing carrying its mask. and ( - self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS + not causal + or self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN or is_quantized_kv_cache(self._kv_cache_dtype_str) ) and max_qo_len >= 1 and max_qo_len <= self._mtp_decode_qlen ) + if ( + not causal + and max_qo_len == 2 + and is_quantized_kv_cache(self._kv_cache_dtype_str) + and not AiterMLAHelper.has_fp8_non_causal_qlen2_kernel( + self._decode_num_heads + ) + ): + # AITER's fp8 dispatch folds (gqa 16, qlen 3 or 4) onto the + # qseqlen-4 kernel but never lists 2. Only 32/64/96/128 heads + # at qlen 2 fold onto that 16-head / 4-token non-causal kernel; + # 16-head qlen 2, and padded 48/80/112, keep Q2 and abort the + # process rather than raising. The bf16 fold has no such hole. + raise ValueError( + "AITER has no non-causal fp8 MLA kernel for this 2-token " + "query block. Pin the draft to TRITON_MLA for this " + "speculative config." + ) if use_persistent_metadata: from aiter import get_mla_metadata_v1 @@ -1067,7 +1180,7 @@ def _build_decode( paged_kv_last_page_len, self._num_attention_heads, 1, - True, + causal, self._mla_work_meta_data, self._mla_work_info_set, self._mla_work_indptr, @@ -1137,6 +1250,11 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> AiterMLAMetadata: + # The common MLA builder owns the causal flag but its _build_decode + # contract is shared with CUDA backends. Keep this ROCm-specific state + # on the per-ubatch Aiter builder instead of changing that shared API; + # super().build invokes _build_decode synchronously. + self._decode_causal = common_attn_metadata.causal attn_metadata = super().build( common_prefix_len, common_attn_metadata, fast_build ) @@ -1226,6 +1344,10 @@ class AiterMLAHelper: # fold that reaches a persistent one is gfx950-only. _ASM_PADDED_MAX_PS_QLEN: Final = 4 _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = () + # Pinned AITER v0.1.21.post2 folds these fp8 qlen-2 counts onto the + # non-causal 16-head / 4-token kernel. Other multiples of 16 (48/80/112) + # fold to H16 while keeping Q2, which has no non-causal fp8 entry. + _AITER_FP8_NON_CAUSAL_QLEN2_HEADS: ClassVar[tuple[int, ...]] = (32, 64, 96, 128) @staticmethod def qo_indptr_for_uniform_qlen( @@ -1263,6 +1385,18 @@ def is_valid_num_heads(num_heads: int) -> bool: ) ) + @staticmethod + def has_fp8_non_causal_qlen2_kernel(num_heads: int) -> bool: + """Whether fp8 (num_heads, qlen=2) folds onto a non-causal ASM kernel. + + Pinned AITER v0.1.21.post2 folds 32/64/96/128 heads at qlen 2 onto the + non-causal 16-head / 4-token kernel. A 16-head (or padded-to-16) qlen-2 + block has no matching entry, and padded 48/80/112 keep Q2 after the + H16 fold. + """ + kernel_heads = AiterMLAHelper.get_actual_mla_num_heads(num_heads) + return kernel_heads in AiterMLAHelper._AITER_FP8_NON_CAUSAL_QLEN2_HEADS + @staticmethod def get_actual_mla_num_heads(num_heads: int) -> int: if num_heads == 24 and _aiter_mla_native_h24_supported(): @@ -1364,6 +1498,7 @@ def use_gluon_verify( max_qo_len: int, kv_cache_dtype: str, dcp_world_size: int = 1, + causal: bool = True, ) -> bool: """Whether a small-head multi-token verify uses native Gluon MTP. @@ -1376,7 +1511,13 @@ def use_gluon_verify( DCP verify is excluded: its per-row causal windows are served by the segmented path, which Gluon's MTP entry cannot express. + + Gluon masks the block causally with no way to turn it off, so a + non-causal block takes the padded asm decode whatever the head count + -- padding is what gives it a kernel there. """ + if not causal: + return False if max_qo_len <= 1 or dcp_world_size > 1: return False if is_quantized_kv_cache(kv_cache_dtype): @@ -1974,6 +2115,7 @@ def forward_mqa( decode.paged_kv_indptr, decode.paged_kv_indices, decode.paged_kv_last_page_len, + causal=attn_metadata.causal, **mla_kwargs, ) diff --git a/vllm/v1/attention/backends/mla/sparse_indexer.py b/vllm/v1/attention/backends/mla/sparse_indexer.py new file mode 100644 index 000000000000..c0b445534d22 --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_indexer.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1 indexer backend for the DeepGEMM sparse MQA-logits path. + +Candidate-consuming indexer layers score only the candidate blocks published +by the candidate-source layer (``fp8_fp4_(paged_)sparse_mqa_logits``) instead +of dense logits over the whole context. This backend extends the dense +indexer metadata with the per-row scratch that path needs, so the layer op +allocates nothing per step, and lets all consumer layers of a step share one +candidate expansion and DeepGEMM schedule (they all read the same candidate +blocks). Opt in with ``AttentionConfig.indexer_sparse_logits``. +""" + +from dataclasses import dataclass, fields + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + check_deep_select_layout, + has_deep_select, + pick_sparse_block_kv, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, + dsa_indexer_uses_fp4, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +@dataclass +class SparseMQARowsMetadata: + """Per-row state for one sparse-logits call: the decode rows, or one + prefill chunk. Buffers are builder-owned views sized to the rows. + + ``kernel_metadata`` (the DeepGEMM schedule) is filled by the first + consumer indexer layer of the step together with ``sparse_indices`` / + ``end``; later consumer layers reuse all three. + """ + + row_ks: torch.Tensor + """[rows] int32 K-range start per row; zeros for paged decode.""" + row_ke: torch.Tensor + """[rows] int32 K-range end (prefill) or compressed context length.""" + sparse_indices: torch.Tensor + """[rows, num_sparse_blocks] int32 DeepGEMM sparse block ids.""" + end: torch.Tensor + """[rows] int32 valid sparse-column count per row.""" + col_indices: torch.Tensor + """[rows, topk] int32 scratch for the sparse-column top-k.""" + block_table: torch.Tensor | None = None + """Decode only: [rows, pages] per-row page table.""" + row_indices: torch.Tensor | None = None + """Decode only: [rows] row -> request index.""" + kernel_metadata: torch.Tensor | None = None + + +@dataclass +class DeepseekV41SparseIndexerMetadata(DeepseekV32IndexerMetadata): + sparse_block_kv: int = 0 + sparse_decode: SparseMQARowsMetadata | None = None + sparse_prefill: list[SparseMQARowsMetadata] | None = None + """Parallel to ``prefill.chunks``.""" + + +class DeepseekV41SparseIndexerBackend(DeepseekV41IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_SPARSE_INDEXER" + + @staticmethod + def get_builder_cls() -> type["DeepseekV41SparseIndexerMetadataBuilder"]: + return DeepseekV41SparseIndexerMetadataBuilder + + +class DeepseekV41SparseIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): + """Dense indexer metadata plus the sparse-logits row scratch. + + Every requirement of the sparse kernels is checked here, at engine start, + rather than falling back to the dense path per step. + """ + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + **kwargs, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) + hf_config = vllm_config.model_config.hf_text_config + candidate_block_size = getattr(hf_config, "candidate_block_size", 0) + candidate_topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0) + topk_tokens = getattr(hf_config, "index_topk", 0) + prefix = "attention_config.indexer_sparse_logits" + if candidate_block_size <= 0 or candidate_topk_blocks <= 0: + raise ValueError( + f"{prefix} requires a two-level (candidate block) indexer model " + "such as DeepSeek-V4.1." + ) + if not dsa_indexer_uses_fp4(vllm_config): + raise ValueError( + f"{prefix} requires the MXFP4 indexer cache " + "(attention_config.indexer_kv_dtype='mxfp4'): the sparse kernels " + "take UE8M0-packed granularity-32 scales." + ) + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + raise ValueError(f"{prefix} requires an SM100-class GPU.") + if not has_deep_gemm_sparse_mqa(): + raise ValueError( + f"{prefix} requires DeepGEMM >= 2.8 " + "(fp8_fp4_sparse_mqa_logits not found)." + ) + if self.dcp_world_size > 1 or self.use_pcp: + raise NotImplementedError( + f"{prefix} is not supported with decode/prefill context parallel." + ) + if not self.supports_varlen: + raise ValueError( + f"{prefix} requires the varlen paged MQA logits decode layout." + ) + # The top-k runs on the kernels' bf16 logits with DeepSelect; there is + # no fp32 fallback, so the other top-k backends cannot be honored. + topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend + if not has_deep_select(): + raise ValueError( + f"{prefix} requires the DeepSelect top-k extension " + "(vllm._deepselect_C) for this GPU." + ) + if topk_backend not in ("auto", "deep_select"): + raise ValueError( + f"{prefix} uses DeepSelect for the top-k; " + f"kernel_config.sparse_indexer_topk_backend='{topk_backend}' " + "is not supported with it." + ) + self.sparse_block_kv = pick_sparse_block_kv(candidate_block_size) + if kv_cache_spec.block_size % self.sparse_block_kv != 0: + raise ValueError( + f"{prefix}: indexer cache pages ({kv_cache_spec.block_size} " + f"tokens) must be a multiple of {self.sparse_block_kv} tokens." + ) + self.num_sparse_blocks = candidate_topk_blocks * ( + candidate_block_size // self.sparse_block_kv + ) + num_sparse_cols = self.num_sparse_blocks * self.sparse_block_kv + check_deep_select_layout(num_sparse_cols, topk_tokens) + # Sparse logits are bf16 [rows, num_sparse_cols]; make the prefill + # chunker (which budgets fp32 [rows, seq_len]) size chunks as if every + # row were this wide. + self.min_split_seq_len = (num_sparse_cols * 2 + 3) // 4 + + # Rows are indexed by batch token (decode rows first, then prefill + # chunks by token range), like ``arange_buffer``. + max_rows = self.arange_buffer.shape[0] + int32 = dict(dtype=torch.int32, device=device) + self.sparse_indices_buffer = torch.zeros( + (max_rows, self.num_sparse_blocks), **int32 + ) + self.end_buffer = torch.zeros(max_rows, **int32) + self.col_indices_buffer = torch.zeros((max_rows, topk_tokens), **int32) + self.zero_row_ks = torch.zeros(max_rows, **int32) + logger.info_once( + "DeepSeek V4.1 indexer: candidate consumers score only the candidate " + "blocks with DeepGEMM sparse MQA logits (%d sparse blocks of %d " + "tokens per row).", + self.num_sparse_blocks, + self.sparse_block_kv, + ) + + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + return seq_lens_cpu.clamp(min=self.min_split_seq_len) + + def _rows(self, start: int, end: int, row_ks, row_ke, **kwargs): + return SparseMQARowsMetadata( + row_ks=row_ks, + row_ke=row_ke, + sparse_indices=self.sparse_indices_buffer[start:end], + end=self.end_buffer[start:end], + col_indices=self.col_indices_buffer[start:end], + **kwargs, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV41SparseIndexerMetadata: + base = super().build(common_prefix_len, common_attn_metadata, fast_build) + metadata = DeepseekV41SparseIndexerMetadata( + **{f.name: getattr(base, f.name) for f in fields(base)}, + sparse_block_kv=self.sparse_block_kv, + ) + if base.decode is not None: + decode = base.decode + rows = decode.seq_lens.shape[0] + # The varlen layout flattens every query to one row, so the base + # decode tensors are already per row. + assert decode.seq_lens.shape[1] == 1 and not decode.requires_padding + assert decode.block_table.shape[0] == rows + metadata.sparse_decode = self._rows( + 0, + rows, + row_ks=self.zero_row_ks[:rows], + row_ke=decode.seq_lens.view(-1), + block_table=decode.block_table, + row_indices=( + decode.indices + if decode.indices is not None + else self.arange_buffer[:rows] + ), + ) + if base.prefill is not None: + metadata.sparse_prefill = [ + self._rows( + chunk.token_start, + chunk.token_end, + row_ks=chunk.cu_seqlen_ks, + row_ke=chunk.cu_seqlen_ke, + ) + for chunk in base.prefill.chunks + ] + return metadata diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 62c3921b5e12..c4f3cd8b4bdd 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -79,6 +79,8 @@ def __init__( cache_config: CacheConfig, backend_cls: "type[AttentionBackend] | None" = None, block_size: int = 64, + packed_bytes_per_token: int = 584, + packed_page_alignment: int = 576, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend @@ -88,6 +90,8 @@ def __init__( self.prefix = prefix self.cache_config = cache_config self.dtype = dtype + self.packed_bytes_per_token = packed_bytes_per_token + self.packed_page_alignment = packed_page_alignment compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -105,8 +109,9 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: self.kv_cache = kv_cache.squeeze(1) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous - # bf16/fp8 cache uses the natural element-size page. + # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode + # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural + # element-size page. uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( block_size=self.block_size, @@ -115,10 +120,11 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales) - state_content_bytes=584 if uses_fp8_ds_mla_layout else None, - # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). - alignment=576 if uses_fp8_ds_mla_layout else 512, + state_content_bytes=( + self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None + ), + # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). + alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index a3adf70d2588..e656343a9189 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -1094,7 +1094,7 @@ def compute_causal_conv1d_metadata( def get_dcp_local_seq_lens( seq_lens: torch.Tensor, dcp_size: int = 1, - dcp_rank: int | None = None, + dcp_rank: int | torch.Tensor | None = None, cp_kv_cache_interleave_size: int = 1, ) -> torch.Tensor: """While using dcp, kv_cache size stored on each rank may be different, @@ -1112,6 +1112,12 @@ def get_dcp_local_seq_lens( dcp_size, ) seq_lens_tiled = seq_lens_i32.unsqueeze(-1) + elif isinstance(dcp_rank, torch.Tensor): + assert dcp_rank.dtype == torch.int32 + assert dcp_rank.device == seq_lens.device + assert dcp_rank.numel() == 1 + rank_offsets = dcp_rank + seq_lens_tiled = seq_lens_i32 else: rank_offsets = torch.tensor(dcp_rank, dtype=torch.int32, device=seq_lens.device) seq_lens_tiled = seq_lens_i32 diff --git a/vllm/v1/attention/ops/dcp.py b/vllm/v1/attention/ops/dcp.py index bd20c9020d3e..58ac1e14a09d 100644 --- a/vllm/v1/attention/ops/dcp.py +++ b/vllm/v1/attention/ops/dcp.py @@ -77,6 +77,13 @@ def mask_dcp_empty_shards_( return assert seq_lens is not None and query_start_loc is not None + # A DCP rank can receive no local sequences during CUDA graph warmup even + # though the padded LSE buffer still has rows. In that case every row is an + # empty shard; avoid indexing the empty seq_lens tensor below. + if seq_lens.shape[0] == 0: + lse.fill_(float("-inf")) + return + row_indices = torch.arange( lse.shape[0], device=lse.device, dtype=query_start_loc.dtype ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 44b6017a057b..0e342cb0e064 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -58,6 +58,7 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 _GFX950_C4A_NATIVE_MAX_ROWS = 256 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 @@ -69,6 +70,8 @@ def _get_aiter_top_k_kernel( compress_ratio: int, num_rows: int, max_valid_seq_len: int | None = None, + num_columns: int | None = None, + topk_tokens: int = 1024, on_gfx950: bool = _ON_GFX950, ) -> Callable[..., None] | None: if compress_ratio <= 1 or not on_gfx950: @@ -76,6 +79,13 @@ def _get_aiter_top_k_kernel( if not is_prefill: assert max_valid_seq_len is not None + if ( + topk_tokens == 512 + and 0 < num_rows <= 384 + and num_columns is not None + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS + ): + return None # AITER v0.1.19 decode is one-block only. This measured gfx950 # FP32/k=1024 compressed-row boundary is independent of the native # split-count boundary in sampler.cu. @@ -1261,6 +1271,8 @@ def rocm_aiter_sparse_attn_indexer( compress_ratio=compress_ratio, num_rows=num_rows, max_valid_seq_len=max_compressed_seq_len, + num_columns=logits.shape[1], + topk_tokens=topk_tokens, ) if aiter_topk_kernel is not None: _launch_aiter_top_k_per_row_decode( diff --git a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py b/vllm/v1/attention/ops/triton_fp8_mqa_logits.py deleted file mode 100644 index 619d0ec50a9d..000000000000 --- a/vllm/v1/attention/ops/triton_fp8_mqa_logits.py +++ /dev/null @@ -1,262 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Temporary gfx942 fallback for AITER's fp8_mqa_logits kernel. - -This module vendors AITER's Triton fp8_mqa_logits kernel with the gfx942 -tile-size workaround from ROCm/aiter#3257. It is used only while vLLM's -pinned AITER version lacks that fix. - -TODO: Remove this vendored copy once vLLM pins an AITER version that includes -ROCm/aiter#3257 bugfix for gfx942. -""" - -import torch - -from vllm.triton_utils import tl, triton - -# gfx942 (MI300X) has 64 KiB of LDS per CU. We accept the default -# (BLOCK_KV=128, num_stages=2) tile only when *both* of these hold: -# -# 1. Occupancy gate. With waves_per_eu=2 and num_warps=4 we target two -# workgroups co-resident on a CU -> per-WG LDS budget = 32 KiB. Triton -# keeps Q in registers (loop-invariant) and the fp32 scores accumulator -# in VGPRs (heavy VALU), so only the double-buffered KV tile is -# expected to live in LDS. A 0.9 safety factor leaves headroom for any -# LDS overhead the compiler may add. -# -# 2. Hardware ceiling. Defensive upper bound that also counts Q and -# scores against the 64 KiB CU limit, in case a Triton version (older -# or future) decides to spill them to LDS. False positives here only -# shrink the tile; false negatives are JIT-aborts, so we lean -# conservative. -_GFX942_CU_LDS_BYTES = 64 * 1024 -_GFX942_PER_WG_LDS_BUDGET_BYTES = _GFX942_CU_LDS_BYTES * 9 // 20 # ~28.8 KiB - - -def _gfx942_default_tile_fits_lds(num_heads: int, head_size: int) -> bool: - """Return True iff (BLOCK_KV=128, num_stages=2) fits in MI300X LDS.""" - BLOCK_KV = 128 - NUM_STAGES = 2 - kv_bytes = head_size * BLOCK_KV * NUM_STAGES - scores_bytes = num_heads * BLOCK_KV * 4 - q_bytes = num_heads * head_size - fits_occupancy = kv_bytes < _GFX942_PER_WG_LDS_BUDGET_BYTES - fits_hardware = q_bytes + kv_bytes + scores_bytes <= _GFX942_CU_LDS_BYTES - return fits_occupancy and fits_hardware - - -@triton.jit -def _fp8_mqa_logits_kernel( - Q_ptr, # fp8e4m3 [seq_len, H, D] - KV_ptr, # fp8e4m3 [seq_len_kv, D] - kv_scales_ptr, # fp32 [seq_len_kv] - weights_ptr, # fp32 [seq_len, H] - cu_start_ptr, # int32 [seq_len] - cu_end_ptr, # int32 [seq_len] - logits_ptr, # fp32 [seq_len, seq_len_kv] - seq_len, - seq_len_kv, - NUM_HEADS: tl.constexpr, - HEAD_SIZE: tl.constexpr, - # strides - stride_q_s: tl.int64, - stride_q_h: tl.constexpr, - stride_q_d: tl.constexpr, - stride_kv_s: tl.int64, - stride_kv_d: tl.constexpr, - stride_w_s: tl.int64, - stride_w_h: tl.constexpr, - stride_logits_s: tl.int64, - stride_logits_k: tl.int64, - # block sizes - BLOCK_KV: tl.constexpr, -): - row_id = tl.program_id(0) - # go from larger to smaller in terms of work - # to reduce the tail effect - row_id = tl.num_programs(0) - row_id - 1 - tl.assume(row_id >= 0) - tl.assume(stride_q_s > 0) - tl.assume(stride_q_h > 0) - tl.assume(stride_q_d > 0) - tl.assume(stride_kv_s > 0) - tl.assume(stride_kv_d > 0) - tl.assume(stride_w_s > 0) - tl.assume(stride_w_h > 0) - - logits_row_ptrs = logits_ptr + row_id * stride_logits_s - - h_inds = tl.arange(0, NUM_HEADS)[:, None] - d_inds = tl.arange(0, HEAD_SIZE) - - # load Q[BLOCK_Q, NUM_HEADS, HEAD_SIZE] - q_ptrs = ( - Q_ptr + row_id * stride_q_s + h_inds * stride_q_h + d_inds[None, :] * stride_q_d - ) - - q_block = tl.load(q_ptrs, cache_modifier=".cg") - w_ptrs = weights_ptr + row_id * stride_w_s + h_inds * stride_w_h - w_block = tl.load(w_ptrs, cache_modifier=".cg").to(tl.float32) - - # Load start/end for each row in this block - start_ind = tl.load(cu_start_ptr + row_id) - end_ind = tl.load(cu_end_ptr + row_id) - - start_ind = tl.maximum(start_ind, 0) - end_ind = tl.minimum(end_ind, seq_len_kv) - shifted_end = end_ind - start_ind - shifted_unmasked_end = shifted_end // BLOCK_KV * BLOCK_KV - - kv_col_offsets = tl.arange(0, BLOCK_KV) + start_ind - kv_ptrs = ( - KV_ptr + kv_col_offsets[None, :] * stride_kv_s + d_inds[:, None] * stride_kv_d - ) - - kv_scales_ptrs = kv_scales_ptr + kv_col_offsets - - logits_ptrs = logits_row_ptrs + kv_col_offsets * stride_logits_k - - # Loop over KV tiles - for _ in tl.range(0, shifted_unmasked_end, BLOCK_KV): - kv_block = tl.load(kv_ptrs) - kv_scales = tl.load(kv_scales_ptrs) - - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV] - scores = tl.dot(q_block, kv_block, input_precision="ieee") - # Multiply by kv_scales (broadcast along rows) - scores = scores * kv_scales[None, :] - # ReLU - scores = tl.maximum(scores, 0.0) - scores = scores * w_block - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ] - scores = tl.sum(scores, axis=0) - tl.store(logits_ptrs, scores) - - kv_ptrs += BLOCK_KV * stride_kv_s - kv_scales_ptrs += BLOCK_KV - logits_ptrs += BLOCK_KV * stride_logits_k - kv_col_offsets += BLOCK_KV - - # masked load - kv_col_mask = kv_col_offsets < end_ind - kv_block = tl.load(kv_ptrs, mask=kv_col_mask[None, :], other=0.0) - kv_scales = tl.load(kv_scales_ptrs, mask=kv_col_mask, other=0.0) - - # [NUM_HEADS, BLOCK_KV] = [NUM_HEADS, HEAD_SIZE] x [HEAD_SIZE, BLOCK_KV] - scores = tl.dot(q_block, kv_block, input_precision="ieee") - # Multiply by kv_scales (broadcast along rows) - scores = scores * kv_scales[None, :] - # ReLU - scores = tl.maximum(scores, 0.0) - scores = scores * w_block - # [NUM_HEADS, BLOCK_KV] -> [BLOCK_KV, ] - scores = tl.sum(scores, axis=0) - # masked store - in_window = (kv_col_offsets >= start_ind) & (kv_col_offsets < end_ind) - tl.store(logits_ptrs, scores, mask=in_window) - - -def fp8_mqa_logits_gfx942( - q: torch.Tensor, - k_fp8: torch.Tensor, - kv_scales: torch.Tensor, - weights: torch.Tensor, - cu_starts: torch.Tensor, - cu_ends: torch.Tensor, -) -> torch.Tensor: - """Compute FP8 MQA logits on MI300X (gfx942) using the vendored kernel. - - Drop-in replacement for ``aiter.ops.triton.attention.fp8_mqa_logits. - fp8_mqa_logits`` on MI300X. Selects ``(BLOCK_KV, num_stages)`` based on - whether the default tile fits within the 64 KiB LDS budget of a gfx942 - CU (see module docstring). - - Args: - q: Query tensor of shape ``[M, H, D]``, FP8 dtype. - k_fp8: Key tensor of shape ``[N, D]``, FP8 dtype. - kv_scales: K scales of shape ``[N]`` (or ``[N, 1]`` -- viewed as - ``[N]``), float32. - weights: Per-head weights of shape ``[M, H]``, float32. - cu_starts: Start indices (inclusive) of shape ``[M]``, int32. - cu_ends: End indices (exclusive) of shape ``[M]``, int32. - - Returns: - Logits of shape ``[M, N]``, float32 -- positions outside - ``[cu_starts[i], cu_ends[i])`` for row ``i`` are pre-filled with - ``-inf`` so the caller can run a top-k without masking. - """ - seq_len, num_heads, head_size = q.shape - seq_len_kv = k_fp8.shape[0] - assert num_heads & (num_heads - 1) == 0, ( - f"num_heads must be a power of two (got {num_heads})" - ) - assert head_size & (head_size - 1) == 0, ( - f"head_size must be a power of two (got {head_size})" - ) - - # The kernel walks ``kv_scales`` as a 1-D contiguous array of size N - # (it indexes by ``kv_scales_ptr + kv_col_offsets``). The vLLM caller - # passes a ``[N, 4]`` uint8 view-cast-to-float32 which lands as - # ``[N, 1]`` contiguous -- byte-identical to ``[N]`` -- but flatten - # explicitly to keep the kernel's pointer arithmetic intent clear. - kv_scales_1d = kv_scales.reshape(-1) - - # Initialise with -inf so positions outside [cu_starts, cu_ends) read - # as ``-inf`` after the masked store path -- this matches AITER's - # ``fp8_mqa_logits`` semantics and is what the top-k consumer expects. - logits = torch.full( - (seq_len, seq_len_kv), - fill_value=-float("inf"), - dtype=torch.float32, - device=q.device, - ) - - if _gfx942_default_tile_fits_lds(num_heads, head_size): - block_kv = 128 - num_stages = 2 - else: - # DSv4 sparse indexer (NUM_HEADS=64, HEAD_SIZE=128) lands here: - # default tile spills past gfx942's 64 KiB LDS budget. (64, 1) - # needs ~33 KiB and clears the per-WG budget with margin. - block_kv = 64 - num_stages = 1 - - # heuristic for MFMA instruction shape, identical to AITER's choice - matrix_instr_nonkdim = 32 - if seq_len <= 1024: - matrix_instr_nonkdim = 16 - - stride_q_s, stride_q_h, stride_q_d = q.stride() - stride_kv_s, stride_kv_d = k_fp8.stride() - stride_w_s, stride_w_h = weights.stride() - stride_logits_s, stride_logits_k = logits.stride() - - _fp8_mqa_logits_kernel[(seq_len,)]( - Q_ptr=q, - KV_ptr=k_fp8, - kv_scales_ptr=kv_scales_1d, - weights_ptr=weights, - cu_start_ptr=cu_starts, - cu_end_ptr=cu_ends, - logits_ptr=logits, - seq_len=seq_len, - seq_len_kv=seq_len_kv, - NUM_HEADS=num_heads, - HEAD_SIZE=head_size, - stride_q_s=stride_q_s, - stride_q_h=stride_q_h, - stride_q_d=stride_q_d, - stride_kv_s=stride_kv_s, - stride_kv_d=stride_kv_d, - stride_w_s=stride_w_s, - stride_w_h=stride_w_h, - stride_logits_s=stride_logits_s, - stride_logits_k=stride_logits_k, - BLOCK_KV=block_kv, - num_warps=4, - num_stages=num_stages, - waves_per_eu=2, - matrix_instr_nonkdim=matrix_instr_nonkdim, - ) - - return logits diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py index 321f42bcb199..68ce0e0de2e2 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py @@ -21,15 +21,11 @@ """ import math -from typing import Any import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2 - -from .triton_turboquant_decode import _use_fp8_e4b15 # ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory # pressure (mirrors the pattern used in triton_decode_attention.py and @@ -66,387 +62,6 @@ def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor: return lut.reshape(N * N, 2).contiguous() -# --------------------------------------------------------------------------- -# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2 -# --------------------------------------------------------------------------- - - -@triton.jit -def _tq_decode_stage1_v2( - Q_rot_ptr, - KV_cache_ptr, - KV_cache_u16_ptr, # uint16 view — Opt#3 SoA loads - Block_table_ptr, - Seq_lens_ptr, - Centroids_ptr, - Pair_lut_ptr, - Mid_o_ptr, - stride_qb, - stride_qh, - stride_cache_block, # bytes per block (bs*H*slot_aligned) - stride_bt_b, - stride_mid_b, - stride_mid_h, - stride_mid_s, - NUM_Q_HEADS: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - PADDED_SLOT: tl.constexpr, - MAX_NUM_BLOCKS: tl.constexpr, - NUM_KV_SPLITS: tl.constexpr, - KV_GROUP_SIZE: tl.constexpr, - MSE_BITS: tl.constexpr, - MSE_BYTES: tl.constexpr, - VQB: tl.constexpr, - VAL_DATA_BYTES: tl.constexpr, - # Opt#3 SoA layout - KEY_DATA_BYTES: tl.constexpr, - META_REGION_OFFSET: tl.constexpr, - NUM_SOA_FIELDS: tl.constexpr, - SOA_K_NORM: tl.constexpr, - SOA_V_SCALE: tl.constexpr, - SOA_V_ZERO: tl.constexpr, - N_CENTROIDS: tl.constexpr, - ATTN_SCALE: tl.constexpr, - BLOCK_D: tl.constexpr, - TILE_SIZE: tl.constexpr, - BLOCK_M: tl.constexpr, - KEY_FP8: tl.constexpr, - NORM_CORRECTION: tl.constexpr = 0, - FP8_E4B15: tl.constexpr = 0, - USE_PAIR_LUT: tl.constexpr = 0, - USE_BF16_DOT: tl.constexpr = 0, -): - # Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group - bid = tl.program_id(0) - kv_hid = tl.program_id(1) - sid = tl.program_id(2) - - seq_len = tl.load(Seq_lens_ptr + bid) - split_len = tl.cdiv(seq_len, NUM_KV_SPLITS) - split_start = split_len * sid - split_end = tl.minimum(split_start + split_len, seq_len) - - if split_start >= split_end: - return - - # Dimension and tile offsets - d_offs = tl.arange(0, BLOCK_D) - d_mask = d_offs < HEAD_DIM - m_offs = tl.arange(0, BLOCK_M) - t_offs = tl.arange(0, TILE_SIZE) - - # Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS) - q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs - q_mask = m_offs < KV_GROUP_SIZE - - # Load Q tile: [BLOCK_M, BLOCK_D] float32 - q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :] - Q = tl.load( - Q_rot_ptr + q_addrs, - mask=q_mask[:, None] & d_mask[None, :], - other=0.0, - ) - - # Pre-warm centroids in L1 - if not KEY_FP8: - _c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS)) # noqa: F841 - - # Loop-invariant bit-extraction constants for MSE keys - if not KEY_FP8: - if MSE_BITS == 4: - pass # 4-bit byte extraction is computed inline in the tile loop - else: - # Generic bit extraction - mse_bit_off = d_offs * MSE_BITS - mse_byte_idx = mse_bit_off // 8 - mse_bit_shift = mse_bit_off % 8 - mse_mask_val = (1 << MSE_BITS) - 1 - - # Loop-invariant bit-extraction for 3-bit values - if VQB == 3: - val_bit_off = d_offs * 3 - val_byte_idx = val_bit_off // 8 - val_bit_shift = val_bit_off % 8 - - # exp2 pre-scaling - RCP_LN2: tl.constexpr = 1.4426950408889634 - LN2: tl.constexpr = 0.6931471805599453 - QK_SCALE = ATTN_SCALE * RCP_LN2 - - # Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L - M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) - L = tl.zeros([BLOCK_M], dtype=tl.float32) - - acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) - - bt_base = bid * stride_bt_b - - # ================================================================ - # TILED LOOP: TILE_SIZE KV tokens per iteration - # ================================================================ - for start_n in range(split_start, split_end, TILE_SIZE): - kv_offs = start_n + t_offs - kv_mask_1d = kv_offs < split_end - - page_idx = kv_offs // BLOCK_SIZE - page_off = kv_offs % BLOCK_SIZE - block_nums = tl.load( - Block_table_ptr + bt_base + page_idx, - mask=kv_mask_1d, - other=0, - ).to(tl.int64) - - # Opt#3 SoA addressing: data region then per-block SoA metadata. - slot_within_block = page_off.to(tl.int64) - block_base = block_nums * stride_cache_block - DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES - data_bases = ( - block_base - + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT) - + tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT - ) - head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast( - kv_hid, tl.int64 - ) * (NUM_SOA_FIELDS * BLOCK_SIZE) - knorm_u16_addrs = ( - head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block - ) - vscale_u16_addrs = ( - head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block - ) - vzero_u16_addrs = ( - head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block - ) - - # ============================================================ - # KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T) - # ============================================================ - if KEY_FP8: - k_addrs = data_bases[:, None] + d_offs[None, :] - k_raw = tl.load( - KV_cache_ptr + k_addrs, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ) - if FP8_E4B15: - K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32) - else: - K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32) - # K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot - K_T = tl.trans(K_f) - # S: [BLOCK_M, TILE_SIZE] (bf16 dot on ROCm for better dynamic range) - if USE_BF16_DOT: - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16)) - else: - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16)) - else: - # --- MSE dequantization --- - if MSE_BITS == 4 and USE_PAIR_LUT: - # OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each - # packed byte loaded EXACTLY ONCE instead of twice (the - # previous code built byte_addrs from d_offs // 2 which has - # the pattern [0,0,1,1,2,2,...], duplicating every byte - # load). We load the packed bytes as a [TILE_SIZE, HALF_D] - # tile, decode both nibbles, then issue a single 3-D gather - # against pair_lut which returns (T[lo], T[hi]) in slots - # 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D] - # interleaves the pairs into the expected dim order - # (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]). - HALF_D: tl.constexpr = BLOCK_D // 2 - half_offs = tl.arange(0, HALF_D) - # Each byte k covers dims (2k, 2k+1); it is in-range iff - # 2k < HEAD_DIM. HEAD_DIM is always even in our layouts, - # so this is equivalent to k < HEAD_DIM // 2. - byte_mask = (half_offs * 2) < HEAD_DIM - - byte_addrs = data_bases[:, None] + half_offs[None, :] - byte_raw = tl.load( - KV_cache_ptr + byte_addrs, - mask=kv_mask_1d[:, None] & byte_mask[None, :], - other=0, - ).to(tl.int32) - - lo_idx = byte_raw & 0xF # [TILE_SIZE, HALF_D] - hi_idx = (byte_raw >> 4) & 0xF # [TILE_SIZE, HALF_D] - pair_key = lo_idx * N_CENTROIDS + hi_idx # [TILE_SIZE, HALF_D] - - # Single 3-D gather: pair_lut[pair_key*2 + {0,1}] - # -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element. - pair_slot = tl.arange(0, 2) - c_pair = tl.load( - Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :], - mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]), - other=0.0, - ) - # Row-major reshape interleaves pairs: - # out[t, 2k] = c_pair[t, k, 0] = T[lo_k] - # out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k] - c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D]) - elif MSE_BITS == 4: - # 4-bit without pair LUT: exploit byte alignment - half_idx = d_offs // 2 - nibble_shift = (d_offs % 2) * 4 - mse_addrs = data_bases[:, None] + half_idx[None, :] - mse_raw = tl.load( - KV_cache_ptr + mse_addrs, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF - c_vals = tl.load( - Centroids_ptr + mse_idx, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0.0, - ) - else: - # Generic 3-bit path: two byte loads → 16-bit → shift/mask - mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :] - mse_raw0 = tl.load( - KV_cache_ptr + mse_addrs0, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - mse_raw1 = tl.load( - KV_cache_ptr + mse_addrs0 + 1, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - raw16 = mse_raw0 | (mse_raw1 << 8) - mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val - c_vals = tl.load( - Centroids_ptr + mse_idx, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0.0, - ) - - # Opt#1: norm-correction is pre-folded into the stored per-token - # scalar at store time, so no per-tile sum+sqrt+divide here. The - # stored vec_norm carries ||k||/||c_vec||; multiply by c_vals - # below gives the original c_vec/||c_vec|| * ||k||. - - # Opt#3: K-norms from per-block SoA region. Single u16 load per - # token; contiguous across tokens within a block. - norm_u16 = tl.load( - KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0 - ) - vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32) - - # Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d] - K_recon = c_vals * vec_norms[:, None] - # K_T: [BLOCK_D, TILE_SIZE] - K_T = tl.trans(K_recon) - - # S: [BLOCK_M, TILE_SIZE] via tensor core MMA - if USE_BF16_DOT: - S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16)) - else: - S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16)) - - # Mask out-of-range positions - S = tl.where(kv_mask_1d[None, :], S, float("-inf")) - - # ============================================================ - # ONLINE SOFTMAX with exp2 - # ============================================================ - m_j = tl.maximum(M, tl.max(S, axis=1)) - m_j = tl.where(m_j > float("-inf"), m_j, 0.0) - P = tl.math.exp2(S - m_j[:, None]) - l_j = tl.sum(P, axis=1) - alpha = tl.math.exp2(M - m_j) - - acc = acc * alpha[:, None] - L = L * alpha + l_j - M = m_j - - # ============================================================ - # VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V) - # Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero - # move to the SoA metadata region (one u16 load each per token). - # ============================================================ - val_bases = data_bases + KEY_DATA_BYTES - - if VQB == 3: - val_addrs0 = val_bases[:, None] + val_byte_idx[None, :] - val_raw0 = tl.load( - KV_cache_ptr + val_addrs0, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - val_raw1 = tl.load( - KV_cache_ptr + val_addrs0 + 1, - mask=kv_mask_1d[:, None] & d_mask[None, :], - other=0, - ).to(tl.int32) - raw16 = val_raw0 | (val_raw1 << 8) - v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32) - else: # VQB == 4 - # OPT#3 (load-halving, value path): same pattern as OPT#1 for - # keys -- load each byte exactly once as [TILE_SIZE, HALF_D], - # decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D] - # with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...]. - V_HALF_D: tl.constexpr = BLOCK_D // 2 - v_half_offs = tl.arange(0, V_HALF_D) - v_byte_mask = (v_half_offs * 2) < HEAD_DIM - - val_addrs = val_bases[:, None] + v_half_offs[None, :] - val_byte = tl.load( - KV_cache_ptr + val_addrs, - mask=kv_mask_1d[:, None] & v_byte_mask[None, :], - other=0, - ).to(tl.int32) - v_lo = (val_byte & 0xF).to(tl.float32) # [T, HALF_D] - v_hi = ((val_byte >> 4) & 0xF).to(tl.float32) # [T, HALF_D] - v_idx = tl.interleave(v_lo, v_hi) # [T, BLOCK_D] - - scale_u16 = tl.load( - KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0 - ) - zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0) - v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32) - v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32) - V = v_idx * v_scales[:, None] + v_zeros[:, None] - - # P·V accumulation via tensor core MMA - # P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D] - if USE_BF16_DOT: - acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16)) - else: - acc += tl.dot(P.to(tl.float16), V.to(tl.float16)) - - # ================================================================ - # EPILOGUE: Store per-Q-head partial results for stage2 - # ================================================================ - # acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M] - # Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1] - safe_L = tl.where(L > 0.0, L, 1.0) - one_over_L = 1.0 / safe_L - acc = acc * one_over_L[:, None] - - # Convert M from log2 space back to natural log for stage2 compatibility - lse = M * LN2 + tl.log(safe_L) - - # Store each Q head's result using 2D scatter - # q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs - out_addrs = ( - bid * stride_mid_b - + q_head_offs[:, None] * stride_mid_h - + sid * stride_mid_s - + d_offs[None, :] - ) - tl.store( - Mid_o_ptr + out_addrs, - acc, - mask=q_mask[:, None] & d_mask[None, :], - ) - # Store LSE per Q head - lse_addrs = ( - bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM - ) - tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask) - - # --------------------------------------------------------------------------- # Launcher v2 # --------------------------------------------------------------------------- @@ -480,190 +95,3 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: shows up on a profile, cache by a hash-of-values fingerprint instead. """ return build_pair_lut(centroids) - - -def triton_turboquant_decode_attention_v2( - query: torch.Tensor, - kv_cache: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - Pi: torch.Tensor, - centroids: torch.Tensor, - scale: float, - mse_bits: int, - key_packed_size: int, - value_quant_bits: int, - value_packed_size: int, - max_seq_len: int = 0, # unused; kept for backward compatibility - key_fp8: bool = False, - norm_correction: bool = False, - PiT: torch.Tensor | None = None, - # Pre-allocated buffers (optional, required for CUDA-graph stability). - mid_o_buf: torch.Tensor | None = None, - output_buf: torch.Tensor | None = None, - lse_buf: torch.Tensor | None = None, - buf_holder: Any = None, - # Fixed split count — MUST be a compile-time constant across iterations - # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract. - max_num_kv_splits: int = 32, -) -> torch.Tensor: - """Launch optimized TQ decode attention (v2 stage1 + shared stage2). - - Follows the same buffer-reuse + fixed-grid contract as the v1 launcher - (triton_turboquant_decode_attention) so the backend can capture a CUDA - graph across both versions. - """ - B, Hq, D = query.shape - Hk = kv_cache.shape[2] - block_size = kv_cache.shape[1] - padded_slot = kv_cache.shape[3] - max_num_blocks = block_table.shape[1] - n_centroids = centroids.shape[0] - kv_group_size = Hq // Hk - device = query.device - del max_seq_len # no longer used: splits is fixed via max_num_kv_splits - - cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size) - - # Opt#3 SoA layout constants (match store-side computation). - key_data_bytes = D if key_fp8 else cfg["mse_bytes"] - data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"] - meta_region_offset = block_size * Hk * data_bytes_per_slot - num_soa_fields = 2 if key_fp8 else 3 - soa_k_norm = 0 - soa_v_scale = 0 if key_fp8 else 1 - soa_v_zero = 1 if key_fp8 else 2 - kv_cache_u16 = kv_cache.view(torch.uint16) - - # Compute q_rot = q @ Pi.T - if key_fp8: - q_rot = query.float().contiguous() - else: - q_float = query.float() - if PiT is None: - PiT = Pi.T.contiguous() - q_rot = (q_float @ PiT).contiguous() - - # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores - BLOCK_M = max(16, triton.next_power_of_2(kv_group_size)) - - # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration - TILE_SIZE = 16 - - # Fixed split count — must be constant across calls for cudagraph replay. - NUM_KV_SPLITS = max_num_kv_splits - - # --- mid_o buffer reuse (same pattern as v1 launcher) --- - if ( - mid_o_buf is not None - and mid_o_buf.shape[0] >= B - and mid_o_buf.shape[2] >= NUM_KV_SPLITS - ): - mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :] - else: - mid_o = torch.empty( - B, - Hq, - NUM_KV_SPLITS, - D + 1, - dtype=torch.float32, - device=device, - ) - if buf_holder is not None: - buf_holder._tq_mid_o_buf = mid_o - - fp8_e4b15 = _use_fp8_e4b15(device.index or 0) - - # Build pair LUT for 4-bit MSE (FLUTE §3.2) - use_pair_lut = mse_bits == 4 and not key_fp8 - pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids - - # Platform-dependent pipelining depth: ROCm prefers num_stages=1. - stage1_num_stages = 1 if _is_hip else 2 - - # --- Stage 1: v2 kernel --- - # Grid over KV heads (not Q heads) — each program handles all Q heads - grid = (B, Hk, NUM_KV_SPLITS) - _tq_decode_stage1_v2[grid]( - q_rot, - kv_cache, - kv_cache_u16, - block_table, - seq_lens, - centroids, - pair_lut, - mid_o, - q_rot.stride(0), - q_rot.stride(1), - kv_cache.stride(0), - block_table.stride(0), - mid_o.stride(0), - mid_o.stride(1), - mid_o.stride(2), - NUM_Q_HEADS=Hq, - NUM_KV_HEADS=Hk, - HEAD_DIM=D, - BLOCK_SIZE=block_size, - PADDED_SLOT=padded_slot, - MAX_NUM_BLOCKS=max_num_blocks, - NUM_KV_SPLITS=NUM_KV_SPLITS, - KV_GROUP_SIZE=kv_group_size, - MSE_BITS=mse_bits, - MSE_BYTES=cfg["mse_bytes"], - VQB=value_quant_bits, - VAL_DATA_BYTES=cfg["val_data_bytes"], - KEY_DATA_BYTES=key_data_bytes, - META_REGION_OFFSET=meta_region_offset, - NUM_SOA_FIELDS=num_soa_fields, - SOA_K_NORM=soa_k_norm, - SOA_V_SCALE=soa_v_scale, - SOA_V_ZERO=soa_v_zero, - N_CENTROIDS=n_centroids, - ATTN_SCALE=scale, - BLOCK_D=cfg["BLOCK_D"], - TILE_SIZE=TILE_SIZE, - BLOCK_M=BLOCK_M, - KEY_FP8=1 if key_fp8 else 0, - NORM_CORRECTION=1 if norm_correction else 0, - FP8_E4B15=fp8_e4b15, - USE_PAIR_LUT=1 if use_pair_lut else 0, - USE_BF16_DOT=1 if _is_hip else 0, - num_warps=4, - num_stages=stage1_num_stages, - ) - - # --- output / lse buffer reuse (same pattern as v1 launcher) --- - if output_buf is not None and output_buf.shape[0] >= B: - output = output_buf[:B, :Hq, :D] - else: - output = torch.empty(B, Hq, D, dtype=torch.float32, device=device) - if buf_holder is not None: - buf_holder._tq_output_buf = output - if lse_buf is not None and lse_buf.shape[0] >= B: - lse = lse_buf[:B, :Hq] - else: - lse = torch.empty(B, Hq, dtype=torch.float32, device=device) - if buf_holder is not None: - buf_holder._tq_lse_buf = lse - - # --- Stage 2: reduce across KV splits (unchanged) --- - grid2 = (B, Hq) - _fwd_kernel_stage2[grid2]( - mid_o, - output, - lse, - seq_lens, - mid_o.stride(0), - mid_o.stride(1), - mid_o.stride(2), - output.stride(0), - output.stride(1), - lse.stride(0), - NUM_KV_SPLITS=NUM_KV_SPLITS, - BLOCK_DV=cfg["BLOCK_D"], - Lv=D, - num_warps=4, - num_stages=2, - ) - - return output.to(query.dtype) diff --git a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py index 03ef36e5b68c..d77592010a4b 100644 --- a/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py +++ b/vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py @@ -876,21 +876,6 @@ def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor: return build_pair_lut(centroids) -# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every -# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is -# tied to the caller's tensor. -_pit_qdtype_cache: dict = {} - - -def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor: - key = (PiT.data_ptr(), qdtype) - cached = _pit_qdtype_cache.get(key) - if cached is None: - cached = PiT.to(qdtype).contiguous() - _pit_qdtype_cache[key] = cached - return cached - - def triton_turboquant_unified_attention( query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16 kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8 diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 28478b69f5e6..c55c06abf1cc 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1608,6 +1608,15 @@ def _get_kv_cache_bytes_per_block( ] if hot_page_sizes: bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes)) + stride_alignments = [ + spec.block_stride_alignment + for group in kv_cache_groups + for layer_name in group.layer_names + if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec) + and spec.block_stride_alignment + ] + if stride_alignments: + bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments)) return bytes_per_block @@ -2669,12 +2678,7 @@ def get_kv_cache_configs( # When speculating with more than 1 speculative module (e.g. multi-layered MTP) # tag every SlidingWindowSpec with how many extra tokens to retain in the window. - extra_retained_tokens = ( - vllm_config.speculative_config.num_speculative_tokens - 1 - if vllm_config.speculative_config is not None - and vllm_config.speculative_config.use_multi_module_mtp() - else 0 - ) + extra_retained_tokens = max(0, vllm_config.num_prefill_lookahead_tokens - 1) for layer_name, layer_spec in merged_kv_cache_specs.items(): if isinstance(layer_spec, SlidingWindowSpec): merged_kv_cache_specs[layer_name] = replace( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 1b17eb3b1518..f1e11e9b1f3d 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,8 +70,7 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats -from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager -from vllm.v1.structured_output.utils import strip_speculative_padding +from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext logger = init_logger(__name__) @@ -263,13 +262,7 @@ def __init__( self.use_eagle_block_drop = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = vllm_config.num_lookahead_tokens - # Positions past the computed tokens that the drafter reads mid-prefill. - # Eagle-family drafters read 1 ahead, but multi-module MTP reads - # num_spec_tokens ahead at chunked-prefill boundaries. Determines the - # encoder scheduling shift, the deferred encoder free, the KV cache - # manager's re-prefillable window (this minus 1), and how many tokens to - # reserve between a chunk boundary and the prefill end. - self.num_prefill_lookahead = 0 + self.num_prefill_lookahead = vllm_config.num_prefill_lookahead_tokens self.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: if speculative_config.num_speculative_tokens_per_batch_size: @@ -279,12 +272,6 @@ def __init__( vllm_num_speculative_tokens=self.num_spec_tokens, ) self.use_eagle = speculative_config.use_eagle() - if self.use_eagle: - self.num_prefill_lookahead = ( - self.num_spec_tokens - if speculative_config.use_multi_module_mtp() - else 1 - ) self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() if self.use_eagle and not self.use_eagle_block_drop: logger.warning( @@ -2052,33 +2039,18 @@ def update_from_output( request.status = RequestStatus.FINISHED_STOPPED stopped = True - if new_token_ids and self.structured_output_manager.should_advance( - request, new_token_ids=new_token_ids + if new_token_ids and not self.structured_output_manager.accept_tokens( + request, new_token_ids ): - struct_output_request = request.structured_output_request - assert struct_output_request is not None - grammar = struct_output_request.grammar - assert isinstance(grammar, StructuredOutputGrammar) - # new_token_ids can be a mixed block of reasoning content, then - # the reasoning end marker, then the start of the grammar content. - # Trim the reasoning content so the grammar only sees grammar content. - advance_token_ids = ( - self.structured_output_manager.trim_reasoning_for_advance( - request, new_token_ids - ) + logger.error( + "Unexpected: grammar rejected tokens %s for request %s. " + "Terminating request.", + new_token_ids, + req_id, ) - if advance_token_ids and not grammar.accept_tokens( - req_id, advance_token_ids - ): - logger.error( - "Unexpected: grammar rejected tokens %s for request %s. " - "Terminating request.", - advance_token_ids, - req_id, - ) - request.status = RequestStatus.FINISHED_ERROR - request.resumable = False - stopped = True + request.status = RequestStatus.FINISHED_ERROR + request.resumable = False + stopped = True routed_experts = None if ( @@ -2101,7 +2073,7 @@ def update_from_output( prompt_start = ( request.sampling_params.routed_experts_prompt_start ) - assert prompt_start < request.num_prompt_tokens + assert prompt_start <= request.num_prompt_tokens else: prompt_start = 0 routed_experts = self.routed_experts_mgr.get( @@ -2434,11 +2406,9 @@ def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: continue # Add newly generated spec token ids to the request. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] - request.spec_token_ids = spec_token_ids + request.spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) def update_draft_token_ids_in_output( self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput @@ -2464,10 +2434,9 @@ def update_draft_token_ids_in_output( # (needed for chunked prefill case for example). del spec_token_ids[orig_num_spec_tokens:] # Filter out spec tokens which do not adhere to the grammar. - if self.structured_output_manager.should_advance(request): - metadata = request.structured_output_request - spec_token_ids = strip_speculative_padding(spec_token_ids) - spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids) # type: ignore[union-attr] + spec_token_ids = self.structured_output_manager.validate_tokens( + request, spec_token_ids + ) # Pad to original number of spec tokens. num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids) if num_invalid_tokens: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index d4e68d8adc86..fc5dc410485d 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2003,7 +2003,9 @@ def cache_blocks( ) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) if self.mamba_cache_mode == "align": - partial_hash = self._cache_partial_tail_block(request, num_tokens) + partial_hash = self._cache_partial_tail_block( + request, num_tokens, retention_interval=retention_interval + ) if partial_hash is not None: self.cached_blocks_this_step.add(partial_hash) if num_cached_blocks_after > num_cached_blocks_before: @@ -2038,6 +2040,7 @@ def _cache_partial_tail_block( self, request: Request, num_tokens: int, + retention_interval: int | None, ) -> BlockHashWithGroupId | None: hash_block_size = self.block_pool.hash_block_size # Re-key the reserved block at its exported checkpoint boundary. @@ -2047,6 +2050,17 @@ def _cache_partial_tail_block( blocks = self.req_to_blocks[request.request_id] assert 0 <= checkpoint_idx < len(blocks) checkpoint_block = blocks[checkpoint_idx] + if ( + retention_interval == 0 + and num_tokens < request.num_prompt_tokens + and checkpoint_position != request.shared_prefix_boundary + ): + # retention_interval == 0 keeps this transient checkpoint + # request-local. The slot may carry a hash from this step's + # full-block pass; that must go too, since the checkpoint + # state is about to overwrite the block. + self.block_pool._maybe_evict_cached_block(checkpoint_block) + return None if checkpoint_block.block_hash_num_tokens == checkpoint_position: return None return self.block_pool.cache_partial_block( diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 5ae9ee0cac83..7d68b6e7978b 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -102,6 +102,8 @@ class EngineCoreReadyResponse: weight_transfer_backend: str | None = None enable_sleep_mode: bool = False supports_draft_weight_updates: bool = False + # Full-attention block size in tokens after initialization, or unavailable. + effective_attention_block_size: int | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/admission_control.py b/vllm/v1/engine/admission_control.py new file mode 100644 index 000000000000..b72d50f3465d --- /dev/null +++ b/vllm/v1/engine/admission_control.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any, Protocol, cast + + +class AdmissionCounterBuffer(Protocol): + def __getitem__(self, index: int) -> int: ... + + def __setitem__(self, index: int, value: int) -> None: ... + + +class SharedAdmissionStats: + """Lock-free admission stats shared by API server processes. + + Each process writes exclusively to its own cache-line-sized slot. Readers + aggregate the slots, avoiding both inter-process locks and cache-line + contention on the request path. + """ + + # Space per-process counters 64 bytes apart to avoid false sharing on CPUs + # with 64-byte cache lines: one process's writes would otherwise invalidate + # the cache line containing another process's counter. + CACHE_LINE_BYTES = 64 + COUNTER_BYTES = 8 + COUNTERS_PER_SLOT = CACHE_LINE_BYTES // COUNTER_BYTES + + def __init__( + self, + client_addresses: dict[str, Any], + client_count: int = 1, + client_index: int = 0, + ): + self._counters = cast( + AdmissionCounterBuffer, + client_addresses.pop("mp_admission_counters"), + ) + self._client_count = client_count + self._slot_offset = client_index * self.COUNTERS_PER_SLOT + + @classmethod + def num_counters(cls, client_count: int) -> int: + return client_count * cls.COUNTERS_PER_SLOT + + def set_num_requests(self, count: int) -> None: + self._counters[self._slot_offset] = count + + def get_num_requests(self) -> int: + return sum( + self._counters[client_index * self.COUNTERS_PER_SLOT] + for client_index in range(self._client_count) + ) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 91618c409362..bf3f5f420ae8 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,6 +44,7 @@ from vllm.utils.async_utils import cancel_task_threadsafe from vllm.utils.collection_utils import as_list from vllm.v1.engine import EngineCoreRequest, PauseMode +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.v1.engine.input_processor import InputProcessor @@ -144,12 +145,20 @@ def __init__( # Convert EngineInput --> EngineCoreRequest. self.input_processor = InputProcessor(self.vllm_config, renderer) + self.admission_stats = ( + SharedAdmissionStats(client_addresses, client_count, client_index) + if client_addresses is not None + and "mp_admission_counters" in client_addresses + else None + ) + # Converts EngineCoreOutputs --> RequestOutput. self.output_processor = OutputProcessor( renderer.tokenizer, log_stats=self.log_stats, stream_interval=self.vllm_config.scheduler_config.stream_interval, tracing_enabled=tracing_endpoint is not None, + admission_stats=self.admission_stats, ) # EngineCore (starts the engine in background process). @@ -291,8 +300,9 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: Both limits return HTTP 503 (Service Unavailable) so that load balancers and client SDKs retry on a different instance. - - ``max_num_queued_reqs``: hard cap on the number of unfinished requests - (waiting + running). A request with ``n > 1`` counts as ``n`` slots. + - ``max_num_queued_reqs``: approximate global cap on the number of + unfinished requests (waiting + running). A request with ``n > 1`` + counts as ``n`` slots. - ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt tokens of requests still in prefill. @@ -312,13 +322,17 @@ def check_admission(self, n: int = 1, request_id: str | None = None) -> None: """ max_num_reqs = self.scheduler_config.max_num_queued_reqs if max_num_reqs is not None: - current = self.get_num_unfinished_requests() - if current + n > max_num_reqs: + current_requests = ( + self.admission_stats.get_num_requests() + if self.admission_stats is not None + else self.get_num_unfinished_requests() + ) + if current_requests + n > max_num_reqs: logger.info( "Request queue full - rejecting request %s " "(current=%d, n=%d, max=%d).", request_id, - current, + current_requests, n, max_num_reqs, ) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ffc9a5de7e68..b3d94de3e7db 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -85,7 +85,7 @@ EngineCoreSentinel, fault_tolerant_wrapper, ) -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig, is_full_attention_spec from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus @@ -168,6 +168,7 @@ def __init__( block_size=scheduler_block_size, hash_block_size=hash_block_size, ) + self._initialize_effective_attention_block_size() self.use_spec_decode = vllm_config.speculative_config is not None self.check_for_draft_tokens = ( self.use_spec_decode or vllm_config.model_config.is_diffusion @@ -382,6 +383,21 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: ) return scheduler_kv_cache_config + def _initialize_effective_attention_block_size(self) -> None: + cache_config = self.vllm_config.cache_config + cache_config.effective_attention_block_size = None + cache_manager = getattr(self.scheduler, "kv_cache_manager", None) + if cache_manager is None: + return + block_sizes = { + manager.block_size + for manager in cache_manager.coordinator.single_type_managers + if is_full_attention_spec(manager.kv_cache_spec) + } + cache_config.effective_attention_block_size = ( + block_sizes.pop() if len(block_sizes) == 1 else None + ) + def get_supported_tasks(self) -> tuple[SupportedTask, ...]: supported_tasks = self.model_executor.supported_tasks self._log_pooler_config(supported_tasks) @@ -1649,6 +1665,9 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, block_size=self.vllm_config.cache_config.block_size, mamba_block_size=self.vllm_config.cache_config.mamba_block_size, + effective_attention_block_size=( + self.vllm_config.cache_config.effective_attention_block_size + ), dp_stats_address=self.frontend_stats_publish_address, dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, @@ -2452,13 +2471,10 @@ def _set_nixl_side_channel_host(): def _set_visible_devices(self, vllm_config: VllmConfig, local_dp_rank: int): from vllm.platforms import current_platform - if current_platform.is_xpu(): - pass - else: - device_control_env_var = current_platform.device_control_env_var - self._set_assigned_physical_gpu_ids( - vllm_config, local_dp_rank, device_control_env_var - ) + device_control_env_var = current_platform.device_control_env_var + self._set_assigned_physical_gpu_ids( + vllm_config, local_dp_rank, device_control_env_var + ) def _set_assigned_physical_gpu_ids( self, diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index c354ba699c75..b63e8016c470 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -566,6 +566,7 @@ def __init__( ): self.vllm_config = vllm_config self._renderer: BaseRenderer | None = renderer + self._effective_attention_block_sizes: set[int | None] = set() # ZMQ setup. sync_ctx = zmq.Context(io_threads=2) @@ -803,10 +804,21 @@ def monitor_engine_cores(): def _apply_ready_response(self, payload: bytes) -> None: """Decode an EngineCoreReadyResponse and sync any post-initialization config changes (e.g. auto-fitted max_model_len) back to the frontend.""" - if not payload: - return vllm_config = self.vllm_config - response = msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) + response = ( + msgspec.msgpack.decode(payload, type=EngineCoreReadyResponse) + if payload + else None + ) + self._effective_attention_block_sizes.add( + response.effective_attention_block_size if response is not None else None + ) + sizes = self._effective_attention_block_sizes + vllm_config.cache_config.effective_attention_block_size = ( + next(iter(sizes)) if len(sizes) == 1 else None + ) + if response is None: + return vllm_config.model_config.max_model_len = min( vllm_config.model_config.max_model_len, response.max_model_len ) diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 4cc433cbfe38..3f86f3c8b239 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -354,12 +354,21 @@ def process_inputs( if isinstance(params, SamplingParams): # TODO: can we avoid cloning here in multiproc case? sampling_params = params.clone() + prompt_len = length_from_prompt_token_ids_or_embeds( + prompt_token_ids, prompt_embeds + ) + if not 0 <= sampling_params.routed_experts_prompt_start <= prompt_len: + raise VLLMValidationError( + f"routed_experts_prompt_start must be between 0 and " + f"the prompt length ({prompt_len}), inclusive.", + parameter="routed_experts_prompt_start", + value=sampling_params.routed_experts_prompt_start, + ) # If unset max tokens, then generate up to the max_model_len. if sampling_params.max_tokens is None: - seq_len = length_from_prompt_token_ids_or_embeds( - prompt_token_ids, prompt_embeds + sampling_params.max_tokens = ( + self.model_config.max_model_len - prompt_len ) - sampling_params.max_tokens = self.model_config.max_model_len - seq_len sampling_params.update_from_generation_config( self.generation_config_fields, @@ -368,12 +377,7 @@ def process_inputs( if self.tokenizer is not None: sampling_params.update_from_tokenizer(self.tokenizer) if sampling_params.trace_decode_token_ids: - self._normalize_trace_replay_params( - sampling_params, - length_from_prompt_token_ids_or_embeds( - prompt_token_ids, prompt_embeds - ), - ) + self._normalize_trace_replay_params(sampling_params, prompt_len) else: pooling_params = params.clone() diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 419e3aa8e99d..257287243dbd 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -5,7 +5,7 @@ from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import numpy as np import torch @@ -41,6 +41,9 @@ ) from vllm.v1.outputs import SamplingMaskLists +if TYPE_CHECKING: + from vllm.v1.engine.admission_control import SharedAdmissionStats + # shared empty CPU tensor used as a placeholder pooling output EMPTY_CPU_TENSOR = torch.empty(0, device="cpu") @@ -455,6 +458,7 @@ def __init__( log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False, + admission_stats: "SharedAdmissionStats | None" = None, ): self.log_stats = log_stats self.tokenizer = tokenizer @@ -464,6 +468,7 @@ def __init__( self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list) self.lora_states = LoRARequestStates(log_stats) self.tracing_enabled = tracing_enabled + self.admission_stats = admission_stats def get_num_unfinished_requests(self): return len(self.request_states) @@ -553,6 +558,7 @@ def abort_requests(self, request_ids: Iterable[str], internal: bool) -> list[str child_reqs = self.abort_requests(child_reqs, internal=True) request_ids_to_abort.extend(child_reqs) self.parent_requests.pop(request_id, None) + self._update_admission_stats() return request_ids_to_abort def add_request( @@ -585,6 +591,7 @@ def add_request( # Track the external_req_id -> [internal_req_id, ...] mapping self.external_req_ids[req_state.external_req_id].append(request_id) + self._update_admission_stats() def _update_streaming_request_state( self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None @@ -763,6 +770,11 @@ def _finish_request(self, req_state: RequestState) -> None: parent_req = req_state.parent_req if parent_req and not parent_req.child_requests: self.parent_requests.pop(parent_req.request_id, None) + self._update_admission_stats() + + def _update_admission_stats(self) -> None: + if self.admission_stats is not None: + self.admission_stats.set_num_requests(self.get_num_unfinished_requests()) def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None): self.lora_states.update_scheduler_stats(scheduler_stats) diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index c23abc9cb4b8..716211fbb30a 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -205,9 +205,12 @@ def bind_hisparse_kv_caches( """Bind existing cache storage and block tables; return the bound handles.""" assert host_pool.registered is not None tensor_configs = { - name: tensor_config + name: ( + tensor_config, + tensor_config.offset + layer_index * tensor_config.layer_stride, + ) for tensor_config in kv_cache_config.kv_cache_tensors - for name in tensor_config.layers + for layer_index, name in enumerate(tensor_config.layers) } resident_source_index = 0 for group_id, group in enumerate(kv_cache_config.kv_cache_groups): @@ -216,12 +219,12 @@ def bind_hisparse_kv_caches( for cache_name in group.layer_names: assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle = _get_hisparse_cache(forward_context, layer_name) cache_handle.bind_cache( kv_caches[cache_name], - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, @@ -249,11 +252,11 @@ def bind_hisparse_kv_caches( raise RuntimeError("HiSparse hot tensors must share one GPU backing.") layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] cache_handle = _get_hisparse_cache(forward_context, layer_name) - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle.runtime.bind_hot_cache( raw_tensor, - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 314d3ba86b7a..5100379083c3 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -16,7 +16,7 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed import get_tp_group -from vllm.forward_context import get_forward_context +from vllm.forward_context import get_forward_context, in_piecewise_cudagraph from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.math_utils import round_up @@ -975,7 +975,7 @@ def _resolve_and_stage_group( ) -> None: group = self.index_group compute_stream = current_stream() - if group.logical_topk_ready is not None: + if group.logical_topk_ready is not None and not in_piecewise_cudagraph(): group.copy_stream.wait_event(group.logical_topk_ready) else: group.copy_stream.wait_stream(compute_stream) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8cd867eb9c60..7e91300b2fe2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -653,6 +653,12 @@ class MLAAttentionSpec(FullAttentionSpec): is_index_group_leader: bool = False storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" + block_stride_alignment: int | None = None + """Required alignment, in bytes, of the distance between consecutive + blocks of this cache. In block-major layouts that distance is the whole + block (all layers' pages), so the allocator rounds the block up to it. + DeepGEMM's paged sparse MQA-logits kernels address pages as + ``base + page * stride`` and need it 512B-aligned.""" # Group capability enabled when any member flattens a non-causal query block # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False @@ -674,6 +680,7 @@ def merge(cls, specs: list[Self]) -> Self: cache_role_set = {spec.cache_role for spec in specs} index_group_leader_set = {spec.is_index_group_leader for spec in specs} storage_block_size_set = set(spec.storage_block_size for spec in specs) + block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 @@ -681,10 +688,11 @@ def merge(cls, specs: list[Self]) -> Self: and len(cache_role_set) == 1 and len(index_group_leader_set) == 1 and len(storage_block_size_set) == 1 + and len(block_stride_alignment_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, cache role, " - "index-sharing role, and storage block size." + "index-sharing role, storage block size and block stride alignment." ) merged_spec = cls( block_size=specs[0].block_size, @@ -701,6 +709,7 @@ def merge(cls, specs: list[Self]) -> Self: cache_role=cache_role_set.pop(), is_index_group_leader=index_group_leader_set.pop(), storage_block_size=storage_block_size_set.pop(), + block_stride_alignment=block_stride_alignment_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs ), diff --git a/vllm/v1/kv_offload/tiering/factory.py b/vllm/v1/kv_offload/tiering/factory.py index 17f8430204d9..3310167b870b 100644 --- a/vllm/v1/kv_offload/tiering/factory.py +++ b/vllm/v1/kv_offload/tiering/factory.py @@ -123,3 +123,9 @@ def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]: "vllm.v1.kv_offload.tiering.obj.manager", "ObjectStoreSecondaryTierManager", ) + +SecondaryTierFactory.register_tier( + "kvcr", + "vllm.v1.kv_offload.tiering.kvcr.manager", + "KVCRSecondaryTierManager", +) diff --git a/vllm/v1/kv_offload/tiering/kvcr/__init__.py b/vllm/v1/kv_offload/tiering/kvcr/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py new file mode 100644 index 000000000000..00186279a4bc --- /dev/null +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -0,0 +1,635 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM secondary-tier adapter for KVCR.""" + +import ctypes +import mmap +import socket +import time +import uuid +from collections.abc import Collection, Iterable, Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +from kvcr import ( + DURATION_METRIC, + KVCR, + ROUTER_HINT_CAPABILITIES, + ROUTER_HINT_KEY, + STATE_METRIC, + TRANSFER_BLOCKS_METRIC, + TRANSFER_BYTES_METRIC, + KVCRBindings, +) +from kvcr.config import ( + FrameworkDramInput, + G3Options, + KeyAdapter, + KVCRBackendConfigs, + KVCRConfig, + KVCRGuardConfig, + LocalDramOptions, + RemoteFWDramOptions, +) +from kvcr.control_channels import ZmqPeerControlChannel +from kvcr.policy import ( + FIFOPolicy, + G3FIFOPolicy, + G3LRUPolicy, + KVCachePolicy, + LRUPolicy, +) +from kvcr.types import ( + BlockKey, + CacheTier, + InventoryEvent, + MemDescriptor, + OpHandle, + PinRequestId, + PinResult, + QueryStatus, +) +from typing_extensions import override + +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( + OffloadingConnectorStats, +) +from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname +from vllm.v1.core.kv_cache_utils import ( + BlockHash, + maybe_convert_block_hash, +) +from vllm.v1.kv_offload.base import ( + LookupResult, + Medium, + OffloadingCounterMetadata, + OffloadingEvent, + OffloadingGaugeMetadata, + OffloadingHistogramMetadata, + OffloadingMetricMetadata, + OffloadKey, + ReqContext, + RequestOffloadingContext, + get_offload_block_hash, +) +from vllm.v1.kv_offload.tiering.base import ( + JobResult, + ParentManager, + SecondaryTierManager, + TransferJob, +) + +if TYPE_CHECKING: + from vllm.v1.kv_offload.base import OffloadingSpec + + +_REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES + +logger = init_logger(__name__) + +_BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = { + "fifo": FIFOPolicy, + "lru": LRUPolicy, + "g3_fifo": G3FIFOPolicy, + "g3_lru": G3LRUPolicy, +} + + +def _resolve_policy(name: str | None) -> KVCachePolicy | None: + if name is None: + return None + policy_type = _BUILTIN_POLICIES.get(name) + if policy_type is None: + if "." not in name: + raise ValueError( + f"Unknown KVCR policy: {name!r}. " + f"Supported: {list(_BUILTIN_POLICIES)}; external policies " + "must use a fully qualified module.Class path." + ) + policy_type = resolve_obj_by_qualname(name) + if not isinstance(policy_type, type) or not issubclass( + policy_type, KVCachePolicy + ): + raise TypeError(f"{name} is not a KVCachePolicy") + return policy_type() + + +def _vllm_metric_name(name: str) -> str: + return f"vllm:{name}" + + +def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]: + return { + _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata( + documentation="KVCR operation and stage duration in seconds.", + labelnames=("scope", "result"), + ), + _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata( + documentation="Bytes transferred by KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata( + documentation="KV blocks transferred by successful KVCR operations.", + labelnames=("operation",), + ), + _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata( + documentation="Current KVCR metadata and operation counts.", + labelnames=("resource",), + ), + } + + +class _FrameworkPinAdapter: + """Adapt KVCR framework pin callbacks to primary-tier store jobs.""" + + def __init__(self, tier: "KVCRSecondaryTierManager") -> None: + self._tier = tier + self._next_request_id = 0 + self._next_pin_request_id = 0 + self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {} + self._completed_pins: list[tuple[PinRequestId, PinResult]] = [] + self._pin_jobs: dict[str, int] = {} + self._pin_job_results: list[JobResult] = [] + + def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId: + # The parent callback is valid only from serve_external_requests(). + request = PinRequestId(self._next_pin_request_id) + self._next_pin_request_id += 1 + self._pending_pins[request] = tuple(keys) + return request + + def process_pending(self, parent: ParentManager) -> None: + pending = self._pending_pins + self._pending_pins = {} + self._completed_pins.extend( + (request, self._resolve_pin(parent, keys)) + for request, keys in pending.items() + ) + + def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]: + completed = self._completed_pins + self._completed_pins = [] + return completed + + def cancel_pin_request(self, request: PinRequestId) -> None: + self._pending_pins.pop(request, None) + + def _resolve_pin( + self, parent: ParentManager, keys: Collection[BlockKey] + ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None: + offload_keys = tuple(OffloadKey(bytes(key)) for key in keys) + request_id = f"kvcr-source:{self._next_request_id}" + self._next_request_id += 1 + req_context = ReqContext(req_id=request_id) + started = False + job: TransferJob | None = None + pin_handle: str | None = None + try: + parent.on_new_request(req_context) + started = True + hit_keys = tuple( + key + for key in offload_keys + if parent.lookup(key, req_context) is LookupResult.HIT + ) + if not hit_keys: + return None + + job = parent.create_store_job(hit_keys, req_context) + job_keys = tuple(job.keys) + chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids) + if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): + return None + + descriptors: dict[BlockKey, list[MemDescriptor] | None] = { + BlockKey(bytes(key)): None for key in offload_keys + } + descriptors.update( + ( + BlockKey(bytes(key)), + [self._tier._make_descriptor(chunk_id)], + ) + for key, chunk_id in zip(job_keys, chunk_ids) + ) + pin_handle = f"kvcr-job-{job.job_id}" + self._pin_jobs[pin_handle] = job.job_id + return pin_handle, descriptors + except Exception: + logger.warning("KVCR source acquisition failed", exc_info=True) + return None + finally: + if started: + try: + parent.on_request_finished(req_context) + except Exception: + logger.warning("KVCR source request cleanup failed", exc_info=True) + if job is not None and pin_handle is None: + self._pin_job_results.append( + JobResult(job_id=job.job_id, success=False) + ) + + def release_pin(self, pin_handle: str) -> bool: + job_id = self._pin_jobs.pop(pin_handle, None) + if job_id is not None: + self._pin_job_results.append(JobResult(job_id=job_id, success=True)) + return True + + def has_active_pins(self) -> bool: + return bool(self._pin_jobs) + + def take_pin_job_results(self) -> list[JobResult]: + results = self._pin_job_results + self._pin_job_results = [] + return results + + +class _VllmKeyAdapter: + """Adapt vLLM's local key format to KVCR.""" + + def encode(self, framework_key: object) -> BlockKey: + if not isinstance(framework_key, bytes): + raise TypeError("vLLM offload keys must be bytes") + return BlockKey(framework_key) + + def decode(self, key: BlockKey) -> int | bytes: + block_hash = BlockHash(get_offload_block_hash(OffloadKey(key))) + return maybe_convert_block_hash(block_hash) + + +# Job ID, remaining blocks, aggregate success, and successful load keys. +_JobState = tuple[int, int, bool, set[OffloadKey] | None] + + +class KVCRSecondaryTierManager(SecondaryTierManager): + """Secondary tier wrapper around the KVCR KV P2P API.""" + + @classmethod + @override + def build_metric_definitions( + cls, extra_config: dict[str, Any] + ) -> dict[str, OffloadingMetricMetadata]: + if not extra_config.get("enable_telemetry", False): + return {} + return _kvcr_metric_definitions() + + def __init__( + self, + offloading_spec: "OffloadingSpec", + primary_kv_view: memoryview, + tier_type: str, + router_capabilities: Iterable[str] | None = None, + control_host: str = "0.0.0.0", + control_ports: list[int] | None = None, + control_advertise_host: str | None = None, + eager_ctrl_connect: bool = True, + opportunistic_query: bool = True, + enable_telemetry: bool = False, + operation_timeout_ms: int = 1000, + abandon_timeout_ms: int = 5000, + metadata_retry_interval_ms: int = 100, + secondary_g2_slots: int = 0, + kvcr_service_socket_path: str | None = None, + compatibility_digest: str | None = None, + policy: str | None = None, + g3: dict[str, Any] | None = None, + local_dram_backend: str = "UCX", + remote_fw_dram_backend: str = "UCX", + ) -> None: + super().__init__(offloading_spec, primary_kv_view, tier_type) + selected_policy = _resolve_policy(policy) + if (kvcr_service_socket_path is None) != (compatibility_digest is None): + raise ValueError( + "kvcr_service_socket_path and compatibility_digest must be " + "configured together" + ) + missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ()) + if missing: + raise ValueError( + f"KVCR secondary tier requires router capabilities: {sorted(missing)}" + ) + events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events + self_describing_events = ( + offloading_spec.kv_events_config.self_describing_kv_events + ) + if ( + events_enabled + and ( + secondary_g2_slots > 0 + or kvcr_service_socket_path is not None + or g3 is not None + ) + and not self_describing_events + ): + raise ValueError( + "KVCR local G2/G3 inventory requires self_describing_kv_events" + ) + self._key_adapter: KeyAdapter = _VllmKeyAdapter() + advertise_host = control_advertise_host or socket.gethostname() + dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local + if dp_local_rank is None: + if control_ports is None or len(control_ports) != 1: + raise ValueError( + "control_ports must contain exactly one port when " + "data_parallel_rank_local is unset" + ) + dp_local_rank = 0 + if control_ports is None or not 0 <= dp_local_rank < len(control_ports): + raise ValueError( + "control_ports must contain a port for " + f"local data-parallel rank {dp_local_rank}" + ) + if any(not 1 <= port <= 65535 for port in control_ports): + raise ValueError("control_ports must be between 1 and 65535") + control = ZmqPeerControlChannel( + control_host, + control_ports[dp_local_rank], + advertise_host, + ) + # Give colocated workers distinct transport listen ports. + with socket.socket() as _s: + _s.bind(("", 0)) + _nixl_listen_port = _s.getsockname()[1] + self._primary_base_addr = ctypes.addressof( + ctypes.c_char.from_buffer(primary_kv_view) + ) + if primary_kv_view.strides is None: + raise ValueError("primary KV memoryview must expose strides") + self._primary_row_stride = int(primary_kv_view.strides[0]) + if secondary_g2_slots < 0: + raise ValueError("secondary_g2_slots must be non-negative") + if ( + secondary_g2_slots or kvcr_service_socket_path is not None + ) and not events_enabled: + logger.warning( + "KVCR local DRAM is enabled, but KV cache events are disabled; " + "local DRAM inventory will not be published" + ) + local_mapping: mmap.mmap | None = None + local_dram: LocalDramOptions | None = None + if secondary_g2_slots: + if kvcr_service_socket_path is not None: + logger.warning( + "secondary_g2_slots is ignored when " + "kvcr_service_socket_path is configured" + ) + else: + local_mapping = mmap.mmap( + -1, secondary_g2_slots * self._primary_row_stride + ) + local_dram = LocalDramOptions( + pools=[ + ( + "", + ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)), + len(local_mapping), + ) + ], + backend=local_dram_backend, + ) + guard_config = ( + KVCRGuardConfig( + kvcr_service_socket_path=kvcr_service_socket_path, + guard_index=dp_local_rank, + compatibility_digest=compatibility_digest, + ) + if kvcr_service_socket_path is not None and compatibility_digest is not None + else None + ) + g3_config = None + if g3 is not None: + g3_values = dict(g3) + g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"]) + g3_config = G3Options(**g3_values) + nixl_agent_name = f"KVCR-{uuid.uuid4()}" + self._framework_pin_adapter = _FrameworkPinAdapter(self) + self._inventory_events: list[OffloadingEvent] = [] + + try: + self._kvcr = KVCR( + KVCRConfig( + pool_layouts=[("", self._primary_row_stride)], + enable_telemetry=enable_telemetry, + operation_timeout_ms=operation_timeout_ms, + abandon_timeout_ms=abandon_timeout_ms, + nixl_agent_name=nixl_agent_name, + nixl_listen_port=_nixl_listen_port, + ), + KVCRBindings( + request_pin=self._framework_pin_adapter.request_pin, + poll_pin_results=self._framework_pin_adapter.poll_pin_results, + release_pin=self._framework_pin_adapter.release_pin, + cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request), + framework_control=control, + key_adapter=self._key_adapter, + inventory_sink=(self._record_inventory if events_enabled else None), + stats_factory=( + OffloadingConnectorStats if enable_telemetry else None + ), + policy=selected_policy, + ), + KVCRBackendConfigs( + framework_dram=FrameworkDramInput( + self._primary_base_addr, primary_kv_view.nbytes + ), + local_dram=local_dram, + g3=g3_config, + remote_fw_dram=RemoteFWDramOptions( + eager_ctrl_connect=eager_ctrl_connect, + opportunistic_query=opportunistic_query, + metadata_retry_interval_ms=metadata_retry_interval_ms, + backend=remote_fw_dram_backend, + ), + ), + guard_config, + ) + except BaseException: + control.close() + if local_mapping is not None: + local_mapping.close() + raise + self._local_dram_mmap = local_mapping + self._finished_jobs: list[JobResult] = [] + self._jobs_by_op: dict[OpHandle, _JobState] = {} + + @override + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + block_key = self._key_adapter.encode(key) + status, _ = self._kvcr.query((block_key,), req_context.req_id)[0] + if status is QueryStatus.FETCHING: + return LookupResult.RETRY + if status in (QueryStatus.HIT, QueryStatus.FETCHABLE): + return LookupResult.HIT + return LookupResult.MISS + + @override + def submit_load(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + + op_handle = self._kvcr.deliver( + blocks, request_id=job_metadata.req_context.req_id + ) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + set(), + ) + + @override + def submit_store(self, job_metadata: TransferJob) -> None: + blocks = { + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + for key, chunk_id in zip( + job_metadata.keys, job_metadata.chunk_ids, strict=True + ) + } + if not blocks: + self._finished_jobs.append( + JobResult(job_id=job_metadata.job_id, success=True) + ) + return + op_handle = self._kvcr.deposit(blocks) + self._jobs_by_op[op_handle] = ( + job_metadata.job_id, + len(blocks), + True, + None, + ) + + @override + def get_finished_jobs(self) -> Iterable[JobResult]: + results = self._finished_jobs + self._finished_jobs = [] + results.extend(self._poll_finished_jobs()) + return results + + @override + def get_stats(self) -> OffloadingConnectorStats | None: + stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats()) + if stats is not None: + # Add the framework namespace only at the vLLM boundary. + for metrics in stats.data.values(): + for name in tuple(metrics): + metrics[_vllm_metric_name(name)] = metrics.pop(name) + return stats + + @override + def serve_external_requests(self, parent: ParentManager) -> None: + self._framework_pin_adapter.process_pending(parent) + + def _poll_finished_jobs(self) -> list[JobResult]: + results: list[JobResult] = [] + for op_handle, entries in self._kvcr.poll_completed(): + job_state = self._jobs_by_op.get(op_handle) + if job_state is None: + continue + job_id, remaining, success, successful_keys = job_state + remaining -= len(entries) + success = success and all(entry.success for entry in entries.values()) + if successful_keys is not None: + successful_keys.update( + OffloadKey(bytes(key)) + for key, entry in entries.items() + if entry.success + ) + if remaining > 0: + self._jobs_by_op[op_handle] = ( + job_id, + remaining, + success, + successful_keys, + ) + continue + self._jobs_by_op.pop(op_handle, None) + results.append( + JobResult( + job_id=job_id, + success=success, + successful_keys=successful_keys if not success else None, + ) + ) + results.extend(self._framework_pin_adapter.take_pin_job_results()) + return results + + @override + def has_pending_work(self) -> bool: + # Keep scheduler polling: control messages may arrive without a load job. + return True + + @override + def take_events(self) -> Iterable[OffloadingEvent]: + events = self._inventory_events + self._inventory_events = [] + return events + + @override + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: + params = getattr(req_context, "kv_transfer_params", None) + if isinstance(params, Mapping): + hint = params.get(ROUTER_HINT_KEY) + if hint is not None: + self._kvcr.submit_hint( + request_id=req_context.req_id, + hints=hint, + ) + return RequestOffloadingContext() + + @override + def on_request_finished(self, req_context: ReqContext) -> None: + self._kvcr.discard_hint(req_context.req_id) + + @override + def drain_jobs(self) -> None: + while self._jobs_by_op or self._framework_pin_adapter.has_active_pins(): + drained = self._poll_finished_jobs() + if drained: + self._finished_jobs.extend(drained) + else: + time.sleep(0.001) + + @override + def shutdown(self) -> None: + self.drain_jobs() + # TODO: Keep registered buffers alive if KVCR.close() fails. + self._kvcr.close() + if self._local_dram_mmap is not None: + self._local_dram_mmap.close() + self._local_dram_mmap = None + + def _record_inventory(self, event: InventoryEvent) -> None: + if event.tier is CacheTier.LOCAL_G2: + medium = Medium.CPU + elif event.tier is CacheTier.G3: + medium = Medium.STORAGE + else: + raise ValueError(f"unsupported KVCR inventory tier: {event.tier}") + self._inventory_events.append( + OffloadingEvent( + keys=[OffloadKey(bytes(key)) for key in event.keys], + medium=medium, + removed=event.removed, + ownership="kvcr", + removal_expected=True, + ) + ) + + def _make_descriptor(self, chunk_id: int) -> MemDescriptor: + return MemDescriptor( + end_point_name=self._kvcr.config.nixl_agent_name, + mem_type="DRAM", + addr=self._primary_base_addr + chunk_id * self._primary_row_stride, + size=self._primary_row_stride, + device_Id=0, + info="", + ) diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 2cd490171006..ba1973b76784 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -47,10 +47,18 @@ # Reap unbound store batches that have been parked without a FetchMsg # binding them to a session for longer than this. Protects against the # prefiller buffering blocks for a decoder that never asks (decoder died, -# network partition, lost kv_request_id). Must be longer than the per-store -# deadline so the store-timeout path fires first for individual jobs. +# network partition, lost kv_request_id). Applies only while a batch is +# unbound: binding hands it to the session, which times the job out under +# _STORE_TIMEOUT_S instead. +# Overridable per tier via the ``unbound_store_timeout_s`` config key. _UNBOUND_STORE_TIMEOUT_S = 60.0 +# How long a reaped kv_request_id is remembered so a fetch that arrives +# after the reap can be rejected immediately. Must exceed the consumer's +# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has +# already given up on its own, so there is nobody left to reject. +_REAPED_ID_RETENTION_S = 60.0 + # Time we wait during shutdown for inflight transfers to drain via # cancel(mode="wait") before falling back to mode="immediate". Bounded # so a wedged peer can't hang shutdown. @@ -208,6 +216,7 @@ def __init__( port: int | None = None, backends: list[str] | None = None, num_threads: int = 4, + unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S, **kwargs: Any, ) -> None: """Initialize the P2P secondary tier manager. @@ -245,9 +254,29 @@ def __init__( num_threads: NIXL agent worker threads for the UCX-only branch. Ignored when ``backends`` contains a non-UCX entry. + unbound_store_timeout_s: Seconds a producer holds stored blocks + for a consumer that has not fetched them yet, before + ``_reap_unbound_stores`` drops them. Raise it for + deployments whose prefills outlast the default; the blocks + keep primary-tier slots pinned for the whole window. **kwargs: Reserved for future tier-specific options. + + Raises: + ValueError: If ``unbound_store_timeout_s`` is not a positive + number, or anything convertible to one. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + try: + timeout_s = float(unbound_store_timeout_s) + except (TypeError, ValueError): + timeout_s = float("nan") + if not timeout_s > 0: + raise ValueError( + f"unbound_store_timeout_s must be a positive number, got " + f"{unbound_store_timeout_s!r}" + ) + self._unbound_store_timeout_s = timeout_s + # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py). # Peers whose seeds differ compute different hashes for identical # content, so lookups silently miss and no KV crosses the wire. The @@ -303,8 +332,16 @@ def __init__( # kv_request_id → list of batches submit_store'd before any peer # asked for that id. Drained into a session by _on_session_fetch # when the corresponding FetchMsg arrives, or surfaced as failures - # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S. + # by _reap_unbound_stores after _unbound_store_timeout_s. self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {} + # kv_request_id → time its parked batches were reaped. Nothing under + # one of these ids can be served any more, so a FetchMsg for it is + # rejected on the spot by _poll_once instead of parking demand until + # the consumer's own load timeout, and submit_store fails rather than + # re-parking under it — a prefill that outran the timeout would + # otherwise keep pinning fresh primary-tier slots, one timeout at a + # time. Entries are pruned after _REAPED_ID_RETENTION_S. + self._reaped_stores: dict[str, float] = {} self._finished_jobs: list[JobResult] = [] # kv_request_ids that hit a transport/session failure; On load lookup() @@ -390,7 +427,7 @@ def on_request_finished(self, req_context: ReqContext) -> None: kv_transfer_params; if a session has bound the id, finish it. If no session has bound the id yet, this is a no-op: parked batches in `_unbound_stores` are left in place and cleaned up only by - `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`. + `_reap_unbound_stores` after `_unbound_store_timeout_s`. """ source = req_context.get_state(P2PSourceInfo) dest = req_context.get_state(P2PDestInfo) @@ -447,6 +484,22 @@ def submit_store(self, job_metadata: TransferJob) -> None: self._finished_jobs.append(JobResult(job_id=job_id, success=False)) return + # This id's earlier batches were already reaped, so the consumer has + # either given up or will be rejected by the fetch path. Parking more + # would pin primary-tier slots for another full timeout on blocks + # nobody can fetch, so fail the job now — a prefill that outran the + # timeout falls back to local recompute on the decoder either way. + if kv_request_id in self._reaped_stores: + logger.warning( + "P2P %s: submit_store for reaped kv_request_id=%s job_id=%d " + "— failing without parking", + self._local_id, + kv_request_id, + job_id, + ) + self._finished_jobs.append(JobResult(job_id=job_id, success=False)) + return + # Fast path: a session has already received FetchMsg for this id, # so we can route the batch straight into its ServerRole. session = self._kv_to_session.get(kv_request_id) @@ -707,13 +760,27 @@ def _reap_unbound_stores(self) -> None: """Time out submit_store batches that no peer has ever fetched. Walks `_unbound_stores` for entries whose oldest batch is older - than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces - every batched job as failed, and adds the id to `_failed_req_ids` - so a late inbound FetchMsg short-circuits to a clean rejection. + than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces + every batched job as failed, and records the id in `_reaped_stores` + so a late inbound FetchMsg is rejected in one round trip and a late + `submit_store` fails instead of re-parking under the reaped id. + + Also expires `_reaped_stores` entries whose consumer can no longer + be waiting, which is why this runs before the empty check below. """ + now = time.monotonic() + if self._reaped_stores: + retention_deadline = now - _REAPED_ID_RETENTION_S + for kid in [ + kid + for kid, reaped_at in self._reaped_stores.items() + if reaped_at <= retention_deadline + ]: + del self._reaped_stores[kid] + if not self._unbound_stores: return - deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S + deadline = now - self._unbound_store_timeout_s expired: list[str] | None = None for kid, batches in self._unbound_stores.items(): # Batches are appended in arrival order, so the head is oldest. @@ -725,7 +792,7 @@ def _reap_unbound_stores(self) -> None: return for kid in expired: batches = self._unbound_stores.pop(kid) - self._failed_req_ids.add(kid) + self._reaped_stores[kid] = now for batch in batches: self._finished_jobs.append( JobResult(job_id=batch.job_id, success=False) @@ -735,7 +802,7 @@ def _reap_unbound_stores(self) -> None: "without a fetch — failing %d job(s)", self._local_id, kid, - _UNBOUND_STORE_TIMEOUT_S, + self._unbound_store_timeout_s, len(batches), ) @@ -779,6 +846,24 @@ def _poll_once(self) -> None: # inline in dispatch, so the replayed add_stored_blocks calls # match that demand and submit transfers immediately. for kv_request_id in result.new_fetch_ids: + if kv_request_id in self._reaped_stores: + # This id's blocks were already reaped, so no + # submit_store will ever satisfy the demand on_fetch + # just recorded. Finalize the round now: the peer gets + # TransferDoneMsg(success=False) on this tick and falls + # back to local prefill, instead of parking until its + # own load timeout expires. The id is deliberately left + # unbound so any later submit_store keeps taking the + # reject path. Batches re-parked between the reap and the + # retention prune are failed here rather than left to a + # second reap, which would pin their slots for another + # full timeout. + session.finish_request(kv_request_id) + for batch in self._unbound_stores.pop(kv_request_id, ()): + self._finished_jobs.append( + JobResult(job_id=batch.job_id, success=False) + ) + continue self._kv_to_session[kv_request_id] = session for batch in self._unbound_stores.pop(kv_request_id, ()): session.add_stored_blocks( @@ -810,6 +895,7 @@ def shutdown(self) -> None: JobResult(job_id=batch.job_id, success=False) ) self._unbound_stores.clear() + self._reaped_stores.clear() self._control.close() self._data.close() diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 5aa83454739e..85b623bd84fa 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -19,6 +19,7 @@ triton_kernel_dispatcher_with_warmup, ) from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput from vllm.v1.sample.logits_processor.builtin import MinTokensLogitsProcessor from vllm.v1.sample.metadata import SamplingMetadata @@ -327,9 +328,7 @@ def apply_logits_processors( num_draft_tokens = torch.tensor(metadata.num_draft_tokens, device="cpu") original_indices = torch.arange(num_requests, device="cpu") repeat_indices_cpu = original_indices.repeat_interleave(num_draft_tokens) - repeat_indices = repeat_indices_cpu.to( - device=logits.device, non_blocking=True - ) + repeat_indices = async_tensor_h2d(repeat_indices_cpu, logits.device) logits = self.apply_penalties( logits, sampling_metadata, metadata, repeat_indices, output_token_ids ) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index c0e456521ca0..1bd1bdc45bbd 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -255,7 +255,6 @@ def __init__( self.backup_next_token_ids = CpuGpuBuffer( self.max_batch_size, dtype=torch.int32, - pin_memory=PIN_MEMORY, device=device, with_numpy=True, ) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index ca6adc736d06..7fd311c873c1 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import itertools import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor @@ -18,6 +17,7 @@ StructuredOutputGrammar, ) from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend +from vllm.v1.structured_output.utils import strip_speculative_padding if TYPE_CHECKING: import numpy as np @@ -217,6 +217,100 @@ def _async_submit_fill_bitmask( ) -> Future: return self.executor_for_fillmask.submit(self._fill_bitmasks, batch) + def _get_constraint_start( + self, + request: "Request", + spec_tokens: Sequence[int], + spec_tokens_committed: bool = False, + ) -> int: + """Return the index into `spec_tokens` where tokens should start being + constrained by the grammar. + Assumes `spec_tokens` are stripped of -1 padding tokens. + Returns `len(spec_tokens) + 1` if no token should be constrained after + accepting all `spec_tokens`. + + `spec_tokens_committed` is True if spec_tokens is already in + `request.all_token_ids`. + """ + if self.enable_in_reasoning: + return 0 + + structured_req = request.structured_output_request + assert structured_req is not None + if structured_req.reasoning_ended: + return 0 + + reasoner = self._get_reasoner(request) + if reasoner is None: + return 0 + + if structured_req.reasoning_ended is None: + # This should be removed here, but since `openai_gptoss` + # is an independent code path, it is kept for now. + # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, + # it can be removed. + if reasoner.is_reasoning_end(request.prompt_token_ids or []): + return 0 + structured_req.reasoning_ended = False + + num_spec_tokens = len(spec_tokens) + if num_spec_tokens <= 0: + return num_spec_tokens + 1 + + # Use `find_reasoning_end_offset` to find constraint start if supported + if ( + isinstance(reasoner, ParserEngineReasoningAdapter) + and reasoner.reasoning_end_token_ids + ): + offset = reasoner.find_reasoning_end_offset(spec_tokens) + if offset is not None: + return offset + 1 + + # Fallback to `find_reasoning_end_offset` + # TODO: Build a read-only Sequence view over all_token_ids + # instead of copying the entire all_token_ids into input_ids. + if spec_tokens_committed: + if not spec_tokens or not reasoner.is_reasoning_end_streaming( + request.all_token_ids, spec_tokens + ): + return num_spec_tokens + 1 + + input_ids = request.all_token_ids.copy() + delta_ids = list(spec_tokens) + else: + input_ids = request.all_token_ids.copy() + input_ids.extend(spec_tokens) + if not reasoner.is_reasoning_end_streaming(input_ids, spec_tokens): + return num_spec_tokens + 1 + delta_ids = list(spec_tokens) + + for i in range(num_spec_tokens - 1, 0, -1): + input_ids.pop() + delta_ids.pop() + if not reasoner.is_reasoning_end_streaming(input_ids, delta_ids): + return i + 1 + return 1 + + def validate_tokens(self, request: "Request", spec_tokens: list[int]) -> list[int]: + """Return the longest unconstrained or grammar-valid prefix of `spec_tokens`.""" + if not request.use_structured_output: + return spec_tokens + + spec_tokens = strip_speculative_padding(spec_tokens) + constraint_start = self._get_constraint_start(request, spec_tokens) + if constraint_start >= len(spec_tokens): + return spec_tokens + + structured_req = request.structured_output_request + if TYPE_CHECKING: + assert structured_req is not None + grammar = structured_req.grammar + if TYPE_CHECKING: + assert isinstance(grammar, StructuredOutputGrammar) + prefix = spec_tokens[:constraint_start] + validated = grammar.validate_tokens(spec_tokens[constraint_start:]) + return prefix + validated + def grammar_bitmask( self, requests: dict[str, "Request"], @@ -264,7 +358,7 @@ def grammar_bitmask( if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) + apply_bitmask = self._get_constraint_start(request, ()) == 0 batch.append((grammar, cumulative_index, apply_bitmask)) if len(batch) == self.fill_bitmask_parallel_batch_size: promises.append(self._async_submit_fill_bitmask(batch)) @@ -288,76 +382,47 @@ def grammar_bitmask( grammar = structured_output_request.grammar if TYPE_CHECKING: assert isinstance(grammar, StructuredOutputGrammar) - apply_bitmask = self.should_fill_bitmask(request) - - reasoner = None if apply_bitmask else self._get_reasoner(request) + req_tokens = scheduled_spec_decode_tokens.get(req_id, list()) + constraint_start = self._get_constraint_start( + request, strip_speculative_padding(req_tokens) + ) state_advancements = 0 - post_reasoning_end_in_window = False - req_tokens = scheduled_spec_decode_tokens.get(req_id, ()) - # Locate the reasoning-end marker in the draft window once. - # Rejected-draft padding (-1) is trailing, so the window is - # cut there and the reasoner never sees a placeholder. - window_start = len(request.all_token_ids) - reasoning_end_index: int | None = None - if reasoner is not None and req_tokens: - window = req_tokens - if -1 in window: - window = window[: window.index(-1)] - reasoning_end_index = self._find_reasoning_end_index( - reasoner, - request.all_token_ids, - window_start, - window, - delta_appended=False, - ) + seen_padding = False + failed = False for i, token in enumerate(req_tokens): + apply_bitmask = ( + not failed and not seen_padding and i >= constraint_start + ) self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),)) - advance_grammar = apply_bitmask if token == -1: - apply_bitmask = False - advance_grammar = False - elif not apply_bitmask and window_start + i == reasoning_end_index: - # Reasoning ended mid-window. Constrain the rest - # of the window via bitmask. Skip grammar advance - # through the marker (it is reasoning content); - # try to advance through subsequent drafts so the - # next bitmask row reflects the post-advance state, - # but tolerate rejection since those drafts predate - # the bitmask and are not guaranteed valid. - apply_bitmask = True - advance_grammar = False - post_reasoning_end_in_window = True - if advance_grammar and not grammar.is_terminated(): - if post_reasoning_end_in_window: - accepted = bool(grammar.validate_tokens([token])) - if accepted: - accepted = grammar.accept_tokens(req_id, [token]) - else: - accepted = grammar.accept_tokens(req_id, [token]) - if accepted: + seen_padding = True + elif apply_bitmask: + if not grammar.is_terminated() and grammar.accept_tokens( + req_id, [token] + ): state_advancements += 1 - elif not post_reasoning_end_in_window: - raise AssertionError( - (token, req_id, scheduled_spec_decode_tokens) + else: + failed = True + logger.error( + "Unexpected: grammar terminated or rejected draft " + "token %s for request %s during bitmask fill.", + token, + req_id, ) cumulative_index += 1 + # Diffusion LLMs don't sample a bonus token after the # scheduled positions, so skip its bitmask in that case. if not (self.vllm_config.model_config.is_diffusion and req_tokens): - # bonus_apply must be True when the bonus-row position - # should be grammar-constrained. Two triggers: - # - should_fill_bitmask(request): reasoning was already - # over at step start (or no reasoner / - # enable_in_reasoning). - # - apply_bitmask: reasoning ended mid-window in this - # call and was flipped True after the marker; - # should_fill_bitmask still returns False here because - # reasoning_ended is only persisted later by - # should_advance. - bonus_apply = self.should_fill_bitmask(request) or apply_bitmask + bonus_apply = ( + not failed + and not seen_padding + and constraint_start <= len(req_tokens) + ) self._fill_bitmasks(((grammar, cumulative_index, bonus_apply),)) cumulative_index += 1 + if state_advancements > 0: grammar.rollback(state_advancements) @@ -370,161 +435,29 @@ def grammar_bitmask( # and deserialization when sending this to the GPU workers. return bitmask_tensor.numpy() - def should_fill_bitmask(self, request: "Request") -> bool: - # NOTE (Hanchen) if enable_in_reasoning is True, it means that - # the model needs to be constrained in reasoning. So we should always - # enable the bitmask filling. - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): + def accept_tokens(self, request: "Request", new_token_ids: list[int]) -> bool: + """Advance grammar with accepted tokens. Returns False on rejection.""" + if not request.use_structured_output or not new_token_ids: return True - reasoner = self._get_reasoner(request) - if reasoner is not None: + structured_req = request.structured_output_request + if TYPE_CHECKING: assert structured_req is not None - if structured_req.reasoning_ended is None: - # This should be removed here, but since `openai_gptoss` - # is an independent code path, it is kept for now. - # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, - # it can be removed. - structured_req.reasoning_ended = reasoner.is_reasoning_end( - request.prompt_token_ids or [] - ) - return structured_req.reasoning_ended - return True - - def should_advance( - self, - request: "Request", - new_token_ids: list[int] | None = None, - ) -> bool: - if not request.use_structured_output: - return False - - # To determine whether we can advance the FSM. - # Supports thinking usage where we skip the reasoning components. + grammar = structured_req.grammar if TYPE_CHECKING: - assert request.structured_output_request is not None - assert request.structured_output_request.grammar is not None - structured_req = request.structured_output_request - if self.enable_in_reasoning or ( - structured_req is not None and structured_req.reasoning_ended is True - ): - return True - - # by default, we should always advance - # for cases that don't use thinking mode. - reasoner = self._get_reasoner(request) - if reasoner is None: - return True + assert isinstance(grammar, StructuredOutputGrammar) - assert structured_req is not None - - # Check if reasoning ends in *this* step. - # When the caller passes new_token_ids (the tokens that were just - # appended this step), use it directly as the delta window. The - # placeholder-derived fallback assumes num_output_placeholders == - # len(new_token_ids), which breaks under async scheduling + spec - # decode when some drafts are rejected (#43388): the placeholder - # count remains > 0 after the step and the computed delta window - # starts past the reasoning-end marker. - all_token_ids = request.all_token_ids - if new_token_ids: - # The tokens were already appended this step, so the step window - # starts exactly len(new_token_ids) from the end. - start = len(all_token_ids) - len(new_token_ids) - delta_ids: Sequence[int] = new_token_ids - else: - # Draft-validation callers arrive here after update_from_output - # has already persisted any boundary found this step, so this - # window only needs to cover tokens not examined yet. - delta_from = request.num_computed_tokens - request.num_output_placeholders - start = ( - delta_from - if delta_from >= 0 - else max(len(all_token_ids) + delta_from, 0) - ) - delta_ids = all_token_ids[start:] - end_index = self._find_reasoning_end_index( - reasoner, all_token_ids, start, delta_ids, delta_appended=True + constraint_start = self._get_constraint_start( + request, new_token_ids, spec_tokens_committed=True ) - if end_index is not None: - structured_req.reasoning_ended = True - # Record the boundary so the scheduler can exclude reasoning tokens. - structured_req.reasoning_end_token_index = end_index + # Early return only when the constraint hasn't started. + # Otherwise, latch structured_req.reasoning_ended. + if constraint_start > len(new_token_ids): return True - - return False - - @staticmethod - def _find_reasoning_end_index( - reasoner: "ReasoningParser", - all_token_ids: Sequence[int], - start: int, - delta_ids: Sequence[int], - *, - delta_appended: bool, - ) -> int | None: - """Locates the token that ends reasoning within ``delta_ids``. - - ``delta_ids`` occupy positions ``start..``; ``delta_appended`` says - whether they are already in ``all_token_ids`` (accepted tokens) or - not (speculative drafts). - - Legacy parsers are probed one token at a time. Accepted tokens get a - whole-delta check as well, and a marker seen only that way reports - the final token. Drafts skip it: some predicates are order-sensitive - over a multi-token window and would miss the marker. - - Returns: - Absolute index of the last reasoning token, or ``None``. - """ - if ( - isinstance(reasoner, ParserEngineReasoningAdapter) - and reasoner.reasoning_end_token_ids - ): - offset = reasoner.find_reasoning_end_offset(delta_ids) - return None if offset is None else start + offset - - if delta_appended and not reasoner.is_reasoning_end_streaming( - all_token_ids, delta_ids - ): - return None - prefix = list(itertools.islice(all_token_ids, start)) - for offset, token in enumerate(delta_ids): - prefix.append(token) - if reasoner.is_reasoning_end_streaming(prefix, [token]): - return start + offset - if delta_appended: - return len(all_token_ids) - 1 - return None - - def trim_reasoning_for_advance( - self, request: "Request", new_token_ids: list[int] - ) -> list[int]: - """Drops reasoning content from tokens about to advance the grammar. - - When reasoning ends mid-step (see should_advance), the step's output - still contains reasoning tokens up to and including the end marker. - Those are not grammar content: feeding them to accept_tokens makes - the grammar reject the marker and kills the request (#44006). - - Returns: - The suffix of ``new_token_ids`` that follows the reasoning-end - marker. Steps fully after the boundary are returned unchanged. - """ - structured_req = request.structured_output_request - if structured_req is None: - return new_token_ids - end_idx = structured_req.reasoning_end_token_index - if end_idx is None: - return new_token_ids - first_idx = len(request.all_token_ids) - len(new_token_ids) - num_reasoning = end_idx + 1 - first_idx - if num_reasoning <= 0: - return new_token_ids - return new_token_ids[num_reasoning:] + structured_req.reasoning_ended = True + return grammar.accept_tokens( + request.request_id, new_token_ids[constraint_start:] + ) def clear_backend(self) -> None: if self.backend is not None: diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 3809ba2d91b0..1ba7e2734d9e 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -236,6 +236,12 @@ def reset(self): } +def _has_pattern_and_length_bounds(schema: dict[str, Any]) -> bool: + return ("pattern" in schema or "format" in schema) and ( + "minLength" in schema or "maxLength" in schema + ) + + def has_xgrammar_unsupported_json_features(schema: dict[str, Any]) -> bool: """Check if JSON schema contains features unsupported by xgrammar.""" @@ -269,16 +275,39 @@ def check_object(obj: dict[str, Any]) -> bool: # the compiled EBNF: pattern/format grammars come out byte-identical # with and without the length keywords, while maxLength alone lowers # to {0, N} correctly. + if obj.get("type") == "string" and _has_pattern_and_length_bounds(obj): + return True + + # propertyNames validates names, so it is a string schema even when it + # omits "type", which is the form that escapes the check above. if ( - obj.get("type") == "string" - and ("pattern" in obj or "format" in obj) - and ("minLength" in obj or "maxLength" in obj) + obj.get("type") == "object" + and isinstance(obj.get("propertyNames"), dict) + and _has_pattern_and_length_bounds(obj["propertyNames"]) ): return True - # Unsupported keywords for objects - if obj.get("type") == "object" and any( - key in obj for key in ("patternProperties", "propertyNames") + # FIXME: propertyNames conflicts with properties/patternProperties/ + # additionalProperties/unevaluatedProperties under xgrammar. + # https://github.com/mlc-ai/xgrammar/issues/826 + if ( + obj.get("type") == "object" + and "propertyNames" in obj + and ( + "properties" in obj + or "patternProperties" in obj + or isinstance(obj.get("additionalProperties"), dict) + or obj.get("unevaluatedProperties", True) is not True + ) + ): + return True + + # FIXME: multiple patternProperties, or patternProperties alongside + # properties, conflict under xgrammar. + if ( + obj.get("type") == "object" + and isinstance(obj.get("patternProperties"), dict) + and ("properties" in obj or len(obj["patternProperties"]) > 1) ): return True diff --git a/vllm/v1/structured_output/request.py b/vllm/v1/structured_output/request.py index b0f4f3c3a8fd..dce0038782ce 100644 --- a/vllm/v1/structured_output/request.py +++ b/vllm/v1/structured_output/request.py @@ -25,12 +25,6 @@ class StructuredOutputRequest: Future[StructuredOutputGrammar] | StructuredOutputGrammar | Exception | None ) = None reasoning_ended: bool | None = None - # Absolute index into the request's all_token_ids of the last reasoning - # token (the reasoning-end marker). Tokens at or before this index are - # reasoning content and must never be fed to the grammar. Only set when - # reasoning ends in a step whose tokens the scheduler advances immediately - # (structural tags + speculative decoding, see #42452). - reasoning_end_token_index: int | None = None reasoning_parser_kwargs: dict[str, Any] | None = None # Cached per request; do not share reasoning parsers across requests because # their behavior can depend on reasoning_parser_kwargs. diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b0411dc9e8b1..facafd28608f 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -137,9 +137,10 @@ def __init__( self.np = self.cpu.numpy() def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: - if n is None: - return self.gpu.copy_(self.cpu, non_blocking=True) - return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True) + cpu, gpu = self.cpu, self.gpu + if n is not None: + cpu, gpu = cpu[:n], gpu[:n] + return gpu.copy_(cpu.pin_memory() if PIN_MEMORY else cpu, non_blocking=True) def copy_to_cpu(self, n: int | None = None) -> torch.Tensor: """NOTE: Because this method is non-blocking, explicit synchronization @@ -209,6 +210,15 @@ def __init__( self.processes: list[BaseProcess] = [] self._address_pipes: list[connection.Connection] = [] + admission_counters = None + if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: + from vllm.v1.engine.admission_control import SharedAdmissionStats + + admission_counters = spawn_context.RawArray( + "q", + SharedAdmissionStats.num_counters(num_servers), + ) + for i, in_addr, out_addr in zip( range(num_servers), input_addresses, output_addresses ): @@ -218,6 +228,8 @@ def __init__( "client_count": num_servers, "client_index": i, } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters if stats_update_address is not None: client_config["stats_update_address"] = stats_update_address if tensor_queue is not None: diff --git a/vllm/v1/worker/cpu/buffer_utils.py b/vllm/v1/worker/cpu/buffer_utils.py index 8ee257a28f40..cdc76ab8b47f 100644 --- a/vllm/v1/worker/cpu/buffer_utils.py +++ b/vllm/v1/worker/cpu/buffer_utils.py @@ -4,13 +4,11 @@ import torch -from vllm.utils.platform_utils import is_uva_available - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): - if not is_uva_available(): - raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu") self.np = self.cpu.numpy() - self.uva = self.cpu + + def uva(self, n: int | None = None): + return self.cpu[:n] if n is not None else self.cpu diff --git a/vllm/v1/worker/cpu/shm.py b/vllm/v1/worker/cpu/shm.py index f92c6c3bf28d..b615f4b5b5c8 100644 --- a/vllm/v1/worker/cpu/shm.py +++ b/vllm/v1/worker/cpu/shm.py @@ -76,14 +76,17 @@ def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]: def async_tensor_h2d( data: list | np.ndarray | torch.Tensor, - device: str | torch.device, + device: str | torch.device | None = None, dtype: torch.dtype | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: if isinstance(data, np.ndarray): data = torch.from_numpy(data) - if isinstance(data, torch.Tensor): + if not isinstance(data, torch.Tensor): + data = torch.tensor(data, dtype=dtype, device="cpu") + elif out is None: return data.to(dtype=dtype) - return torch.tensor(data, dtype=dtype, device="cpu") + return data if out is None else out.copy_(data) torch_utils.async_tensor_h2d = async_tensor_h2d diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 2c1dac2d0710..0731072e5968 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -6,6 +6,7 @@ import numpy as np import torch +from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.utils.platform_utils import is_uva_available from vllm.utils.torch_utils import ( @@ -13,6 +14,8 @@ get_accelerator_view_from_cpu_tensor, ) +logger = init_logger(__name__) + # Default round-robin depth for the UVA buffer pools. Must be >= the number of # concurrent in-flight steps (engine batch_queue_size). _DEFAULT_MAX_CONCURRENCY = 2 @@ -23,31 +26,36 @@ def set_default_max_concurrency(n: int) -> None: _DEFAULT_MAX_CONCURRENCY = max(2, n) -def async_copy_to_gpu( - x: torch.Tensor | np.ndarray, - out: torch.Tensor | None = None, - device: torch.device | None = None, -) -> torch.Tensor: - if isinstance(x, np.ndarray): - x = torch.from_numpy(x) - assert x.is_cpu - - if out is None: - assert device is not None - out = torch.empty_like(x, device=device) - - # pin_memory() is no-op if the memory is already pinned. - pinned = x.pin_memory() - return out.copy_(pinned, non_blocking=True) - - class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): if not is_uva_available(): raise RuntimeError("UVA is not available") self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True) self.np = self.cpu.numpy() - self.uva = get_accelerator_view_from_cpu_tensor(self.cpu) + self._uva = get_accelerator_view_from_cpu_tensor(self.cpu) + + def uva(self, n: int | None = None) -> torch.Tensor: + return self._uva[:n] if n is not None else self._uva + + +class NonUvaBuffer: + """Explicit-copy fallback for platforms without pinned memory.""" + + def __init__(self, size: int | Sequence[int], dtype: torch.dtype): + from vllm.platforms import current_platform + + logger.warning_once( + "Pinned memory is not available on this platform; falling back " + "to device memory for UVA buffers." + ) + self.cpu = torch.zeros(size, dtype=dtype, device="cpu") + self.np = self.cpu.numpy() + self._uva = torch.zeros(size, dtype=dtype, device=current_platform.device_type) + + def uva(self, n: int | None = None) -> torch.Tensor: + if n is None: + return self._uva.copy_(self.cpu) + return self._uva[:n].copy_(self.cpu[:n]) class UvaBufferPool: @@ -69,7 +77,8 @@ def __init__( self.max_concurrency = max_concurrency # UVA buffers for concurrency - self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)] + self._buffer_cls = UvaBuffer if is_uva_available() else NonUvaBuffer + self._uva_bufs = [self._buffer_cls(size, dtype) for _ in range(max_concurrency)] # Current buffer index self._curr = 0 @@ -80,12 +89,12 @@ def copy_to_uva(self, x: torch.Tensor | np.ndarray | list) -> torch.Tensor: n = len(x) if n > buf.cpu.shape[0]: capacity = 1 << (n - 1).bit_length() - buf = UvaBuffer((capacity, *buf.cpu.shape[1:]), self.dtype) + buf = self._buffer_cls((capacity, *buf.cpu.shape[1:]), self.dtype) self._uva_bufs[self._curr] = buf # CPU-to-CPU copy dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np dst[:n] = x - return buf.uva[:n] + return buf.uva(n) def copy_to_gpu( self, @@ -141,6 +150,9 @@ def __init__( self.device = device self.max_concurrency = max_concurrency + # Without UVA, large-but-cold tensors live on the GPU directly. + uva_instead_of_gpu = uva_instead_of_gpu and is_uva_available() + if not uva_instead_of_gpu: # Create a GPU tensor (default) self.gpu = torch.zeros(size, dtype=dtype, device=device) @@ -148,7 +160,7 @@ def __init__( # For a large but not-frequently-accessed tensor, we can use UVA instead of # GPU to save GPU memory self._uva_buf = UvaBuffer(size, dtype) - self.gpu = self._uva_buf.uva + self.gpu = self._uva_buf.uva() self._staged_write_indices: list[int] = [] self._staged_write_starts: list[int] = [] diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d197dbbc6df1..9657c28fe601 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -777,11 +777,12 @@ def prepare_inputs_to_capture( input_batch = InputBatch.make_dummy( num_reqs, num_tokens, input_buffers, max_query_len=max_query_len ) - input_block_tables = block_tables.get_dummy_block_tables(num_reqs) - slot_mapping_provider: BlockTables | PCPManager = block_tables if pcp_manager is not None: - slot_mapping_provider = pcp_manager - slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens) + input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) + + block_table_provider = pcp_manager or block_tables + input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) + slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens) slot_mappings_by_layer = build_slot_mappings_by_layer( slot_mappings, kv_cache_config ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 5cbd642d2529..d3c65bd06f11 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING @@ -217,9 +218,12 @@ def set_dummy_context( context_len: int, num_kv_blocks: int, max_model_len: int, + input_block_tables: Sequence[torch.Tensor] | None = None, ) -> None: """Give each dummy request context_len of context, used when profiling step cost.""" - if not block_tables.input_block_tables: + if input_block_tables is None: + input_block_tables = block_tables.input_block_tables + if not input_block_tables: # Attention-free models have no KV context to fabricate. return num_reqs = input_batch.num_reqs @@ -241,7 +245,7 @@ def set_dummy_context( seq_len = context_len + query_len for block_table, block_size, bpk in zip( - block_tables.input_block_tables, + input_block_tables, block_tables.kernel_block_sizes, block_tables.blocks_per_kv_block, ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 761893303a5c..82cc697e731c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -64,7 +64,7 @@ from vllm.tasks import SupportedTask from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib -from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE +from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -99,10 +99,7 @@ init_kv_cache, ) from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import ( - async_copy_to_gpu, - set_default_max_concurrency, -) +from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, @@ -299,15 +296,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None - # Multi-module MTP feeds its modules the next num_speculative_steps prefill - # tokens during chunked prefill. Other speculators only read the immediate - # next one. - num_prefill_lookahead = ( - self.num_speculative_steps - if self.speculative_config is not None - and self.speculative_config.use_multi_module_mtp() - else 1 - ) + num_prefill_lookahead = max(1, vllm_config.num_prefill_lookahead_tokens) self.step_timing = StepTimingCollector() @@ -1289,7 +1278,7 @@ def prepare_inputs( req_ids = batch_req_state.req_ids num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens idx_mapping_np = batch_req_state.idx_mapping_np - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) num_reqs = len(req_ids) # Get the number of draft tokens for each request. @@ -1324,7 +1313,7 @@ def prepare_inputs( cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) cu_num_logits_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_np[1:]) - cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) + cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) adaptive_verification = ( self.adaptive_verification if num_draft_tokens_per_req is not None else None @@ -1352,7 +1341,7 @@ def prepare_inputs( # Some attention backends like FA3 require query_start_loc to be non-decreasing. query_start_loc_np[num_reqs + 1 :] = num_tokens query_start_loc = self.input_buffers.query_start_loc - async_copy_to_gpu(query_start_loc_np, out=query_start_loc) + async_tensor_h2d(query_start_loc_np, out=query_start_loc) if adaptive_verification is not None: cu_num_logits, query_start_loc, total_num_draft_tokens = ( adaptive_verification.reallocate_drafts(req_ids, idx_mapping) @@ -1468,7 +1457,7 @@ def prepare_inputs( input_batch = pcp.maybe_partition_pcp_batch( self.pcp_manager, input_batch, - padded_num_tokens=batch_desc.num_tokens, + batch_desc, ) return input_batch @@ -1496,7 +1485,8 @@ def prepare_attn( def prepare_dummy_attn( self, input_batch: InputBatch, valid_state_slots: bool = False ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: - block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) + block_table_provider = self.pcp_manager or self.block_tables + block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs) if valid_state_slots: state_slots = torch.arange( 1, @@ -1745,25 +1735,29 @@ def execute_model( self.input_buffers, max_query_len=batch_desc.max_query_len, ) - if not skip_attn_for_dummy_run: - block_tables, slot_mappings = self.prepare_dummy_attn( - input_batch, valid_dummy_state_slots - ) - if context_len: - set_dummy_context( - input_batch, - self.block_tables, - context_len, - self.kv_cache_config.num_blocks, - self.max_model_len, - ) - else: + if self.pcp_manager is not None: + input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) + if skip_attn_for_dummy_run: assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( "Attention metadata must be prepared for dummy runs when using " "FULL cudagraph mode." ) block_tables = None slot_mappings = None + else: + block_tables, slot_mappings = self.prepare_dummy_attn( + input_batch, valid_dummy_state_slots + ) + if not skip_attn_for_dummy_run and context_len: + assert block_tables is not None + set_dummy_context( + input_batch, + self.block_tables, + context_len, + self.kv_cache_config.num_blocks, + self.max_model_len, + input_block_tables=block_tables, + ) attn_metadata = None slot_mappings_by_layer = None diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 3a51532f57a1..376e0f251012 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterator from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import numpy as np import torch @@ -9,14 +10,17 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens from vllm.v1.worker.gpu.block_table import BlockTables -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import ( InputBatch, InputBuffers, ) +if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + logger = init_logger(__name__) @@ -169,8 +173,14 @@ def validate_config( "MRV2 sparse MLA PCP does not support CUDA graphs yet. " "Set -cc.cudagraph_mode=NONE." ) - if cudagraph_mode.has_full_cudagraphs(): - raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + if ( + cudagraph_mode.has_full_cudagraphs() + and not cudagraph_mode.separate_routine() + ): + raise NotImplementedError( + "MRV2 PCP supports full CUDA graphs for decode-only routines. " + "Use FULL_DECODE_ONLY, FULL_AND_PIECEWISE, PIECEWISE, or NONE." + ) if ( parallel_config.decode_context_parallel_size > 1 and parallel_config.dcp_comm_backend != "ag_rs" @@ -348,13 +358,13 @@ def _build_batch_layout( dtype=np.int64, ) - self._hidden_restore_idx = async_copy_to_gpu( + self._hidden_restore_idx = async_tensor_h2d( hidden_restore_idx, device=self.device ) - self._padded_gather_idx = async_copy_to_gpu( + self._padded_gather_idx = async_tensor_h2d( padded_gather_idx, device=self.device ) - self._gathered_kv_write_mask = async_copy_to_gpu( + self._gathered_kv_write_mask = async_tensor_h2d( gathered_kv_write_mask, device=self.device ) return segments_by_rank, per_rank_num_tokens @@ -375,6 +385,22 @@ def get_num_tokens_for_dispatch( for rank in range(self.pcp_world_size) ) + @staticmethod + def _resolve_num_reqs_after_padding( + input_batch: InputBatch, + padded_num_reqs: int | None, + num_local_reqs: int, + ) -> int: + if padded_num_reqs is None: + return num_local_reqs + if input_batch.has_prefill: + raise RuntimeError("PCP FULL graphs require a decode-only batch.") + assert padded_num_reqs >= num_local_reqs, ( + "PCP graph request capacity must cover the rank-local batch: " + f"{padded_num_reqs} < {num_local_reqs}." + ) + return padded_num_reqs + @property def input_buffers(self) -> InputBuffers: assert self._input_buffers is not None @@ -384,6 +410,7 @@ def partition_batch( self, input_batch: InputBatch, padded_num_tokens: int | None = None, + padded_num_reqs: int | None = None, ) -> InputBatch: assert self._input_buffers is not None input_buffers = self._input_buffers @@ -414,10 +441,15 @@ def partition_batch( ] num_local_reqs = len(local_segments) - if num_local_reqs > input_buffers.max_num_reqs: + num_reqs_after_padding = self._resolve_num_reqs_after_padding( + global_batch, + padded_num_reqs, + num_local_reqs, + ) + if num_reqs_after_padding > input_buffers.max_num_reqs: raise RuntimeError( - "PCP local request count exceeds the MRV2 input buffer size: " - f"{num_local_reqs} > {input_buffers.max_num_reqs}." + "PCP padded local request count exceeds the MRV2 input buffer size: " + f"{num_reqs_after_padding} > {input_buffers.max_num_reqs}." ) local_to_global_batch_req_idx_np = np.fromiter( @@ -504,22 +536,26 @@ def partition_batch( local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1] np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens - async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) - local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1] + async_tensor_h2d(local_query_start_loc_np, out=input_buffers.query_start_loc) + local_query_start_loc = input_buffers.query_start_loc[ + : num_reqs_after_padding + 1 + ] - local_to_global_req_idx = async_copy_to_gpu( + local_to_global_req_idx = async_tensor_h2d( local_to_global_req_idx_np, device=self.device ) - seq_lens = input_buffers.seq_lens[:num_local_reqs] + seq_lens = input_buffers.seq_lens[:num_reqs_after_padding] + real_seq_lens = seq_lens[:num_local_reqs] if num_local_tokens > 0: local_end_positions = torch.index_select( input_buffers.positions, 0, - local_query_start_loc[1:] - 1, + local_query_start_loc[1 : num_local_reqs + 1] - 1, ) - seq_lens.copy_(local_end_positions + 1) + real_seq_lens.copy_(local_end_positions + 1) else: - seq_lens.zero_() + real_seq_lens.zero_() + seq_lens[num_local_reqs:].zero_() is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) @@ -552,11 +588,15 @@ def partition_batch( local_num_computed_prefill_tokens_np = np.minimum( local_start_pos_np, local_prefill_len_np ) - local_is_prefilling_np = ( + real_local_is_prefilling_np = ( local_num_computed_prefill_tokens_np < local_prefill_len_np ) - seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) - seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens + local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_) + local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32) + seq_lens_cpu_upper_bound_np[:num_local_reqs] = ( + local_start_pos_np + local_num_scheduled_tokens + ) dcp_local_seq_lens_cpu_upper_bound = None if self.dcp_world_size > 1: # The largest DCP shard of each row's whole request, identical on @@ -564,18 +604,24 @@ def partition_batch( request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[ local_to_global_batch_req_idx_np ] - dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( + real_dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( torch.from_numpy(request_seq_lens.astype(np.int32)), self.dcp_world_size, 0, self.cp_interleave, ) + dcp_local_seq_lens_cpu_upper_bound = torch.zeros( + num_reqs_after_padding, dtype=torch.int32 + ) + dcp_local_seq_lens_cpu_upper_bound[:num_local_reqs].copy_( + real_dcp_local_seq_lens_cpu_upper_bound + ) self._local_batch = replace( input_batch, req_ids=local_req_ids, num_reqs=num_local_reqs, - num_reqs_after_padding=num_local_reqs, + num_reqs_after_padding=num_reqs_after_padding, idx_mapping=local_to_global_req_idx, idx_mapping_np=local_to_global_req_idx_np, expanded_idx_mapping=local_to_global_req_idx, @@ -588,7 +634,7 @@ def partition_batch( num_draft_tokens=0, num_draft_tokens_per_req=None, query_start_loc=local_query_start_loc, - query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1], + query_start_loc_np=local_query_start_loc_np[: num_reqs_after_padding + 1], seq_lens=seq_lens, seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np), dcp_local_seq_lens=None, @@ -608,6 +654,31 @@ def partition_batch( ) return self._local_batch + def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch: + """Stage a capture or dummy batch in persistent PCP input buffers.""" + input_buffers = self.input_buffers + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + input_batch = replace( + input_batch, + input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids), + positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions), + is_padding=input_buffers.is_padding[:num_tokens].copy_( + input_batch.is_padding + ), + query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_( + input_batch.query_start_loc + ), + seq_lens=input_buffers.seq_lens[:num_reqs].copy_(input_batch.seq_lens), + ) + return input_batch + + def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: + assert self._local_block_tables is not None + return tuple( + block_table[:num_reqs].zero_() for block_table in self._local_block_tables + ) + def prepare_attn( self, input_batch: InputBatch ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: @@ -728,13 +799,22 @@ def restore_for_sampling( def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, - padded_num_tokens: int | None = None, + batch_desc: "BatchExecutionDescriptor", ) -> InputBatch: if manager is None: return input_batch + + padded_num_tokens = None + padded_num_reqs = None + if batch_desc.cg_mode != CUDAGraphMode.NONE: + padded_num_tokens = batch_desc.num_tokens + if batch_desc.cg_mode == CUDAGraphMode.FULL: + padded_num_reqs = batch_desc.num_reqs + return manager.partition_batch( input_batch, padded_num_tokens=padded_num_tokens, + padded_num_reqs=padded_num_reqs, ) diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 8c65a23ef7ec..ab1b991d1a8d 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -11,7 +11,7 @@ from vllm.distributed.parallel_state import get_pp_group from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.worker.gpu.input_batch import InputBatch @@ -138,7 +138,7 @@ def get_prev_sampled_outputs( return None # Filter excluded request indices. idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np) - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) self.main_stream.wait_event(slot.event) if slot.draft_tokens is not None and draft_tokens_to_update is not None: @@ -148,7 +148,7 @@ def get_prev_sampled_outputs( keep = ~exclude_mask keep_t = torch.as_tensor(keep, device=self.device) draft_tokens = draft_tokens[keep_t] - draft_idx_mapping = async_copy_to_gpu( + draft_idx_mapping = async_tensor_h2d( slot.idx_mapping_np[keep], device=self.device ) draft_tokens_to_update[draft_idx_mapping] = draft_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index af58b8993bfa..245f304ac4e3 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py @@ -14,6 +14,7 @@ from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream @@ -21,7 +22,6 @@ get_attn_cg_support, get_query_lens_mismatch_unsupported_backend, ) -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) _PROFILE_REPLAYS = 5 @@ -413,7 +413,7 @@ def reallocate_drafts( if draft_budget == 0: capacities.zero_() else: - async_copy_to_gpu(scheduled_drafts, out=capacities) + async_tensor_h2d(scheduled_drafts, out=capacities) if draft_budget < int(scheduled_drafts.sum()): _assign_draft_token_budget_compiled( self._confidence_probs, @@ -424,10 +424,7 @@ def reallocate_drafts( ) num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs] - async_copy_to_gpu( - num_non_draft_tokens, - out=num_non_draft_tokens_gpu, - ) + async_tensor_h2d(num_non_draft_tokens, out=num_non_draft_tokens_gpu) self._cu_num_logits[:1].zero_() torch.cumsum( capacities + self.num_bonus_tokens, diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index db1e17a9f23b..bf1a5761b2ef 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -3,7 +3,7 @@ import torch.nn as nn -from vllm.config import ModelConfig, VllmConfig, replace +from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace from vllm.logger import init_logger from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config @@ -31,6 +31,29 @@ def _resolve_dspark_attention_backend( return None +def _get_dspark_parallel_config( + parallel_config: ParallelConfig, + tensor_parallel_size: int, +) -> ParallelConfig: + if parallel_config.enable_eplb: + logger.warning_once( + "EPLB is disabled for the DSpark draft model. EPLB remains enabled " + "for the target model." + ) + + return replace( + parallel_config, + pipeline_parallel_size=1, + tensor_parallel_size=tensor_parallel_size, + enable_eplb=False, + eplb_config=replace( + parallel_config.eplb_config, + num_redundant_experts=0, + ), + enable_elastic_ep=False, + ) + + def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: speculative_config = vllm_config.speculative_config assert speculative_config is not None @@ -54,12 +77,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo draft_vllm_config = replace( vllm_config, - parallel_config=replace( + parallel_config=_get_dspark_parallel_config( vllm_config.parallel_config, - pipeline_parallel_size=1, - tensor_parallel_size=( - speculative_config.draft_parallel_config.tensor_parallel_size - ), + speculative_config.draft_parallel_config.tensor_parallel_size, ), attention_config=replace( vllm_config.attention_config, diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 8437b0aed9a4..057d6f7154c7 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -5,8 +5,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import PIN_MEMORY -from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.worker.gpu.input_batch import InputBatch @@ -68,7 +67,7 @@ def apply_grammar_bitmask( # Asynchronously copy the bitmask to GPU. with torch.cuda.stream(self.copy_stream): - bitmask = async_copy_to_gpu( + bitmask = async_tensor_h2d( grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]] ) diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 27e850b66a67..5e7f1f70c56d 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/vllm/v1/worker/lora_model_runner_mixin.py @@ -6,7 +6,6 @@ from collections.abc import Callable from contextlib import contextmanager -from typing import TypeAlias import numpy as np import torch @@ -19,10 +18,7 @@ from vllm.lora.request import LoRARequest from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.model_executor.models import supports_lora -from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch -from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch - -InputBatch: TypeAlias = TPUInputBatch | GPUInputBatch +from vllm.v1.worker.gpu_input_batch import InputBatch logger = init_logger(__name__) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 724c714057f7..1deb89bda65c 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -26,8 +26,7 @@ UniformTypeKVCacheSpecs, ) from vllm.v1.utils import CpuGpuBuffer -from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch logger = init_logger(__name__) @@ -1436,7 +1435,7 @@ def preprocess_mamba( kv_cache_config: KVCacheConfig, cache_config: CacheConfig, mamba_state_idx: dict[str, int], - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, @@ -1541,7 +1540,7 @@ def preprocess_mamba( def postprocess_mamba_all( scheduler_output: SchedulerOutput, kv_cache_config: KVCacheConfig, - input_batch: GPUInputBatch, + input_batch: InputBatch, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], num_spec_tokens: int, @@ -1572,7 +1571,7 @@ def postprocess_mamba_all( def preprocess_mamba_all_specdec( scheduler_output: SchedulerOutput, - input_batch: GPUInputBatch, + input_batch: InputBatch, mamba_state_idx: dict[str, int], num_reqs: int, prev_last_scheduled_idx_buf: CpuGpuBuffer, @@ -1591,7 +1590,7 @@ def postprocess_mamba_align_gpu( num_reqs: int, num_accepted_tokens_gpu: torch.Tensor, num_accepted_tokens_cpu_tensor: torch.Tensor, - input_batch: GPUInputBatch, + input_batch: InputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, diff --git a/vllm/v1/worker/tpu_input_batch.py b/vllm/v1/worker/tpu_input_batch.py deleted file mode 100644 index c16156a5e3ef..000000000000 --- a/vllm/v1/worker/tpu_input_batch.py +++ /dev/null @@ -1,559 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Datastructures defining a TPU input batch - -from typing import cast - -import numpy as np -import torch - -from vllm.lora.request import LoRARequest -from vllm.sampling_params import SamplingType -from vllm.utils import length_from_prompt_token_ids_or_embeds -from vllm.utils.collection_utils import swap_dict_values -from vllm.v1.outputs import LogprobsTensors -from vllm.v1.worker.block_table import MultiGroupBlockTable -from vllm.v1.worker.gpu_input_batch import CachedRequestState - -_SAMPLING_EPS = 1e-5 - - -class InputBatch: - def __init__( - self, - max_num_reqs: int, - max_model_len: int, - max_num_batched_tokens: int, - device: torch.device, - pin_memory: bool, - vocab_size: int, - block_sizes: list[int], # The block_size of each kv cache group - kernel_block_sizes: list[int], - max_num_blocks_per_req: list[int], - ): - self.max_num_reqs = max_num_reqs - self.max_model_len = max_model_len - self.max_num_batched_tokens = max_num_batched_tokens - self.device = device - self.pin_memory = pin_memory - self.vocab_size = vocab_size - - self._req_ids: list[str | None] = [] - self.req_id_to_index: dict[str, int] = {} - - # TODO(woosuk): This buffer could be too large if max_model_len is big. - # Find a way to reduce the CPU memory usage. - # This buffer is not directly transferred to the GPU, so it does not - # need to be pinned. - self.token_ids_cpu_tensor = torch.zeros( - (max_num_reqs, max_model_len), - device="cpu", - dtype=torch.int32, - pin_memory=False, - ) - self.token_ids_cpu = self.token_ids_cpu_tensor.numpy() - self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32) - self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32) - self.num_computed_tokens_cpu_tensor = torch.zeros( - (max_num_reqs,), - device="cpu", - dtype=torch.int32, - pin_memory=pin_memory, - ) - self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy() - - # Block table. - self.block_table = MultiGroupBlockTable( - max_num_reqs=max_num_reqs, - max_num_batched_tokens=max_num_batched_tokens, - pin_memory=pin_memory, - device=device, - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks=max_num_blocks_per_req, - ) - - # Sampling-related. - self.temperature = torch.empty( - (max_num_reqs,), dtype=torch.float32, device=device - ) - self.temperature_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.temperature_cpu = self.temperature_cpu_tensor.numpy() - self.greedy_reqs: set[str] = set() - self.random_reqs: set[str] = set() - - self.top_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) - self.top_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.top_p_cpu = self.top_p_cpu_tensor.numpy() - self.top_p_reqs: set[str] = set() - - self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device) - self.top_k_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory - ) - self.top_k_cpu = self.top_k_cpu_tensor.numpy() - self.top_k_reqs: set[str] = set() - - self.min_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device) - self.min_p_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory - ) - self.min_p_cpu = self.min_p_cpu_tensor.numpy() - self.min_p_reqs: set[str] = set() - - # Frequency penalty related data structures - self.frequency_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.frequency_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy() - self.frequency_penalties_reqs: set[str] = set() - - # Presence penalty related data structures - self.presence_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.presence_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy() - self.presence_penalties_reqs: set[str] = set() - - # Repetition penalty related data structures - self.repetition_penalties = torch.empty( - (max_num_reqs,), dtype=torch.float, device=device - ) - self.repetition_penalties_cpu_tensor = torch.empty( - (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory - ) - self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy() - self.repetition_penalties_reqs: set[str] = set() - - # req_index -> (min_tokens, stop_token_ids) - self.min_tokens: dict[int, tuple[int, set[int]]] = {} - - # lora related - self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64) - self.lora_id_to_request_ids: dict[int, set[str]] = {} - self.lora_id_to_lora_request: dict[int, LoRARequest] = {} - - # req_index -> generator - # NOTE(woosuk): The indices of the requests that do not have their own - # generator should not be included in the dictionary. - self.generators: dict[int, torch.Generator] = {} - - self.num_logprobs: dict[str, int] = {} - - # To accumulate prompt logprobs tensor chunks across prefill steps. - self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {} - - self.logit_bias: list[dict[int, float] | None] = [None] * max_num_reqs - self.has_allowed_token_ids: set[str] = set() - # NOTE(lufang): In the mask tensor, if the corresponding token allowed, - # the value is False. Since we use masked_fill_ to set -inf. - self.allowed_token_ids_mask: torch.Tensor | None = None - self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None - - # req_index -> bad_words_token_ids - self.bad_words_token_ids: dict[int, list[list[int]]] = {} - - self.req_output_token_ids: list[list[int] | None] = [] - - @property - def req_ids(self) -> list[str]: - # None elements should only be present transiently - # while performing state updates to the batch. - return cast(list[str], self._req_ids) - - def add_request( - self, - request: "CachedRequestState", - req_index: int | None = None, - ) -> None: - if req_index is None: - req_index = self.num_reqs - assert req_index < self.max_num_reqs - - req_id = request.req_id - if req_index == len(self._req_ids): - self._req_ids.append(req_id) - self.req_output_token_ids.append(request.output_token_ids) - else: - self._req_ids[req_index] = req_id - self.req_output_token_ids[req_index] = request.output_token_ids - - self.req_id_to_index[req_id] = req_index - - # Copy the prompt token ids and output token ids. - num_prompt_tokens = length_from_prompt_token_ids_or_embeds( - request.prompt_token_ids, request.prompt_embeds - ) - # TODO: copy prompt_embeds - self.num_prompt_tokens[req_index] = num_prompt_tokens - self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids - start_idx = num_prompt_tokens - end_idx = start_idx + len(request.output_token_ids) - self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids - # Number of tokens without spec decode tokens. - self.num_tokens_no_spec[req_index] = request.num_tokens - - self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens - self.block_table.add_row(request.block_ids, req_index) - - sampling_params = request.sampling_params - assert sampling_params is not None, "pooling requests not supported yet" - if sampling_params.sampling_type == SamplingType.GREEDY: - # Should avoid division by zero later when apply_temperature. - self.temperature_cpu[req_index] = 0.0 - self.greedy_reqs.add(req_id) - else: - self.temperature_cpu[req_index] = sampling_params.temperature - self.random_reqs.add(req_id) - - self.top_p_cpu[req_index] = sampling_params.top_p - if sampling_params.top_p < 1: - self.top_p_reqs.add(req_id) - top_k = sampling_params.top_k - if 0 < top_k < self.vocab_size: - self.top_k_reqs.add(req_id) - else: - top_k = self.vocab_size - self.top_k_cpu[req_index] = top_k - self.min_p_cpu[req_index] = sampling_params.min_p - self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty - if sampling_params.min_p > _SAMPLING_EPS: - self.min_p_reqs.add(req_id) - if sampling_params.frequency_penalty != 0.0: - self.frequency_penalties_reqs.add(req_id) - self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty - if sampling_params.presence_penalty != 0.0: - self.presence_penalties_reqs.add(req_id) - self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty - if sampling_params.repetition_penalty != 1.0: - self.repetition_penalties_reqs.add(req_id) - if sampling_params.min_tokens: - self.min_tokens[req_index] = ( - sampling_params.min_tokens, - sampling_params.all_stop_token_ids, - ) - - # NOTE(woosuk): self.generators should not include the requests that - # do not have their own generator. - if request.generator is not None: - self.generators[req_index] = request.generator - - if sampling_params.logprobs is not None: - self.num_logprobs[req_id] = sampling_params.logprobs - if sampling_params.logit_bias is not None: - self.logit_bias[req_index] = sampling_params.logit_bias - - if sampling_params.allowed_token_ids: - self.has_allowed_token_ids.add(req_id) - if self.allowed_token_ids_mask_cpu_tensor is None: - # Lazy allocation for this tensor, which can be large. - # False means we don't fill with -inf. - self.allowed_token_ids_mask = torch.zeros( - self.max_num_reqs, - self.vocab_size, - dtype=torch.bool, - device=self.device, - ) - self.allowed_token_ids_mask_cpu_tensor = torch.zeros( - self.max_num_reqs, self.vocab_size, dtype=torch.bool, device="cpu" - ) - self.allowed_token_ids_mask_cpu_tensor[req_index] = True - # False means we don't fill with -inf. - self.allowed_token_ids_mask_cpu_tensor[req_index][ - sampling_params.allowed_token_ids - ] = False - - if sampling_params.bad_words_token_ids: - self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids - - # Add request lora ID - if request.lora_request: - lora_id = request.lora_request.lora_int_id - if lora_id not in self.lora_id_to_request_ids: - self.lora_id_to_request_ids[lora_id] = set() - - self.request_lora_mapping[req_index] = lora_id - self.lora_id_to_request_ids[lora_id].add(request.req_id) - self.lora_id_to_lora_request[lora_id] = request.lora_request - else: - # No LoRA - self.request_lora_mapping[req_index] = 0 - - def remove_request(self, req_id: str) -> int | None: - """This method must always be followed by a call to condense().""" - - req_index = self.req_id_to_index.pop(req_id, None) - if req_index is None: - return None - self._req_ids[req_index] = None - self.req_output_token_ids[req_index] = None - - self.greedy_reqs.discard(req_id) - self.random_reqs.discard(req_id) - self.top_p_reqs.discard(req_id) - self.top_k_reqs.discard(req_id) - self.min_p_reqs.discard(req_id) - self.min_tokens.pop(req_index, None) - self.frequency_penalties_reqs.discard(req_id) - self.presence_penalties_reqs.discard(req_id) - self.repetition_penalties_reqs.discard(req_id) - self.generators.pop(req_index, None) - self.num_logprobs.pop(req_id, None) - self.in_progress_prompt_logprobs_cpu.pop(req_id, None) - - # LoRA - lora_id = self.request_lora_mapping[req_index] - if lora_id != 0: - self.lora_id_to_request_ids[lora_id].discard(req_id) - if len(self.lora_id_to_request_ids[lora_id]) == 0: - self.lora_id_to_request_ids.pop(lora_id) - self.lora_id_to_lora_request.pop(lora_id) - self.request_lora_mapping[req_index] = 0 - - self.logit_bias[req_index] = None - self.has_allowed_token_ids.discard(req_id) - if self.allowed_token_ids_mask_cpu_tensor is not None: - # False means we don't fill with -inf. - self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False) - self.bad_words_token_ids.pop(req_index, None) - return req_index - - def swap_states(self, i1: int, i2: int) -> None: - old_id_i1 = self._req_ids[i1] - old_id_i2 = self._req_ids[i2] - self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa - self.req_output_token_ids[i1], self.req_output_token_ids[i2] = ( - self.req_output_token_ids[i2], - self.req_output_token_ids[i1], - ) - assert old_id_i1 is not None and old_id_i2 is not None - self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = ( - self.req_id_to_index[old_id_i2], - self.req_id_to_index[old_id_i1], - ) - self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = ( - self.num_tokens_no_spec[i2], - self.num_tokens_no_spec[i1], - ) - self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = ( - self.num_prompt_tokens[i2], - self.num_prompt_tokens[i1], - ) - self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = ( - self.num_computed_tokens_cpu[i2], - self.num_computed_tokens_cpu[i1], - ) - self.temperature_cpu[i1], self.temperature_cpu[i2] = ( - self.temperature_cpu[i2], - self.temperature_cpu[i1], - ) - self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1] - self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1] - self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = ( - self.frequency_penalties_cpu[i2], - self.frequency_penalties_cpu[i1], - ) - self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = ( - self.presence_penalties_cpu[i2], - self.presence_penalties_cpu[i1], - ) - self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = ( - self.repetition_penalties_cpu[i2], - self.repetition_penalties_cpu[i1], - ) - self.min_p_cpu[i1], self.min_p_cpu[i2] = self.min_p_cpu[i2], self.min_p_cpu[i1] - - # NOTE: the following is unsafe - # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\ - # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...] - # instead, we need to temporarily copy the data for one of the indices - n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2]) - tmp = self.token_ids_cpu[i1, :n].copy() - self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n] - self.token_ids_cpu[i2, :n] = tmp - - swap_dict_values(self.generators, i1, i2) - swap_dict_values(self.min_tokens, i1, i2) - swap_dict_values(self.bad_words_token_ids, i1, i2) - - self.request_lora_mapping[i1], self.request_lora_mapping[i2] = ( - self.request_lora_mapping[i2], - self.request_lora_mapping[i1], - ) - self.logit_bias[i1], self.logit_bias[i2] = ( - self.logit_bias[i2], - self.logit_bias[i1], - ) - - if self.allowed_token_ids_mask_cpu_tensor is not None: - ( - self.allowed_token_ids_mask_cpu_tensor[i1], - self.allowed_token_ids_mask_cpu_tensor[i2], - ) = ( - self.allowed_token_ids_mask_cpu_tensor[i2], - self.allowed_token_ids_mask_cpu_tensor[i1], - ) - self.block_table.swap_row(i1, i2) - - def condense(self, empty_req_indices: list[int]) -> None: - """Move non-empty requests down into lower, empty indices. - - Args: - empty_req_indices: empty batch indices, sorted descending. - """ - num_reqs = self.num_reqs - if num_reqs == 0: - # The batched states are empty. - self._req_ids.clear() - self.req_output_token_ids.clear() - return - - # NOTE(woosuk): This function assumes that the empty_req_indices - # is sorted in descending order. - last_req_index = num_reqs + len(empty_req_indices) - 1 - while empty_req_indices: - # Find the largest non-empty index. - while last_req_index in empty_req_indices: - last_req_index -= 1 - - # Find the smallest empty index. - empty_index = empty_req_indices.pop() - if empty_index >= last_req_index: - break - - # Swap the states. - req_id = self._req_ids[last_req_index] - output_token_ids = self.req_output_token_ids[last_req_index] - assert req_id is not None - self._req_ids[empty_index] = req_id - self._req_ids[last_req_index] = None - self.req_output_token_ids[empty_index] = output_token_ids - self.req_output_token_ids[last_req_index] = None - self.req_id_to_index[req_id] = empty_index - - num_tokens = self.num_tokens_no_spec[last_req_index] - self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[ - last_req_index, :num_tokens - ] - self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[ - last_req_index - ] - self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index] - self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[ - last_req_index - ] - self.block_table.move_row(last_req_index, empty_index) - self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index] - self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index] - self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index] - self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[ - last_req_index - ] - self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[ - last_req_index - ] - self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[ - last_req_index - ] - self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index] - generator = self.generators.pop(last_req_index, None) - if generator is not None: - self.generators[empty_index] = generator - - min_token = self.min_tokens.pop(last_req_index, None) - if min_token is not None: - self.min_tokens[empty_index] = min_token - - self.request_lora_mapping[empty_index] = self.request_lora_mapping[ - last_req_index - ] - - self.logit_bias[empty_index] = self.logit_bias[last_req_index] - - if self.allowed_token_ids_mask_cpu_tensor is not None: - self.allowed_token_ids_mask_cpu_tensor[empty_index] = ( - self.allowed_token_ids_mask_cpu_tensor[last_req_index] - ) - - bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None) - if bad_words_token_ids is not None: - self.bad_words_token_ids[empty_index] = bad_words_token_ids - # Decrement last_req_index since it is now empty. - last_req_index -= 1 - - # Trim lists to the batch size. - del self._req_ids[self.num_reqs :] - del self.req_output_token_ids[self.num_reqs :] - - def make_lora_inputs( - self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray - ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]: - """ - Given the num_scheduled_tokens for each request in the batch, return - datastructures used to activate the current LoRAs. - Returns: - 1. prompt_lora_mapping: A tuple of size self.num_reqs where, - prompt_lora_mapping[i] is the LoRA id to use for the ith prompt. - 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens) - where, token_lora_mapping[i] is the LoRA id to use for ith token. - 3. lora_requests: Set of relevant LoRA requests. - """ - - req_lora_mapping = self.request_lora_mapping[: self.num_reqs] - prompt_lora_mapping = tuple(req_lora_mapping) - token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens)) - active_lora_requests: set[LoRARequest] = set( - self.lora_id_to_lora_request.values() - ) - - return prompt_lora_mapping, token_lora_mapping, active_lora_requests - - @property - def num_reqs(self) -> int: - return len(self.req_id_to_index) - - @property - def all_greedy(self) -> bool: - return len(self.random_reqs) == 0 - - @property - def all_random(self) -> bool: - return len(self.greedy_reqs) == 0 - - @property - def no_top_p(self) -> bool: - return len(self.top_p_reqs) == 0 - - @property - def no_top_k(self) -> bool: - return len(self.top_k_reqs) == 0 - - @property - def no_min_p(self) -> bool: - return len(self.min_p_reqs) == 0 - - @property - def no_penalties(self) -> bool: - return ( - len(self.presence_penalties_reqs) == 0 - and len(self.frequency_penalties_reqs) == 0 - and len(self.repetition_penalties_reqs) == 0 - ) - - @property - def max_num_logprobs(self) -> int | None: - return max(self.num_logprobs.values()) if self.num_logprobs else None - - @property - def no_allowed_token_ids(self) -> bool: - return len(self.has_allowed_token_ids) == 0 diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index 68d6bde3a049..5c5d125f4bfa 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -115,27 +115,55 @@ def init_device(self): self.local_rank += dp_local_rank * replica_world_size - assert self.local_rank < visible_device_count, ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " - ) - assert parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." - ) - device = self.device_config.device if ( isinstance(device, torch.device) and device.type == "xpu" and current_platform.is_xpu() ): - self.device = torch.device(f"xpu:{self.local_rank}") + assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" + ) + # NOTE: local_world_size is derived from parallel_config.nnodes, + # which is only set for the "mp" multi-node backend. With the + # "ray"/"external_launcher" backends nnodes stays 1, so + # local_world_size collapses to the full world_size and this + # check wrongly fires on cross-node deployments. + # assigned_physical_gpu_ids is already per-node and the + # local_rank bound above fully validates the mapping for + # these backends, so skip the check for them. + if parallel_config.distributed_executor_backend not in ( + "ray", + "external_launcher", + ): + assert parallel_config.local_world_size <= len( + assigned_physical_gpu_ids + ), ( + f"local_world_size ({parallel_config.local_world_size})" + " exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: + assert self.local_rank < torch.accelerator.device_count(), ( + f"DP adjusted local rank {self.local_rank} is out of " + f"bounds for {torch.accelerator.device_count()} devices." + ) + + visible_device_index = ( + current_platform.logical_device_id_to_visible_device_id(self.local_rank) + ) + self.device = torch.device(f"xpu:{visible_device_index}") torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) torch.accelerator.empty_cache() self.init_gpu_memory = torch.xpu.get_device_properties( - self.local_rank + visible_device_index ).total_memory else: raise RuntimeError(f"Unsupported device type: {self.device_config.device}")