diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 9b9f70b58d13..efd3e30f6f89 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -11,6 +11,7 @@ steps: - CMakeLists.txt - vllm/_custom_ops.py - tests/kernels/attention/test_cpu_attn.py + - tests/v1/attention/test_group_head_counts.py - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/moe/test_cpu_quant_fused_moe.py - tests/kernels/test_onednn.py @@ -27,6 +28,7 @@ steps: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s tests/kernels/attention/test_cpu_attn.py + pytest -x -v -s tests/v1/attention/test_group_head_counts.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py diff --git a/.buildkite/intel_jobs/kernels_intel.yaml b/.buildkite/intel_jobs/kernels_intel.yaml index e914c10d8bfe..43f936ee30cd 100644 --- a/.buildkite/intel_jobs/kernels_intel.yaml +++ b/.buildkite/intel_jobs/kernels_intel.yaml @@ -24,3 +24,27 @@ steps: 'cd tests && pytest -v -s ir && pytest -v -s kernels/ir' + +- label: MoE WNA16 Tensor-Descriptor Tests + timeout_in_minutes: 20 + device: intel_gpu + agent_tags: + label: production + gpu: 1+ + mem: 16+ + no_plugin: true + working_dir: "." + env: + REGISTRY: "public.ecr.aws/q9t5s3a7" + REPO: "vllm-ci-test-repo" + VLLM_TEST_DEVICE: "xpu" + source_file_dependencies: + - vllm/model_executor/layers/fused_moe/fused_moe.py + - vllm/model_executor/layers/fused_moe/utils.py + - vllm/model_executor/layers/fused_moe/experts/triton_moe.py + - tests/kernels/moe/test_fused_moe_kernel_gptq_awq.py + commands: + - >- + bash .buildkite/scripts/hardware_ci/run-intel-test.sh + 'cd tests && + pytest -v -s kernels/moe/test_fused_moe_kernel_gptq_awq.py' diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 4c7186e24709..7b9cac7b5257 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -164,7 +164,21 @@ steps: env: DOCKER_BUILDKIT: "1" - - label: "Publish XPU Triton shim index" + - label: "Build wheel - x86_64 - XPU" + depends_on: ~ + id: build-wheel-x86-xpu + agents: + queue: cpu_queue_release + commands: + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.xpu ." + - "mkdir artifacts" + - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels' + env: + DOCKER_BUILDKIT: "1" + + - label: "Publish and stage XPU Triton shim" key: publish-xpu-triton-shim depends_on: ~ agents: diff --git a/.buildkite/scripts/generate-and-upload-nightly-index.sh b/.buildkite/scripts/generate-and-upload-nightly-index.sh index 1fa75994c01a..ee13f316881a 100755 --- a/.buildkite/scripts/generate-and-upload-nightly-index.sh +++ b/.buildkite/scripts/generate-and-upload-nightly-index.sh @@ -53,11 +53,10 @@ if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \ aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/" fi -# detect version from any wheel in the commit directory -# download the first wheel we find to extract version metadata -first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))") +# detect version from a vLLM wheel in the commit directory +first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].rsplit('/', 1)[-1].startswith('vllm-') and c['Key'].endswith('.whl')), ''))") if [[ -z "$first_wheel_key" ]]; then - echo "Error: No wheels found in $S3_COMMIT_PREFIX" + echo "Error: No vLLM wheel found in $S3_COMMIT_PREFIX" exit 1 fi first_wheel=$(basename "$first_wheel_key") diff --git a/.buildkite/scripts/generate-nightly-index.py b/.buildkite/scripts/generate-nightly-index.py index 9397825bddd9..67adf9c92b80 100644 --- a/.buildkite/scripts/generate-nightly-index.py +++ b/.buildkite/scripts/generate-nightly-index.py @@ -91,9 +91,9 @@ def parse_from_filename(file: str) -> WheelFileInfo: else: if "+" in version: version_part, suffix = version.split("+", 1) - # Only treat known patterns as variants (rocmXXX, cuXXX, cpu) + # Only treat known patterns as variants (rocmXXX, cuXXX, cpu, xpu) # Git hashes and other suffixes are NOT variants - if suffix.startswith(("rocm", "cu", "cpu")): + if suffix.startswith(("rocm", "cu", "cpu", "xpu")): variant = suffix version = version_part # Otherwise keep the full version string (variant stays None) @@ -429,7 +429,11 @@ def generate_index_and_metadata( if PY_VERSION_RE.match(version): # upload-wheels.sh ensures no "dev" is in args.version wheel_files = list( - filter(lambda x: version in x and "dev" not in x, wheel_files) + filter( + lambda x: (version in x and "dev" not in x) + or (x.startswith("triton-") and "+xpu-" in x), + wheel_files, + ) ) print(f"Non-nightly version detected, wheel files used: {wheel_files}") else: diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 42fe3f45460a..86e4b34869bb 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -261,6 +261,8 @@ validate_native_workspace() { } prepare_native_workspace() { + local test_commands="${1:-}" + if [[ "${VLLM_CI_USE_ARTIFACTS:-0}" != "1" ]]; then echo "Native CI requires VLLM_CI_USE_ARTIFACTS=1" return 1 @@ -281,6 +283,8 @@ prepare_native_workspace() { local recorded_base="" local recorded_commit="" local recorded_wheel="" + local checkout="" + local checkout_commit="" local workspace_dir="${VLLM_CI_WORKSPACE:-/vllm-workspace}" local wheel_dir="" local attempt=0 @@ -400,6 +404,53 @@ prepare_native_workspace() { return 1 fi + # The ROCm artifact intentionally contains only the installed wheel and the + # test workspace. The Python-only compilation job also needs setup.py and the + # vllm source tree, so overlay the matching Buildkite checkout for that job. + if [[ "${test_commands}" == *python_only_compile.sh* ]]; then + checkout="${BUILDKITE_BUILD_CHECKOUT_PATH:-}" + if [[ -z "${checkout}" || ! -d "${checkout}" ]]; then + echo "Python-only native CI requires BUILDKITE_BUILD_CHECKOUT_PATH" >&2 + return 1 + fi + if ! git -c "safe.directory=${checkout}" -C "${checkout}" \ + rev-parse --is-inside-work-tree >/dev/null 2>&1; then + echo "Buildkite checkout is not a Git worktree: ${checkout}" >&2 + return 1 + fi + checkout_commit=$( + git -c "safe.directory=${checkout}" -C "${checkout}" rev-parse HEAD + ) || return 1 + if [[ "${checkout_commit}" != "${recorded_commit}" ]]; then + echo "Buildkite checkout ${checkout_commit} does not match ROCm artifact ${recorded_commit}" >&2 + return 1 + fi + + # setup.py normally derives this from .git via setuptools-scm. The native + # source overlay deliberately excludes Git metadata, so preserve the exact + # version from the already installed, artifact-matched wheel. + VLLM_VERSION_OVERRIDE=$( + python3 -c 'import importlib.metadata as m; print(m.version("vllm"))' + ) || return 1 + export VLLM_VERSION_OVERRIDE + VLLM_PRECOMPILED_WHEEL_LOCATION="${wheels[0]}" + export VLLM_PRECOMPILED_WHEEL_LOCATION + echo "INFO: native Python-only wheel=${VLLM_PRECOMPILED_WHEEL_LOCATION}" + + echo "--- Overlaying full source checkout for Python-only compilation" + # Archive the verified commit instead of copying the worktree so dirty or + # untracked agent files cannot contaminate the artifact-matched workspace. + git -c "safe.directory=${checkout}" -C "${checkout}" \ + archive --format=tar "${recorded_commit}" \ + | tar --no-same-owner -C "${workspace_dir}" -xf - || return 1 + for required_source in setup.py pyproject.toml vllm; do + if [[ ! -e "${workspace_dir}/${required_source}" ]]; then + echo "Full source checkout is missing ${required_source}" >&2 + return 1 + fi + done + fi + return 0 } @@ -434,6 +485,12 @@ initialize_native_environment() { TIKTOKEN_RS_CACHE_DIR="${HF_HOME}/tiktoken-rs-cache" : "${HF_HUB_DOWNLOAD_TIMEOUT:=300}" : "${HF_HUB_ETAG_TIMEOUT:=60}" + if [[ "${VLLM_CI_EXPECTED_GPU_COUNT:-1}" == "0" ]]; then + # CPU-only native jobs intentionally reuse the ROCm wheel. Make that target + # explicit so platform selection does not depend on wheel metadata. + VLLM_TARGET_DEVICE=cpu + export VLLM_TARGET_DEVICE + fi export TMPDIR VLLM_RPC_BASE_PATH export TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT XDG_CACHE_HOME export HF_HOME HF_DATASETS_CACHE HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT @@ -948,7 +1005,13 @@ if is_native_runtime; then echo "Failed to initialize the native test environment" exit 1 fi - if ! prepare_native_workspace; then + if [[ "${commands}" == *python_only_compile.sh* ]]; then + # This no-GPU job validates the ROCm precompiled/editable install path, + # rather than CPU runtime platform selection. + VLLM_TARGET_DEVICE=rocm + export VLLM_TARGET_DEVICE + fi + if ! prepare_native_workspace "${commands}"; then echo "Failed to prepare native test workspace" exit 1 fi diff --git a/.buildkite/scripts/xpu/publish-triton-shim.sh b/.buildkite/scripts/xpu/publish-triton-shim.sh index d85a95a5b4d7..099b8d974969 100644 --- a/.buildkite/scripts/xpu/publish-triton-shim.sh +++ b/.buildkite/scripts/xpu/publish-triton-shim.sh @@ -11,12 +11,17 @@ readonly WHEEL_SHA256="3c822f73e9870512f59a6ecf5dc305a4bcab11fa623f9ce91011f6043 readonly WHEEL_FILENAME="${WHEEL_URL##*/}" readonly ENCODED_WHEEL_FILENAME="${WHEEL_FILENAME/+/%2B}" readonly S3_PREFIX="s3://${BUCKET}/${PREFIX}/" +readonly COMMIT="${BUILDKITE_COMMIT:-}" readonly DRY_RUN="${DRY_RUN:-0}" if [[ "$DRY_RUN" != "0" && "$DRY_RUN" != "1" ]]; then echo "DRY_RUN must be 0 or 1" >&2 exit 2 fi +if [[ "$DRY_RUN" == "0" && ! "$COMMIT" =~ ^[0-9a-f]{40}$ ]]; then + echo "BUILDKITE_COMMIT must be a full lowercase commit hash" >&2 + exit 2 +fi cd "$(dirname "${BASH_SOURCE[0]}")/../../.." @@ -70,9 +75,9 @@ sed 's/import regex as re/import re/' \ # shellcheck disable=SC2086 $PYTHON "$index_generator" \ --version "$PREFIX" \ - --wheel-dir "$PREFIX" \ + --wheel-dir "$work_dir/$PREFIX" \ --current-objects "$objects_path" \ - --output-dir "$index_output_dir" \ + --output-dir "$work_dir" \ --comment "XPU Triton shim" grep -Fq 'href="triton/"' "$index_output_dir/index.html" @@ -87,4 +92,7 @@ if [[ "$DRY_RUN" == "1" ]]; then else aws s3 cp --recursive "$index_output_dir/" "$S3_PREFIX" echo "Published XPU Triton shim index to https://wheels.vllm.ai/$PREFIX/" -fi \ No newline at end of file + aws s3 cp "$S3_PREFIX$WHEEL_FILENAME" \ + "s3://$BUCKET/$COMMIT/$WHEEL_FILENAME" + echo "Staged XPU Triton shim for https://wheels.vllm.ai/$COMMIT/xpu/" +fi diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index df65ee762dba..b477b298bb10 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -411,8 +411,9 @@ steps: - label: Python-only Installation # TBD timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: @@ -2092,7 +2093,7 @@ steps: - tests/models/multimodal/generation - tests/models/multimodal/test_mapping.py commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model @@ -3671,7 +3672,8 @@ steps: - vllm/ - tests/models/multimodal/generation commands: - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model - pytest -v -s models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model @@ -3771,6 +3773,7 @@ steps: - tests/v1/executor - tests/v1/kv_offload - tests/v1/worker + - tests/v1/cudagraph - tests/v1/kv_connector/unit - tests/v1/metrics - tests/entrypoints/openai/correctness/test_lmeval.py @@ -3780,6 +3783,7 @@ steps: - pytest -v -s v1/executor - pytest -v -s v1/kv_offload - pytest -v -s v1/worker + - pytest -v -s v1/cudagraph/test_encoder_cudagraph.py - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit - pytest -v -s -m 'not cpu_test' v1/metrics - pip install -U git+https://github.com/vllm-project/lm-evaluation-harness.git@streaming-api diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 0c6fbdae21e7..fb0329d58355 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -242,6 +242,7 @@ steps: commands: - pytest -v -s tests/distributed/test_context_parallel.py - pytest -v -s tests/distributed/test_nccl_symm_mem.py + - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs.py - pytest -v -s tests/v1/distributed/test_dbo.py - pytest -v -s tests/distributed/test_mnnvl_alltoall.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4e73d7620e51..887ed03e372a 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -339,6 +339,28 @@ steps: # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py +- label: B12X Linear Kernels (DGX Spark) Nightly + key: b12x-linear-kernels-dgx-spark-nightly + timeout_in_minutes: 30 + device: dgx-spark + optional: true + num_devices: 1 + depends_on: + - arm64-image-build + source_file_dependencies: + - setup.py + - vllm/model_executor/kernels/linear/ + - vllm/model_executor/warmup/b12x_warmup.py + - vllm/utils/b12x.py + - tests/model_executor/kernels/test_b12x_linear.py + - tests/model_executor/test_b12x_warmup.py + - tests/kernels/quantization/test_block_fp8.py + commands: + - uv pip install --system b12x==1.2.4 + - pytest -v -s model_executor/kernels/test_b12x_linear.py + model_executor/test_b12x_warmup.py + kernels/quantization/test_block_fp8.py -k b12x + - label: Kernels Helion Test key: kernels-helion-test timeout_in_minutes: 115 diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 84ebf0d13aff..f9c21a54217b 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -178,13 +178,14 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt -- label: MoE Refactor Integration Test (B200 - TEMPORARY) +- label: MoE Refactor Integration Test (B200 - TEMPORARY) %N key: moe-refactor-integration-test-b200-temporary device: b200-k8s optional: true num_devices: 2 + parallelism: 4 commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200-shard-$$BUILDKITE_PARALLEL_JOB.txt - label: MoE Refactor Integration Test (B200 DP - TEMPORARY) key: moe-refactor-integration-test-b200-dp-temporary @@ -194,12 +195,13 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -- label: LM Eval Humming f16 (A100 - TEMPORARY) +- label: LM Eval Humming f16 (A100 - TEMPORARY) %N key: lm-eval-humming-f16-a100 timeout_in_minutes: 75 device: a100 optional: true num_devices: 1 + parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -207,7 +209,7 @@ steps: - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-a100-shard-$$BUILDKITE_PARALLEL_JOB.txt - label: LM Eval Humming Act int8 (A100 - TEMPORARY) key: lm-eval-humming-act-a100 @@ -224,12 +226,13 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt -- label: LM Eval Humming f16 (H100 - TEMPORARY) +- label: LM Eval Humming f16 (H100 - TEMPORARY) %N key: lm-eval-humming-f16-h100 timeout_in_minutes: 70 device: h100 optional: true num_devices: 1 + parallelism: 3 source_file_dependencies: - vllm/model_executor/layers/quantization/humming.py - vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -237,7 +240,7 @@ steps: - vllm/model_executor/layers/fused_moe/oracle/ - vllm/model_executor/kernels/linear/ commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-h100-shard-$$BUILDKITE_PARALLEL_JOB.txt - label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY) key: lm-eval-humming-act-h100 diff --git a/.buildkite/test_areas/lora.yaml b/.buildkite/test_areas/lora.yaml index a79196ffbd8d..4287816b90c2 100644 --- a/.buildkite/test_areas/lora.yaml +++ b/.buildkite/test_areas/lora.yaml @@ -48,3 +48,4 @@ steps: - pytest -v -s -x lora/test_olmoe_tp.py - pytest -v -s -x lora/test_gptoss_tp.py - pytest -v -s -x lora/test_qwen35_densemodel_lora.py + - pytest -v -s -x lora/test_gemma4_tp.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 8c80fead290f..8f168cbd367f 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -127,6 +127,7 @@ steps: - pytest -v -s v1/test_kv_cache_spec_registry.py - pytest -v -s v1/cudagraph/test_cudagraph_manager.py - pytest -v -s -m 'cpu_test' v1/kv_connector/unit + - pytest -v -s -m 'cpu_test' v1/ec_connector/unit - pytest -v -s -m 'cpu_test' v1/metrics - label: Extract Hidden States Integration @@ -274,7 +275,8 @@ steps: - bash standalone_tests/python_only_compile.sh mirror: amd: - device: mi250_1 + dind: false + device: mi300_1 timeout_in_minutes: 55 soft_fail: true depends_on: diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index ab1d9bebaa00..d0eff968f7c8 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -100,7 +100,7 @@ steps: - label: Language Models Test (Extended Generation) # 80min device: h200_35gb key: language-models-test-extended-generation - timeout_in_minutes: 65 + timeout_in_minutes: 80 optional: true source_file_dependencies: - vllm/ @@ -124,24 +124,28 @@ steps: commands: - pytest -v -s models/language/generation_ppl_test -- label: Language Models Test (Extended Pooling) +- label: Language Models Test (Extended Pooling) %N device: h200_35gb key: language-models-test-extended-pooling timeout_in_minutes: 120 + parallelism: 4 optional: true source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" - tests/models/language/pooling commands: - - pytest -v -s models/language/pooling -m 'not core_model' + - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB mirror: amd: dind: false device: mi300_1 timeout_in_minutes: 95 + parallelism: 2 depends_on: - image-build-amd + commands: + - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: Language Models Test (MTEB) key: language-models-test-mteb diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 4a84abd5baa7..2ea3621342fc 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -169,16 +169,17 @@ steps: depends_on: - image-build-amd -- label: Multi-Modal Models (Extended Generation 2) +- label: Multi-Modal Models (Extended Generation 2) %N device: h200_35gb key: multi-modal-models-extended-generation-2 + parallelism: 4 optional: true source_file_dependencies: - vllm/ - "!vllm/distributed/kv_transfer/" - tests/models/multimodal/generation commands: - - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' + - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - label: Multi-Modal Models (Extended Generation 3) device: h200_35gb diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index ecaf28795654..4ba339207196 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -75,6 +75,9 @@ steps: - vllm/model_executor/layers/quantization - tests/plugins_tests/bitsandbytes commands: + # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so + # there is nothing to wrap on the vLLM side. + - unset VLLM_GPU_SYNC_CHECK - pip install "vllm-bnb-plugin >= 0.0.1" - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed' @@ -89,5 +92,8 @@ steps: - vllm/model_executor/layers/quantization - tests/plugins_tests/bitsandbytes commands: + # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so + # there is nothing to wrap on the vLLM side. + - unset VLLM_GPU_SYNC_CHECK - pip install "vllm-bnb-plugin >= 0.0.1" - pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)' diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index a7376945396c..b5c7c7a66dac 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -2,10 +2,11 @@ group: Quantization depends_on: - image-build steps: -- label: Quantization +- label: Quantization %N device: h200_35gb key: quantization - timeout_in_minutes: 75 + timeout_in_minutes: 40 + parallelism: 4 env: VLLM_USE_V2_MODEL_RUNNER: "0" source_file_dependencies: @@ -18,7 +19,7 @@ steps: - uv pip install --system conch-triton-kernels # The SM90-only checkpoint currently contains a removed weight_chan_scale # parameter. It was not exercised by the previous L4 job. - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - label: Quantized Fusions device: h200_35gb diff --git a/.github/mergify.yml b/.github/mergify.yml index a7e25ad92a68..4e6de638ab59 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -68,6 +68,23 @@ pull_request_rules: add: - ci/build +- name: label-cohere + description: Automatically apply cohere label + conditions: + - label != stale + - or: + - files~=^examples/.*cohere.*\.py + - files~=^tests/.*cohere.*\.py + - files~=^vllm/model_executor/models/.*cohere.*\.py + - files~=^vllm/tool_parsers/.*cohere.*\.py + - files~=^vllm/reasoning/.*cohere.*\.py + - files~=^vllm/transformers_utils/.*cohere.*\.py + - title~=(?i)Cohere + actions: + label: + add: + - cohere + - name: label-deepseek description: Automatically apply deepseek label conditions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 5e07d228b7f6..4a38aa2fbd2e 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -386,6 +386,10 @@ jobs: users: ['hongxiayang', 'tjtanaa', 'vllmellm', 'giuseppegrossi'], message: 'CC {users} for ROCm-related issue', }, + cohere: { + users: ['walterbm', 'ekagra-ranjan', 'andrewbcohere', 'jasonozuzu-cohere'], + message: 'CC {users} for Cohere-related issue', + }, mistral: { users: ['patrickvonplaten', 'juliendenize', 'andylolu2', 'NickLucche'], message: 'CC {users} for Mistral-related issue', diff --git a/CMakeLists.txt b/CMakeLists.txt index cdcb41bef605..64de3baa9fdf 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -222,7 +222,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") # `+PTX` in TORCH_CUDA_ARCH_LIST is not preserved here. If a kernel really # needs PTX, add `+PTX` to that kernel's component-specific arch list below. # - clear_cuda_arches(CUDA_ARCH_FLAGS) + clear_cuda_gencode_flags(CUDA_ARCH_FLAGS) + warn_if_ptx_arch_requested("${CUDA_ARCH_FLAGS}") extract_unique_cuda_archs_ascending(CUDA_ARCHS "${CUDA_ARCH_FLAGS}") message(STATUS "CUDA target architectures: ${CUDA_ARCHS}") # Filter the target architectures by the supported supported archs @@ -1118,6 +1119,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(FUSED_KDA_DECODE_ARCHS "9.0a;10.0f;12.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(FUSED_GDN_DECODE_ARCHS + "8.0;8.6;8.9;9.0a;10.0f;12.0f" "${CUDA_ARCHS}") endif() if(FUSED_KDA_DECODE_ARCHS) set(FUSED_KDA_DECODE_SRC @@ -1131,6 +1134,18 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") message(STATUS "Building fused KDA decode for archs: ${FUSED_KDA_DECODE_ARCHS}") endif() + if(FUSED_GDN_DECODE_ARCHS) + set(FUSED_GDN_DECODE_SRC + "csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu") + set_gencode_flags_for_srcs( + SRCS "${FUSED_GDN_DECODE_SRC}" + CUDA_ARCHS "${FUSED_GDN_DECODE_ARCHS}") + set_property(SOURCE ${FUSED_GDN_DECODE_SRC} APPEND PROPERTY + COMPILE_OPTIONS "$<$:--use_fast_math>") + list(APPEND VLLM_STABLE_EXT_SRC "${FUSED_GDN_DECODE_SRC}") + message(STATUS + "Building fused GDN decode for archs: ${FUSED_GDN_DECODE_ARCHS}") + endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS @@ -1213,6 +1228,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_FUSED_KDA_DECODE=1) endif() + if(FUSED_GDN_DECODE_ARCHS) + target_compile_definitions(_C_stable_libtorch PRIVATE + VLLM_ENABLE_FUSED_GDN_DECODE=1) + endif() if(KIMI_K3_ATTN_RES_ARCHS) target_compile_definitions(_C_stable_libtorch PRIVATE VLLM_ENABLE_KIMI_K3_ATTN_RES=1) diff --git a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs.py b/benchmarks/kernels/benchmark_kimi_k3_gemm_rs.py new file mode 100644 index 000000000000..cb04bf86e303 --- /dev/null +++ b/benchmarks/kernels/benchmark_kimi_k3_gemm_rs.py @@ -0,0 +1,433 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the SM100 Kimi-K3 GEMM-RS kernel. + +All ranks must belong to one NVLink domain. For example, run a TP8 sweep with: + + torchrun --nproc-per-node=8 \ + benchmarks/kernels/benchmark_kimi_k3_gemm_rs.py +""" + +import argparse +import os +import statistics +from collections.abc import Callable +from dataclasses import dataclass + +import pandas as pd +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.distributed.parallel_state import ( + get_tp_group, + init_distributed_environment, + initialize_model_parallel, +) +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS + +# Shared-expert down-proj and attention O-proj. +_KIMI_K3_PROJECTION_K = (6144, 12288) + + +@dataclass +class Candidate: + name: str + runs: list[Callable[[], torch.Tensor]] + check_correctness: bool = True + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--m", + type=int, + nargs="+", + default=[128, 512, 2048, 8192, 32768], + help="Global token counts to benchmark.", + ) + parser.add_argument( + "--k", + type=int, + nargs="+", + help=( + "Per-rank input dimensions. By default, derive the Kimi-K3 " + "shared-expert down-proj and O-proj dimensions from the TP " + "world size." + ), + ) + parser.add_argument("--n", type=int, default=7168) + parser.add_argument( + "--num-workspaces", + type=int, + default=10, + help="Pointer-distinct inputs and CUDA graphs to rotate.", + ) + parser.add_argument("--warmup-replays", type=int, default=5) + parser.add_argument("--samples", type=int, default=20) + return parser.parse_args() + + +def capture_graph( + op: Callable[[], torch.Tensor], + stream: torch.cuda.Stream, + cpu_group: dist.ProcessGroup, +) -> tuple[torch.cuda.CUDAGraph, list[torch.Tensor | None]]: + result: list[torch.Tensor | None] = [None] + stream.wait_stream(torch.cuda.current_stream()) + dist.barrier(group=cpu_group) + with torch.cuda.stream(stream): + for _ in range(3): + result[0] = op() + stream.synchronize() + dist.barrier(group=cpu_group) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result[0] = op() + torch.cuda.current_stream().wait_stream(stream) + dist.barrier(group=cpu_group) + return graph, result + + +def benchmark_graphs( + candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]], + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + device_barrier: Callable[[], None], +) -> dict[str, float]: + candidate_names = list(candidate_graphs) + for round_index in range(warmup_replays): + for candidate_index in range(len(candidate_names)): + candidate_id = (round_index + candidate_index) % len(candidate_names) + name = candidate_names[candidate_id] + graphs = candidate_graphs[name] + device_barrier() + graphs[round_index % len(graphs)].replay() + torch.accelerator.synchronize() + + timings: dict[str, list[float]] = {name: [] for name in candidate_names} + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for sample_index in range(samples): + for candidate_index in range(len(candidate_names)): + candidate_id = (sample_index + candidate_index) % len(candidate_names) + name = candidate_names[candidate_id] + graphs = candidate_graphs[name] + device_barrier() + start.record() + graphs[sample_index % len(graphs)].replay() + end.record() + end.synchronize() + + elapsed = torch.tensor( + start.elapsed_time(end) * 1000, + dtype=torch.float64, + device=torch.accelerator.current_device_index(), + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) + timings[name].append(elapsed.item()) + return {name: statistics.median(values) for name, values in timings.items()} + + +def valid_rows(M: int, local_M: int, rank: int) -> int: + return min(max(M - rank * local_M, 0), local_M) + + +def benchmark_shape( + gemm_rs: GemmRS, + M: int, + N: int, + K: int, + num_workspaces: int, + warmup_replays: int, + samples: int, + device_group: dist.ProcessGroup, + cpu_group: dist.ProcessGroup, + device_barrier: Callable[[], None], +) -> dict[str, float | int]: + world_size = dist.get_world_size(device_group) + rank = dist.get_rank(device_group) + device = torch.device("cuda", torch.accelerator.current_device_index()) + padded_M = (M + world_size - 1) // world_size * world_size + local_M = padded_M // world_size + + rng = torch.Generator(device=device) + rng.manual_seed(1000 + rank * 10 + M + K) + inputs = [ + torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng) + for _ in range(num_workspaces) + ] + weights = [ + torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng) + for _ in range(num_workspaces) + ] + + partial = torch.empty((padded_M, N), dtype=torch.bfloat16, device=device) + symm_partial = symm_mem.empty((padded_M, N), dtype=torch.bfloat16, device=device) + symm_partial_handle = symm_mem.rendezvous(symm_partial, device_group) + rs_inputs = [torch.empty_like(partial) for _ in range(num_workspaces)] + symm_rs_inputs = [] + symm_rs_handles = [] + for _ in range(num_workspaces): + rs_input = symm_mem.empty( + (padded_M, N), + dtype=torch.bfloat16, + device=device, + ) + symm_rs_inputs.append(rs_input) + symm_rs_handles.append(symm_mem.rendezvous(rs_input, device_group)) + torch_output = torch.empty((local_M, N), dtype=torch.bfloat16, device=device) + symm_output = torch.empty_like(torch_output) + gemm_output = torch.empty((M, N), dtype=torch.bfloat16, device=device) + + if padded_M > M: + partial[M:].zero_() + symm_partial[M:].zero_() + + def make_torch_ring_ll_gemm_rs( + x: torch.Tensor, weight: torch.Tensor + ) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + torch.mm(x, weight.T, out=partial[:M]) + dist.reduce_scatter_single(torch_output, partial, group=device_group) + return torch_output + + return run + + def make_torch_ldmc_gemm_rs( + x: torch.Tensor, weight: torch.Tensor + ) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + torch.mm(x, weight.T, out=symm_partial[:M]) + dist.reduce_scatter_single( + symm_output, + symm_partial, + group=device_group, + ) + return symm_output + + return run + + def make_fused_gemm_rs( + x: torch.Tensor, weight: torch.Tensor + ) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + return gemm_rs(x, weight) + + return run + + def make_torch_gemm( + x: torch.Tensor, + weight: torch.Tensor, + ) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + return torch.mm(x, weight.T, out=gemm_output) + + return run + + def make_ring_ll_rs(rs_input: torch.Tensor) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + dist.reduce_scatter_single(torch_output, rs_input, group=device_group) + return torch_output + + return run + + def make_ldmc_rs( + rs_input: torch.Tensor, + ) -> Callable[[], torch.Tensor]: + def run() -> torch.Tensor: + dist.reduce_scatter_single(symm_output, rs_input, group=device_group) + return symm_output + + return run + + candidates = ( + Candidate( + "ring_ll_us", + [make_torch_ring_ll_gemm_rs(x, w) for x, w in zip(inputs, weights)], + ), + Candidate( + "ldmc_us", + [make_torch_ldmc_gemm_rs(x, w) for x, w in zip(inputs, weights)], + ), + Candidate( + "gemm_rs_us", + [make_fused_gemm_rs(x, w) for x, w in zip(inputs, weights)], + ), + Candidate( + "torch_gemm_us", + [make_torch_gemm(x, w) for x, w in zip(inputs, weights)], + check_correctness=False, + ), + Candidate( + "ring_ll_rs_us", + [make_ring_ll_rs(x) for x in rs_inputs], + check_correctness=False, + ), + Candidate( + "ldmc_rs_us", + [make_ldmc_rs(x) for x in symm_rs_inputs], + check_correctness=False, + ), + ) + + expected = candidates[0].runs[0]() + rows = valid_rows(M, local_M, rank) + for candidate in candidates[1:]: + if not candidate.check_correctness: + continue + actual = candidate.runs[0]() + torch.accelerator.synchronize(device) + torch.testing.assert_close( + actual[:rows], + expected[:rows], + rtol=5e-2, + atol=4.0, + ) + + candidate_graphs = {} + graph_keepalive: list[object] = [symm_partial_handle, *symm_rs_handles] + for candidate in candidates: + stream = torch.cuda.Stream() + bundles = [capture_graph(run, stream, cpu_group) for run in candidate.runs] + candidate_graphs[candidate.name] = [graph for graph, _ in bundles] + graph_keepalive.extend(bundles) + graph_keepalive.append(stream) + + times = benchmark_graphs( + candidate_graphs, + warmup_replays, + samples, + device_group, + device_barrier, + ) + + best_nccl_rs_us = min(times["ring_ll_rs_us"], times["ldmc_rs_us"]) + return { + "M": M, + "N": N, + "K": K, + **times, + "best_nccl_rs_us": best_nccl_rs_us, + "speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_us"], + "speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_us"], + } + + +def print_results(results: list[dict[str, float | int]]) -> None: + results_df = pd.DataFrame(results) + end_to_end = results_df[ + [ + "M", + "N", + "K", + "ring_ll_us", + "ldmc_us", + "gemm_rs_us", + "speedup_vs_ring_ll", + "speedup_vs_ldmc", + ] + ].rename( + columns={ + "ring_ll_us": "Torch GEMM + NCCL RS (RING_LL) (us)", + "ldmc_us": "Torch GEMM + NCCL RS (LDMC) (us)", + "gemm_rs_us": "GEMM-RS (us)", + "speedup_vs_ring_ll": "Speedup vs RING_LL", + "speedup_vs_ldmc": "Speedup vs LDMC", + } + ) + end_to_end = end_to_end.round( + { + "Torch GEMM + NCCL RS (RING_LL) (us)": 2, + "Torch GEMM + NCCL RS (LDMC) (us)": 2, + "GEMM-RS (us)": 2, + "Speedup vs RING_LL": 3, + "Speedup vs LDMC": 3, + } + ) + components = results_df[ + ["M", "N", "K", "torch_gemm_us", "best_nccl_rs_us", "gemm_rs_us"] + ].rename( + columns={ + "torch_gemm_us": "Torch GEMM (us)", + "best_nccl_rs_us": "NCCL RS (best) (us)", + "gemm_rs_us": "GEMM-RS (us)", + } + ) + components = components.round(2) + + print("### End-to-end latency") + print(end_to_end.to_markdown(index=False)) + print("\n### Component latency") + print(components.to_markdown(index=False)) + print("\nNCCL RS (best) is the faster of RING_LL and LDMC for each shape.") + + +def main() -> None: + args = parse_args() + assert args.m and min(args.m) >= 128 + assert args.n % 256 == 0 + assert args.num_workspaces > 0 + assert args.warmup_replays >= 0 + assert args.samples > 0 + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(local_rank) + init_distributed_environment() + world_size = dist.get_world_size() + if args.k is None: + assert all(K % world_size == 0 for K in _KIMI_K3_PROJECTION_K) + K_values = [K // world_size for K in _KIMI_K3_PROJECTION_K] + else: + K_values = args.k + assert all(K % 64 == 0 for K in K_values) + # Reserve symmetric memory for the NCCL-managed benchmark allocations. + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" + # NCCL-managed symmetric allocations select the NVLS/LDMC collective path. + symm_mem.set_backend("NCCL") + with set_current_vllm_config(VllmConfig()): + initialize_model_parallel(tensor_model_parallel_size=world_size) + + tp_group = get_tp_group() + group_warmup = torch.zeros(1, device=torch.accelerator.current_device_index()) + dist.all_reduce(group_warmup, group=tp_group.device_group) + pynccl_comm = tp_group.device_communicator.pynccl_comm + assert pynccl_comm is not None + sync_input = torch.zeros(1, device=torch.accelerator.current_device_index()) + sync_output = torch.empty_like(sync_input) + + def device_barrier() -> None: + pynccl_comm.all_reduce(sync_input, sync_output) + + gemm_rs = GemmRS(max_M=max(args.m), N=args.n) + results = [ + benchmark_shape( + gemm_rs, + M, + args.n, + K, + args.num_workspaces, + args.warmup_replays, + args.samples, + tp_group.device_group, + tp_group.cpu_group, + device_barrier, + ) + for K in K_values + for M in args.m + ] + + if tp_group.rank_in_group == 0: + print_results(results) + + dist.barrier(group=tp_group.cpu_group) + del gemm_rs + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_selective_state_update.py b/benchmarks/kernels/benchmark_selective_state_update.py index a8b73da2aa9a..ddf1efd018c8 100644 --- a/benchmarks/kernels/benchmark_selective_state_update.py +++ b/benchmarks/kernels/benchmark_selective_state_update.py @@ -34,6 +34,7 @@ override_ssm_config, selective_state_update, ) +from vllm.platforms import current_platform from vllm.triton_utils import triton # bf16 shares configs with fp16 - same bit width. @@ -93,10 +94,10 @@ def _make_inputs( ngroups: int, dtype: torch.dtype, state_dtype: torch.dtype | None = None, - device: str = "cuda", ): if state_dtype is None: state_dtype = dtype + device = current_platform.device_type state = torch.randn(batch, nheads, dim, dstate, dtype=state_dtype, device=device) x = torch.randn(batch, nheads, dim, dtype=dtype, device=device) dt = torch.randn(batch, nheads, dim, dtype=dtype, device=device) @@ -127,9 +128,9 @@ def benchmark_config( Time one (BLOCK_SIZE_M, num_warps) config for selective_state_update. Returns elapsed time in microseconds, or None on error. - Uses CUDA graph capture-and-replay to isolate kernel time from Python - eager-mode dispatch / kwarg-resolution overhead, mirroring the timing - methodology in benchmarks/kernels/benchmark_moe.py. + Uses accelerator graph capture-and-replay to isolate kernel time from + Python eager-mode dispatch / kwarg-resolution overhead, mirroring the + timing methodology in benchmarks/kernels/benchmark_moe.py. """ state, x, dt, A, B, C, D, dt_bias, out = _make_inputs( batch, nheads, dim, dstate, ngroups, dtype, state_dtype=state_dtype @@ -157,10 +158,15 @@ def _call_kernel() -> None: _call_kernel() torch.accelerator.synchronize() - # Capture graph_batch_size invocations into a CUDA graph so the + # Capture graph_batch_size invocations into a device graph so the # timed region runs without Python dispatch overhead per call. - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): + # Capture via graph(), not the graph object: CUDA needs a side stream. + graph = ( + torch.cuda.CUDAGraph() + if current_platform.is_cuda_alike() + else torch.xpu.XPUGraph() + ) + with current_platform.graph(graph): for _ in range(graph_batch_size): _call_kernel() torch.accelerator.synchronize() @@ -170,8 +176,8 @@ def _call_kernel() -> None: graph.replay() torch.accelerator.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) latencies: list[float] = [] for _ in range(num_iters): start.record() @@ -671,8 +677,8 @@ def main(): dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 state_dtype = _SSM_CACHE_DTYPE_MAP[args.mamba_ssm_cache_dtype] device_name = get_ssm_device_name() - cap = torch.cuda.get_device_capability() - is_blackwell = cap[0] >= 10 + cap = current_platform.get_device_capability() + is_blackwell = cap is not None and cap[0] >= 10 # Mirror all output to a results file (like Unix tee). buf = StringIO() @@ -690,7 +696,8 @@ def flush(self): sys.stdout = _Tee() # type: ignore[assignment] try: - print(f"Device : {device_name} (sm_{cap[0]}{cap[1]})") + cap_str = f"sm_{cap[0]}{cap[1]}" if cap is not None else "n/a" + print(f"Device : {device_name} ({cap_str})") print(f"Blackwell: {is_blackwell}") print(f"dtype : {args.dtype}") print(f"ssm_cache_dtype: {args.mamba_ssm_cache_dtype}") diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index 1d3d163c61bf..988a27d7c755 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -13,7 +13,7 @@ else() FetchContent_Declare( flashkda GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git - GIT_TAG b5d11010ff01c1d4a683c0dde42e76cbeaa8107f + GIT_TAG 053de1b716ef3255873e02d2d28f4adf09951978 GIT_PROGRESS TRUE GIT_SUBMODULES cutlass ) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a8e9cc781bb..b70db3adeb90 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG f3e1a4f74c99145c0717709860bf765de1703779 + GIT_TAG 617264c1c7955c9e84817654ebeedff069f3c5f1 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/cmake/utils.cmake b/cmake/utils.cmake index bbae89c1f57c..6ded936213c6 100644 --- a/cmake/utils.cmake +++ b/cmake/utils.cmake @@ -220,11 +220,11 @@ endmacro() # # Example: # CMAKE_CUDA_FLAGS="-Wall -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75" -# clear_cuda_arches(CUDA_ARCH_FLAGS) +# clear_cuda_gencode_flags(CUDA_ARCH_FLAGS) # CUDA_ARCH_FLAGS="-gencode arch=compute_70,code=sm_70;-gencode arch=compute_75,code=sm_75" # CMAKE_CUDA_FLAGS="-Wall" # -macro(clear_cuda_arches CUDA_ARCH_FLAGS) +macro(clear_cuda_gencode_flags CUDA_ARCH_FLAGS) # Extract all `-gencode` flags from `CMAKE_CUDA_FLAGS` string(REGEX MATCHALL "-gencode arch=[^ ]+" CUDA_ARCH_FLAGS ${CMAKE_CUDA_FLAGS}) @@ -235,6 +235,26 @@ macro(clear_cuda_arches CUDA_ARCH_FLAGS) ${CMAKE_CUDA_FLAGS}) endmacro() +# +# Warn when a caller requested PTX code generation through global CUDA arch +# flags. vLLM removes those flags and reapplies per-source gencode flags, so the +# user's global PTX request will not be preserved. +# +function(warn_if_ptx_arch_requested CUDA_ARCH_FLAGS) + foreach(_ARCH_FLAG ${CUDA_ARCH_FLAGS}) + if(_ARCH_FLAG MATCHES "code=.*compute_[0-9]+[af]?") + message(WARNING + "PTX code generation requested in CUDA architecture flags " + "(${_ARCH_FLAG}), but vLLM does not preserve global PTX requests " + "when normalizing per-source CUDA architectures. Remove '+PTX' from " + "TORCH_CUDA_ARCH_LIST or rely on vLLM's built-in per-kernel PTX " + "selection.") + return() + endif() + endforeach() +endfunction() + + # # Extract unique CUDA architectures from a list of compute capabilities codes in # the form `[]`, convert them to the form sort diff --git a/csrc/cpu/cpu_types_riscv_impl.hpp b/csrc/cpu/cpu_types_riscv_impl.hpp index 106bf44a2411..9a14d1c33866 100644 --- a/csrc/cpu/cpu_types_riscv_impl.hpp +++ b/csrc/cpu/cpu_types_riscv_impl.hpp @@ -427,7 +427,10 @@ struct FP32Vec8 : public Vec { explicit FP32Vec8(const float* ptr) : reg(RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, VEC_ELEM_NUM)) {}; explicit FP32Vec8(fixed_fp32x8_t data) : reg(data) {}; - explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {}; + // Not explicit: copy-initialisation (`auto v = pair.first`) in + // cpu_attn_vec.hpp requires a converting copy constructor, and every other + // CPU backend leaves this implicit. + FP32Vec8(const FP32Vec8& data) : reg(data.reg) {}; explicit FP32Vec8(const FP16Vec8& v) : reg(RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(v.reg, VEC_ELEM_NUM)) {}; explicit FP32Vec8(fixed_fp16x8_t v) @@ -628,7 +631,8 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(const FP32Vec8& data) : reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)( data.reg, data.reg)) {}; - explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; + // Not explicit: see FP32Vec8's copy constructor. + FP32Vec16(const FP32Vec16& data) : reg(data.reg) {}; explicit FP32Vec16(int64_t value, const FP32Vec16& lut) { // Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on // VLEN=128 / m4 on VLEN=256), which causes heavy register spilling. diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 04b696abf50f..c32a01ecd505 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -204,7 +204,6 @@ inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values)); const __m512i abs_mask = _mm512_set1_epi16(0x7FFF); - const __m512i zero = _mm512_setzero_si512(); // expand values to 16-bit integers __m512i x0 = _mm512_cvtepu8_epi16(a); @@ -214,17 +213,23 @@ inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i x0 = _mm512_permutexvar_epi16(x0, lut); x1 = _mm512_permutexvar_epi16(x1, lut); - // check for zeros - __mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ); - __mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ); - - // emulate bf16 mul with scale factor - x0 = _mm512_add_epi16(x0, s0); - x1 = _mm512_add_epi16(x1, s1); + // Emulate the bf16 multiply by the E8M0 scale as an integer add on the + // exponent field, forcing zeros to stay zero. + // + // vptestmw sets a lane's mask bit when (x & 0x7FFF) != 0, i.e. for the lanes + // that are not +0.0 or -0.0. That is the *complement* of the mask the old + // and/cmp/add/blend sequence built, which selected the zero lanes; the + // zero-masking form of vpaddw then writes 0 where the mask is clear, so the + // two agree lane by lane. That includes the LUT's -0.0 entry, which both + // versions turn into +0.0. + // + // 2 instructions per vector instead of 4, no extra ISA requirement: + // vptestmw is AVX512BW, which vpermw above already needs. + __mmask32 mask0 = _mm512_test_epi16_mask(x0, abs_mask); + __mmask32 mask1 = _mm512_test_epi16_mask(x1, abs_mask); - // blend with zero - x0 = _mm512_mask_blend_epi16(mask0, x0, zero); - x1 = _mm512_mask_blend_epi16(mask1, x1, zero); + x0 = _mm512_maskz_add_epi16(mask0, x0, s0); + x1 = _mm512_maskz_add_epi16(mask1, x1, s1); return std::make_tuple(__m512bh(x0), __m512bh(x1)); } diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh index b43b9b8447d6..3312c127a762 100644 --- a/csrc/libtorch_stable/cooperative_topk.cuh +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -484,7 +484,10 @@ __device__ void large_topk(const float* __restrict__ row_input, template __device__ void cooperative_topk_body(CooperativeTopKParams params) { const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x; - const auto sl = params.lengths[row]; + // Clamp at 0: `sl` is compared signed here but cast to uint32 below, so a + // negative length would otherwise emit indices 0..TopK-1 as valid instead + // of the -1 padding. + const int32_t sl = params.lengths[row] > 0 ? params.lengths[row] : 0; int32_t* out = params.output + row * TopK; const float* in = params.input + row * params.stride; diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index db4805dde0fd..2e2c4f7cf998 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -123,14 +123,10 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, int rope_elem_base = 0) { uint4 const v = *reinterpret_cast(src); if constexpr (FP8 || APPLY_ROPE) { -#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) - // _typeConvert is unavailable on pre-Ampere. Kimi K3 uses - // bf16 inputs, so discard unsupported conversion paths in those builds. - if constexpr (std::is_same_v) { + using Converter = vllm::_typeConvert; + if constexpr (!Converter::exists) { return; } else { -#endif - using Converter = vllm::_typeConvert; auto const* p = reinterpret_cast(&v); float f[kVecElems]; @@ -175,17 +171,15 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, } *reinterpret_cast(dst) = out; #else - uint8_t out[kVecElems]; + uint8_t out[kVecElems]; #pragma unroll - for (int i = 0; i < kVecElems; i++) { - float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max); - out[i] = rocm_cvt_float_to_fp8_e4m3(s); - } - *reinterpret_cast(dst) = *reinterpret_cast(out); + for (int i = 0; i < kVecElems; i++) { + float s = fminf(fmaxf(f[i] * scale_inv, -kFp8Max), kFp8Max); + out[i] = rocm_cvt_float_to_fp8_e4m3(s); + } + *reinterpret_cast(dst) = *reinterpret_cast(out); #endif -#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) } -#endif } else { *reinterpret_cast(dst) = v; } @@ -197,28 +191,32 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, template __device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst, const scalar_t* src) { -#ifndef USE_ROCM - uint4 const input = *reinterpret_cast(src); using Converter = vllm::_typeConvert; - auto const* input2 = - reinterpret_cast(&input); - uint2 output; - auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output); + if constexpr (!Converter::exists) { + return; + } else { +#ifndef USE_ROCM + uint4 const input = *reinterpret_cast(src); + auto const* input2 = + reinterpret_cast(&input); + uint2 output; + auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output); #pragma unroll - for (int i = 0; i < 4; ++i) { - if constexpr (std::is_same_v) { - output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2( - static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE, - __NV_E4M3); - } else { - output2[i] = __nv_cvt_halfraw2_to_fp8x2( - static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3); + for (int i = 0; i < 4; ++i) { + if constexpr (std::is_same_v) { + output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2( + static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE, + __NV_E4M3); + } else { + output2[i] = __nv_cvt_halfraw2_to_fp8x2( + static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3); + } } - } - *reinterpret_cast(dst) = output; + *reinterpret_cast(dst) = output; #else - copyChunk8(dst, src, 1.0f); + copyChunk8(dst, src, 1.0f); #endif + } } // Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192). @@ -360,12 +358,10 @@ __device__ __forceinline__ void writeDsMlaCache( scalar_t* row16 = reinterpret_cast(row); scalar_t* rope_dst = row16 + kKvLoraRank / 2 + 8 + laneId * 2; if constexpr (APPLY_ROPE) { -#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) - if constexpr (std::is_same_v) { + using Converter = vllm::_typeConvert; + if constexpr (!Converter::exists) { return; } else { -#endif - using Converter = vllm::_typeConvert; using packed_t = typename Converter::packed_hip_type; packed_t const src = *reinterpret_cast(pe + laneId * 2); float2 const xy = Converter::convert(src); @@ -374,9 +370,7 @@ __device__ __forceinline__ void writeDsMlaCache( static_cast(cos_sin[laneId + kQkRopeHeadDim / 2]); *reinterpret_cast(rope_dst) = Converter::convert( make_float2(xy.x * cos - xy.y * sin, xy.x * sin + xy.y * cos)); -#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) } -#endif } else { *reinterpret_cast(rope_dst) = *reinterpret_cast(pe + laneId * 2); diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu new file mode 100644 index 000000000000..31a16fcf9965 --- /dev/null +++ b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu @@ -0,0 +1,557 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * SPDX-FileCopyrightText: Copyright contributors to the vLLM project + */ + +#include +#include +#include +#include + +#include "../torch_utils.h" +#include "../../cuda_compat.h" + +namespace { + +template +__device__ __forceinline__ void cp_async_16b(StateT* smem_ptr, + const StateT* gmem_ptr) { + const uint32_t smem_addr = + static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "r"(smem_addr), "l"(gmem_ptr)); +} + +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::); +} + +__device__ __forceinline__ void cp_async_wait_all() { + asm volatile("cp.async.wait_all;\n" ::: "memory"); +} + +template +__device__ __forceinline__ void copy_state_chunk(StateT* shared_state, + const StateT* state, int chunk, + int thread, int threads) { + constexpr int kElementsPerCopy = 16 / sizeof(StateT); + constexpr int kCopiesPerChunk = ChunkV * DimK / kElementsPerCopy; + const int stage = chunk % Stages; + for (int copy = thread; copy < kCopiesPerChunk; copy += threads) { + const int element = copy * kElementsPerCopy; + cp_async_16b(shared_state + stage * ChunkV * DimK + element, + state + chunk * ChunkV * DimK + element); + } + cp_async_commit(); +} + +template +__device__ __forceinline__ float4 load_state4(const StateT* state); + +template <> +__device__ __forceinline__ float4 load_state4(const float* state) { + return *reinterpret_cast(state); +} + +template <> +__device__ __forceinline__ float4 +load_state4<__nv_bfloat16>(const __nv_bfloat16* state) { + const __nv_bfloat162 lo = *reinterpret_cast(state); + const __nv_bfloat162 hi = *reinterpret_cast(state + 2); + return make_float4(__bfloat162float(lo.x), __bfloat162float(lo.y), + __bfloat162float(hi.x), __bfloat162float(hi.y)); +} + +template +__device__ __forceinline__ void store_state4(StateT* state, float4 value); + +template <> +__device__ __forceinline__ void store_state4(float* state, + float4 value) { + *reinterpret_cast(state) = value; +} + +template <> +__device__ __forceinline__ void store_state4<__nv_bfloat16>( + __nv_bfloat16* state, float4 value) { + *reinterpret_cast<__nv_bfloat162*>(state) = + __floats2bfloat162_rn(value.x, value.y); + *reinterpret_cast<__nv_bfloat162*>(state + 2) = + __floats2bfloat162_rn(value.z, value.w); +} + +constexpr int kDimK = 128; +constexpr int kDimV = 128; +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kChunkV = 32; +constexpr int kNumChunks = kDimV / kChunkV; +constexpr int kRowsPerWarp = kChunkV / kWarps; +constexpr int kMaxMtpTokens = 8; +constexpr int kDtBiasFloat32 = 0; +constexpr int kDtBiasBFloat16 = 1; +constexpr int kDtBiasFloat16 = 2; + +struct GdnDecodeStrides { + int64_t mixed_row; + int64_t a_row; + int64_t b_row; + int64_t gate_row; + int64_t state_slot; +}; + +__device__ __forceinline__ float sigmoid_fast(float x) { + return 1.0f / (1.0f + __expf(-x)); +} + +__device__ __forceinline__ float silu_fast(float x) { + return x * sigmoid_fast(x); +} + +__device__ __forceinline__ float softplus_fast(float x) { + return x > 20.0f ? x : log1pf(__expf(x)); +} + +__device__ __forceinline__ float load_dt_bias(const void* dt_bias, int head, + int dt_bias_type) { + if (dt_bias_type == kDtBiasBFloat16) { + return __bfloat162float(static_cast(dt_bias)[head]); + } + if (dt_bias_type == kDtBiasFloat16) { + return __half2float(static_cast(dt_bias)[head]); + } + return static_cast(dt_bias)[head]; +} + +__device__ __forceinline__ float warp_reduce_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_xor_sync(0xffffffffu, value, offset); + } + return value; +} + +struct Sum2 { + float x; + float y; +}; + +__device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffffu, x, offset); + y += __shfl_xor_sync(0xffffffffu, y, offset); + } + return {x, y}; +} + +template +__global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( + const __nv_bfloat16* __restrict__ mixed_qkv, + const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b, + const float* __restrict__ a_log, const void* __restrict__ dt_bias, + const int* __restrict__ state_indices, const int* __restrict__ cu_seqlens, + const int* __restrict__ num_accepted_tokens, StateT* __restrict__ state, + const __nv_bfloat16* __restrict__ output_gate, + const void* __restrict__ norm_weight, __nv_bfloat16* __restrict__ out, + int H, int HV, int state_indices_width, int dt_bias_type, + bool norm_weight_is_bf16, float scale, float norm_eps, + GdnDecodeStrides strides) { + const int request = blockIdx.x; + const int value_head = blockIdx.y; + const int tid = threadIdx.x; + const int lane = tid & 31; + const int warp = tid >> 5; + const int bos = cu_seqlens[request]; + const int eos = cu_seqlens[request + 1]; + const int num_tokens = eos - bos; + if (num_tokens <= 0) { + return; + } + + const int accepted = num_accepted_tokens[request]; + const int source_slot = + accepted > 0 && accepted <= state_indices_width + ? state_indices[request * state_indices_width + accepted - 1] + : 0; + if (source_slot <= 0 || num_tokens > kMaxMtpTokens) { + for (int linear = tid; linear < num_tokens * kDimV; linear += kThreads) { + const int token = bos + linear / kDimV; + const int value = linear % kDimV; + const int64_t out_offset = + (static_cast(token) * HV + value_head) * kDimV + value; + out[out_offset] = __float2bfloat16(0.0f); + } + return; + } + + const int key_head = value_head / 8; + __shared__ StateT shared_state[2][kChunkV][kDimK]; + __shared__ float shared_q[kMaxMtpTokens][kDimK]; + __shared__ float shared_k[kMaxMtpTokens][kDimK]; + __shared__ __nv_bfloat16 shared_v[kMaxMtpTokens][kDimV]; + __shared__ __nv_bfloat16 shared_out[kMaxMtpTokens][kDimV]; + __shared__ float shared_decay[kMaxMtpTokens]; + __shared__ float shared_beta[kMaxMtpTokens]; + + StateT* source_state = + state + static_cast(source_slot) * strides.state_slot + + value_head * kDimV * kDimK; + copy_state_chunk(&shared_state[0][0][0], + source_state, 0, tid, kThreads); + + if (warp < num_tokens) { + const int t = warp; + const int token = bos + t; + const int64_t mixed_base = static_cast(token) * strides.mixed_row; + float q_values[4]; + float k_values[4]; + float q_square = 0.0f; + float k_square = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int dim = lane + i * 32; + q_values[i] = + __bfloat162float(mixed_qkv[mixed_base + key_head * kDimK + dim]); + k_values[i] = __bfloat162float( + mixed_qkv[mixed_base + H * kDimK + key_head * kDimK + dim]); + shared_v[t][dim] = + mixed_qkv[mixed_base + 2 * H * kDimK + value_head * kDimV + dim]; + q_square += q_values[i] * q_values[i]; + k_square += k_values[i] * k_values[i]; + } + const Sum2 qk_sums = warp_reduce_sum_pair(q_square, k_square); + const float q_scale = __shfl_sync( + 0xffffffffu, lane == 0 ? rsqrtf(qk_sums.x + 1.0e-6f) * scale : 0.0f, 0); + const float k_scale = __shfl_sync( + 0xffffffffu, lane == 0 ? rsqrtf(qk_sums.y + 1.0e-6f) : 0.0f, 0); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int dim = lane + i * 32; + shared_q[t][dim] = q_values[i] * q_scale; + shared_k[t][dim] = k_values[i] * k_scale; + } + if (lane == 0) { + const float a_value = __bfloat162float( + a[static_cast(token) * strides.a_row + value_head]); + const float b_value = __bfloat162float( + b[static_cast(token) * strides.b_row + value_head]); + const float g = -__expf(a_log[value_head]) * + softplus_fast(a_value + load_dt_bias(dt_bias, value_head, + dt_bias_type)); + shared_decay[t] = __expf(g); + shared_beta[t] = sigmoid_fast(b_value); + } + } + __syncthreads(); + + const int k_base = lane * 4; + int rows[kRowsPerWarp]; +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + rows[row] = warp + row * kWarps; + } + +#pragma unroll + for (int chunk = 0; chunk < kNumChunks; ++chunk) { + cp_async_wait_all(); + __syncthreads(); + if (chunk + 1 < kNumChunks) { + copy_state_chunk( + &shared_state[0][0][0], source_state, chunk + 1, tid, kThreads); + } + + float h[kRowsPerWarp][4]; +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const float4 state_value = + load_state4(&shared_state[chunk & 1][rows[row]][k_base]); + h[row][0] = state_value.x; + h[row][1] = state_value.y; + h[row][2] = state_value.z; + h[row][3] = state_value.w; + } + + for (int t = 0; t < num_tokens; ++t) { + const float4 q4 = *reinterpret_cast(&shared_q[t][k_base]); + const float4 k4 = *reinterpret_cast(&shared_k[t][k_base]); + const float q_values[4] = {q4.x, q4.y, q4.z, q4.w}; + const float k_values[4] = {k4.x, k4.y, k4.z, k4.w}; + + float dot_hk[kRowsPerWarp] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { +#pragma unroll + for (int i = 0; i < 4; ++i) { + h[row][i] *= shared_decay[t]; + dot_hk[row] += h[row][i] * k_values[i]; + } + } + const Sum2 dot_hk_01 = warp_reduce_sum_pair(dot_hk[0], dot_hk[1]); + const Sum2 dot_hk_23 = warp_reduce_sum_pair(dot_hk[2], dot_hk[3]); + const float reduced_hk[kRowsPerWarp] = {dot_hk_01.x, dot_hk_01.y, + dot_hk_23.x, dot_hk_23.y}; + + float dot_hq[kRowsPerWarp] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int value = chunk * kChunkV + rows[row]; + const float delta = + (__bfloat162float(shared_v[t][value]) - reduced_hk[row]) * + shared_beta[t]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + h[row][i] += k_values[i] * delta; + dot_hq[row] += h[row][i] * q_values[i]; + } + } + const Sum2 dot_hq_01 = warp_reduce_sum_pair(dot_hq[0], dot_hq[1]); + const Sum2 dot_hq_23 = warp_reduce_sum_pair(dot_hq[2], dot_hq[3]); + if (lane == 0) { + shared_out[t][chunk * kChunkV + rows[0]] = + __float2bfloat16(dot_hq_01.x); + shared_out[t][chunk * kChunkV + rows[1]] = + __float2bfloat16(dot_hq_01.y); + shared_out[t][chunk * kChunkV + rows[2]] = + __float2bfloat16(dot_hq_23.x); + shared_out[t][chunk * kChunkV + rows[3]] = + __float2bfloat16(dot_hq_23.y); + } + + const int destination_slot = + state_indices[request * state_indices_width + t]; + if (destination_slot > 0) { + StateT* destination_state = + state + + static_cast(destination_slot) * strides.state_slot + + value_head * kDimV * kDimK; +#pragma unroll + for (int row = 0; row < kRowsPerWarp; ++row) { + const int value = chunk * kChunkV + rows[row]; + store_state4(destination_state + value * kDimK + k_base, + make_float4(h[row][0], h[row][1], h[row][2], h[row][3])); + } + } + } + } + __syncthreads(); + + if (warp < num_tokens) { + const int t = warp; + float output_values[4]; + float sum_square = 0.0f; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int value = lane + i * 32; + output_values[i] = __bfloat162float(shared_out[t][value]); + sum_square += output_values[i] * output_values[i]; + } + sum_square = warp_reduce_sum(sum_square); + const float rstd = + rsqrtf(sum_square / static_cast(kDimV) + norm_eps); + const int token = bos + t; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int value = lane + i * 32; + const float gate = silu_fast(__bfloat162float( + output_gate[static_cast(token) * strides.gate_row + + value_head * kDimV + value])); + const float weight = + norm_weight_is_bf16 + ? __bfloat162float( + static_cast(norm_weight)[value]) + : static_cast(norm_weight)[value]; + const int64_t out_offset = + (static_cast(token) * HV + value_head) * kDimV + value; + out[out_offset] = + __float2bfloat16(output_values[i] * rstd * weight * gate); + } + } +} + +template +void launch_gdn_decode_post_conv_mtp( + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, + torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& num_accepted_tokens, + torch::stable::Tensor& state, torch::stable::Tensor const& norm_weight, + torch::stable::Tensor& out, const __nv_bfloat16* a, const __nv_bfloat16* b, + const __nv_bfloat16* output_gate, int num_key_heads, int num_value_heads, + double scale, double norm_eps, GdnDecodeStrides strides) { + using torch::headeronly::ScalarType; + + const auto dt_bias_scalar_type = dt_bias.scalar_type(); + const int dt_bias_type = + dt_bias_scalar_type == ScalarType::Float + ? kDtBiasFloat32 + : (dt_bias_scalar_type == ScalarType::BFloat16 ? kDtBiasBFloat16 + : kDtBiasFloat16); + torch::stable::accelerator::DeviceGuard const device_guard( + mixed_qkv.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(mixed_qkv.get_device_index()); + const int num_requests = static_cast(state_indices.size(0)); + const dim3 grid(num_requests, num_value_heads); + gdn_decode_post_conv_mtp_kernel<<>>( + static_cast(mixed_qkv.data_ptr()), a, b, + static_cast(a_log.data_ptr()), dt_bias.data_ptr(), + static_cast(state_indices.data_ptr()), + static_cast(cu_seqlens.data_ptr()), + static_cast(num_accepted_tokens.data_ptr()), + static_cast(state.data_ptr()), output_gate, + norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), + num_key_heads, num_value_heads, static_cast(state_indices.size(1)), + dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16, + static_cast(scale), static_cast(norm_eps), strides); + const cudaError_t error = cudaGetLastError(); + STD_TORCH_CHECK(error == cudaSuccess, + "GDN decode MTP post-conv kernel launch failed: ", + cudaGetErrorString(error)); +} + +} // namespace + +void fused_gdn_decode_post_conv_mtp( + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, + torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& num_accepted_tokens, + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, + double scale, double norm_eps) { + using torch::headeronly::ScalarType; + + STD_TORCH_CHECK( + mixed_qkv.is_cuda() && mixed_qkv.scalar_type() == ScalarType::BFloat16, + "mixed_qkv must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(a.is_cuda() && a.scalar_type() == ScalarType::BFloat16, + "a must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(b.is_cuda() && b.scalar_type() == ScalarType::BFloat16, + "b must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(a_log.is_cuda() && a_log.scalar_type() == ScalarType::Float, + "A_log must be a CUDA float32 tensor"); + const auto dt_bias_scalar_type = dt_bias.scalar_type(); + STD_TORCH_CHECK( + dt_bias.is_cuda() && (dt_bias_scalar_type == ScalarType::Float || + dt_bias_scalar_type == ScalarType::BFloat16 || + dt_bias_scalar_type == ScalarType::Half), + "dt_bias must be a CUDA float32, bfloat16, or float16 tensor"); + STD_TORCH_CHECK( + state_indices.is_cuda() && state_indices.scalar_type() == ScalarType::Int, + "state_indices must be a CUDA int32 tensor"); + STD_TORCH_CHECK( + cu_seqlens.is_cuda() && cu_seqlens.scalar_type() == ScalarType::Int, + "cu_seqlens must be a CUDA int32 tensor"); + STD_TORCH_CHECK(num_accepted_tokens.is_cuda() && + num_accepted_tokens.scalar_type() == ScalarType::Int, + "num_accepted_tokens must be a CUDA int32 tensor"); + const auto state_scalar_type = state.scalar_type(); + STD_TORCH_CHECK( + state.is_cuda() && (state_scalar_type == ScalarType::Float || + state_scalar_type == ScalarType::BFloat16), + "state must be a CUDA float32 or bfloat16 tensor"); + STD_TORCH_CHECK(output_gate.is_cuda() && + output_gate.scalar_type() == ScalarType::BFloat16, + "output_gate must be a CUDA bfloat16 tensor"); + STD_TORCH_CHECK(norm_weight.is_cuda() && + (norm_weight.scalar_type() == ScalarType::Float || + norm_weight.scalar_type() == ScalarType::BFloat16), + "norm_weight must be a CUDA float32 or bfloat16 tensor"); + STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16, + "out must be a CUDA bfloat16 tensor"); + + STD_TORCH_CHECK(mixed_qkv.dim() == 2, + "mixed_qkv must have shape [L, 2 * H * 128 + HV * 128]"); + const int num_tokens = static_cast(mixed_qkv.size(0)); + STD_TORCH_CHECK(num_tokens > 0, + "GDN decode MTP fusion requires at least one token"); + STD_TORCH_CHECK( + state.dim() == 4 && state.size(2) == kDimV && state.size(3) == kDimK, + "state must have shape [slots, HV, 128, 128]"); + const int num_value_heads = static_cast(state.size(1)); + const int64_t key_width = + mixed_qkv.size(1) - static_cast(num_value_heads) * kDimV; + STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0, + "mixed_qkv width is inconsistent with state"); + const int num_key_heads = static_cast(key_width / (2 * kDimK)); + STD_TORCH_CHECK(num_value_heads == 8 * num_key_heads, + "GDN decode MTP fusion requires HV/H=8"); + + STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 && + state_indices.size(1) > 0 && + state_indices.size(1) <= kMaxMtpTokens, + "state_indices must have shape [N, S] with 1 <= S <= 8"); + const int num_requests = static_cast(state_indices.size(0)); + STD_TORCH_CHECK( + cu_seqlens.dim() == 1 && cu_seqlens.numel() == num_requests + 1, + "cu_seqlens must have N + 1 elements"); + STD_TORCH_CHECK(num_accepted_tokens.dim() == 1 && + num_accepted_tokens.numel() == num_requests, + "num_accepted_tokens must have N elements"); + STD_TORCH_CHECK( + a.dim() == 2 && a.size(0) == num_tokens && a.size(1) == num_value_heads, + "a must have shape [L, HV]"); + STD_TORCH_CHECK( + b.dim() == 2 && b.size(0) == num_tokens && b.size(1) == num_value_heads, + "b must have shape [L, HV]"); + STD_TORCH_CHECK(a_log.is_contiguous() && a_log.numel() == num_value_heads, + "A_log must be contiguous with HV elements"); + STD_TORCH_CHECK(dt_bias.is_contiguous() && dt_bias.numel() == num_value_heads, + "dt_bias must be contiguous with HV elements"); + STD_TORCH_CHECK(state_indices.is_contiguous(), + "state_indices must be contiguous"); + STD_TORCH_CHECK(cu_seqlens.is_contiguous(), "cu_seqlens must be contiguous"); + STD_TORCH_CHECK(num_accepted_tokens.is_contiguous(), + "num_accepted_tokens must be contiguous"); + STD_TORCH_CHECK(output_gate.dim() == 3 && output_gate.size(0) == num_tokens && + output_gate.size(1) == num_value_heads && + output_gate.size(2) == kDimV, + "output_gate must have shape [L, HV, 128]"); + STD_TORCH_CHECK(norm_weight.is_contiguous() && norm_weight.numel() == kDimV, + "norm_weight must be contiguous with 128 elements"); + STD_TORCH_CHECK(out.dim() == 3 && out.size(0) == num_tokens && + out.size(1) == num_value_heads && out.size(2) == kDimV, + "out must have shape [L, HV, 128]"); + STD_TORCH_CHECK(mixed_qkv.stride(1) == 1, + "mixed_qkv channels must be contiguous"); + STD_TORCH_CHECK(a.stride(1) == 1 && b.stride(1) == 1, + "a and b heads must be contiguous"); + STD_TORCH_CHECK(state.stride(0) >= num_value_heads * kDimV * kDimK && + state.stride(1) == kDimV * kDimK && + state.stride(2) == kDimK && state.stride(3) == 1, + "state must have contiguous [HV, 128, 128] slot contents"); + const int state_elements_per_copy = + state_scalar_type == ScalarType::Float ? 4 : 8; + STD_TORCH_CHECK(reinterpret_cast(state.data_ptr()) % 16 == 0 && + state.stride(0) % state_elements_per_copy == 0, + "state slots must preserve 16-byte alignment"); + STD_TORCH_CHECK(output_gate.stride(2) == 1 && output_gate.stride(1) == kDimV, + "output_gate head rows must be contiguous"); + STD_TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); + STD_TORCH_CHECK(norm_eps >= 0.0, "norm_eps must be non-negative"); + + const GdnDecodeStrides strides{mixed_qkv.stride(0), a.stride(0), b.stride(0), + output_gate.stride(0), state.stride(0)}; + const auto* a_ptr = static_cast(a.data_ptr()); + const auto* b_ptr = static_cast(b.data_ptr()); + const auto* output_gate_ptr = + static_cast(output_gate.data_ptr()); + if (state_scalar_type == ScalarType::Float) { + launch_gdn_decode_post_conv_mtp( + mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, + num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, + output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, + strides); + } else { + launch_gdn_decode_post_conv_mtp<__nv_bfloat16>( + mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, + num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, + output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, + strides); + } +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 85c875402c10..5c595bfaf49c 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -400,6 +400,18 @@ void fused_kda_decode( torch::stable::Tensor& out, std::optional lower_bound, std::optional output_gate, std::optional norm_weight, double norm_eps); + +void fused_gdn_decode_post_conv_mtp( + torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_log, + torch::stable::Tensor const& dt_bias, + torch::stable::Tensor const& state_indices, + torch::stable::Tensor const& cu_seqlens, + torch::stable::Tensor const& num_accepted_tokens, + torch::stable::Tensor& state, torch::stable::Tensor const& output_gate, + torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out, + double scale, double norm_eps); + #endif #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 91717f53b231..9ae0fd4bcbaa 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -906,7 +906,26 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) uint32_t row_idx = group_id + iter * num_groups; if (row_idx >= params.num_rows) break; - const uint32_t seq_len = params.lengths[row_idx]; + // Clamp the row length before any decision is made on it. + // + // `lengths` is int32 and is consumed here as uint32, so a negative value + // (e.g. a padded decode slot whose per-token context length underflowed) + // would reinterpret as ~4e9 and sail past every threshold below. Any + // value beyond the row width would also read into the next row. + // + // Clamping to max_seq_len additionally keeps this per-row decision + // consistent with the `cta_in_group != 0` early exit above, which is + // taken from the host-side scalar: when max_seq_len <= RADIX_THRESHOLD + // the non-leader CTAs return immediately, so a leader that reached the + // cooperative radix path would wait on the inter-CTA barrier for peers + // that no longer exist and spin until the kernel is killed. + const int32_t raw_len = params.lengths[row_idx]; + const uint32_t row_bound = + params.stride < params.max_seq_len ? params.stride : params.max_seq_len; + const uint32_t non_negative_len = + raw_len > 0 ? static_cast(raw_len) : 0u; + const uint32_t seq_len = + non_negative_len < row_bound ? non_negative_len : row_bound; int32_t* row_output = params.output + row_idx * params.top_k; const float* row_input = params.input + row_idx * params.stride; diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 801aed03d910..28090ce5dcaf 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -534,6 +534,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? norm_weight=None, float norm_eps=1e-5) -> ()"); #endif +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.def( + "fused_gdn_decode_post_conv_mtp(" + "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, " + "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, " + "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, " + "float scale, float norm_eps=1e-5) -> ()"); +#endif + #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES ops.def( "kimi_k3_attn_res(" @@ -797,6 +806,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("fused_kda_decode", TORCH_BOX(&fused_kda_decode)); #endif +#ifdef VLLM_ENABLE_FUSED_GDN_DECODE + ops.impl("fused_gdn_decode_post_conv_mtp", + TORCH_BOX(&fused_gdn_decode_post_conv_mtp)); +#endif + #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res)); #endif diff --git a/docker/Dockerfile b/docker/Dockerfile index 0b68f89dde8e..804a3a4bb3ef 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -965,6 +965,9 @@ ENV HF_XET_HIGH_PERFORMANCE 1 # increase timeout for hf downloads (for testing) ENV HF_HUB_DOWNLOAD_TIMEOUT 60 +# Catch GPU<->CPU syncs in execute_model/sample_tokens +ENV VLLM_GPU_SYNC_CHECK=error + # Copy in the v1 package for testing (it isn't distributed yet) COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 4595daa56f3e..d3504d1a3dfb 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -397,7 +397,7 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ && rm -rf /var/lib/apt/lists/* RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \ + uv pip install --system meson meson-python pyyaml types-PyYAML \ auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0" RUN --mount=type=cache,target=/root/.cache/ccache \ @@ -820,6 +820,9 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac # Hide source under src/ so it won't shadow the installed package in tests. RUN mkdir src && mv vllm src/vllm +# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. +ENV VLLM_GPU_SYNC_CHECK=error + # ----------------------- # Final vLLM image FROM mori_base AS final diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 1b847d216742..31efca624151 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -129,6 +129,35 @@ CMD ["/bin/bash"] ######################### UCX + NIXL BUILD STAGE ######################### # Build UCX and NIXL in a dedicated stage so compiler/autotools layers are # never included in the final runtime image (mirrors ROCm's build_rixl stage). +######################### WHEEL BUILD STAGE ######################### +# Produces a standalone dist/*.whl artifact (used by the release pipeline's +# "Build wheel - XPU" step), independent of the runtime image below. +FROM vllm-base AS vllm-build + +ARG GIT_REPO_CHECK=0 + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \ + --mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \ + uv pip install grpcio-tools protobuf nanobind && \ + uv pip install -r /workspace/vllm/requirements/xpu.txt + +# Keep source-dependent layers near the end so frequent code-only changes +# don't invalidate heavy dependency layers. +COPY . . + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build /workspace/vllm/vllm-rs vllm/vllm-rs +COPY --from=rust-build /workspace/vllm/_rust_*.so vllm/ + +RUN --mount=type=bind,source=.git,target=.git \ + if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh; fi + +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=bind,source=.git,target=.git \ + VLLM_TARGET_DEVICE=xpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 + FROM vllm-base AS ucx-nixl-build ARG UCX_VERSION=v1.21.0-rc2 @@ -201,9 +230,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \ uv pip install grpcio-tools protobuf nanobind && \ uv pip install -r /workspace/vllm/requirements/xpu.txt && \ - uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt && \ - uv pip uninstall triton triton-xpu && \ - uv pip install triton-xpu==3.7.2 + uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt # Keep source-dependent layers near the end so frequent code-only changes # don't invalidate heavy dependency and UCX/NIXL layers. diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index 510bf7e2b1c3..be2a4cc4f389 100644 Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 476171f73226..2437392d9be6 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -205,7 +205,7 @@ vllm bench serve --port 9001 --save-result --save-detailed \ --endpoint /v1/completions \ --dataset-name custom \ --dataset-path \ - --custom-skip-chat-template \ + --skip-chat-template \ --num-prompts 80 \ --max-concurrency 1 \ --temperature=0.3 \ @@ -213,7 +213,7 @@ vllm bench serve --port 9001 --save-result --save-detailed \ --result-dir "./log/" ``` -You can skip applying chat template if your data already has it by using `--custom-skip-chat-template`. +You can skip applying chat template if your data already has it by using `--skip-chat-template`. #### Custom Audio Dataset diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 89acc6b7f5ae..25e7ecbdf1c7 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -71,6 +71,7 @@ uv pip install -e . --no-build-isolation For more details about installing from source and installing for other hardware, check out the [installation instructions](../getting_started/installation/README.md) for your hardware and head to the "Build wheel from source" section. For an optimized workflow when iterating on C++/CUDA kernels, see the [Incremental Compilation Workflow](./incremental_build.md) for recommendations. +For JIT kernel warmup conventions, see [JIT Kernel Warmup](./jit_kernel_warmup.md). !!! tip vLLM is compatible with Python versions 3.10 to 3.13. However, vLLM's default [Dockerfile](../../docker/Dockerfile) ships with Python 3.12 and tests in CI (except `mypy`) are run with Python 3.12. diff --git a/docs/contributing/jit_kernel_warmup.md b/docs/contributing/jit_kernel_warmup.md new file mode 100644 index 000000000000..4cf021279057 --- /dev/null +++ b/docs/contributing/jit_kernel_warmup.md @@ -0,0 +1,308 @@ +# JIT Kernel Warmup + +vLLM uses JIT-generated kernels from Triton, CuTeDSL, TileLang, and other backends. This contract makes their required specializations available during startup, before the first request, by warming the kernel's **compile-key space** without dummy runtime launches or real tensor allocation. + +Use it when adding a warmable JIT kernel or migrating an existing warmup path. + +## In This Guide + +- [1. Quickstart](#1-quickstart): for contributors adding or migrating a warmable kernel. +- [2. Search-Space Reference](#2-search-space-reference): additional details regarding warmup input expansion and traced dispatch rules. + +## 1. Quickstart + +Each warmable kernel defines its compile-key mapping and compile-only entry point beside its normal runtime implementation. The startup registry then warms only the wrappers selected by the current engine configuration. + +### Define the Kernel Wrapper + +Here, a **kernel wrapper** (or just **wrapper**) is an instance of a concrete `VllmJitKernel` subclass. + +Expose one wrapper near the kernel's normal runtime entry point. Prefer this shape: + +```python +class MyKernel(VllmJitKernel["MyKernel.CompileKey"]): + + @dataclass(frozen=True) + class CompileKey: + ... + + @staticmethod + def kernel(...): + ... + + def dispatch(self, ...) -> CompileKey: + return self.CompileKey(...) + + def get_warmup_keys(self, ...) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)(...) + + def compile(self, compile_key: CompileKey) -> None: + ... + + def __call__(self, ...): + return self.kernel(...) + + +MY_KERNEL = MyKernel() +``` + +`CompileKey`, `dispatch(...)`, and `get_warmup_keys(...)` are backend-agnostic. Backend-specific behavior belongs in `kernel(...)`, `compile(...)`, and `__call__(...)`. + +The module-level singleton should be used by warmup and by the runtime call path. This keeps dispatch behavior shared instead of duplicated. + +`VllmJitKernel.warmup(...)` compiles every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it. + +### Choose Compile-Key Fields + +`CompileKey` must be frozen and hashable. Include only fields on which the backend specializes, such as tile sizes, head dimensions, dtypes, pointer alignment classes, or backend selectors; exclude runtime-only values. When unsure, inspect the backend cache key, specialization arguments, or verbose JIT-monitor output. + +### Generate Warmup Keys + +Use `_trace_dispatch(self.dispatch)` to describe representative inputs. The tracer maps them through the same specialization logic and deduplicates equal keys: + +```python +def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + ) +``` + +Use independent ranges or alternatives for cartesian products, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. The complete syntax is documented in [Search-Space Reference](#2-search-space-reference). + +### Compile Without Launching + +`compile(compile_key)` means "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss. + +`compile(...)` should not launch a real inference workload or allocate real tensors. Each DSL should expose fake tensor/spec descriptors suitable for compilation only. + +### Register the Selected Wrapper + +Register the wrapper where the runtime implementation is selected: + +```python +MY_KERNEL.register_warmup() +``` + +Registration records metadata only. It does not compile or launch the kernel. Repeated registrations from equivalent layers are allowed and deduplicated later. + +### Review Checklist + +- Warm actual compile keys rather than representative non-key inputs. +- Keep specialization mapping in `dispatch(...)` instead of duplicating it in warmup code. +- Use fake tensors or backend compile-only descriptors; never perform a dummy runtime launch. +- Keep registration metadata-only so model construction remains cheap and side-effect free. +- Compile registered kernels only through `kernel_warmup()`. +- Keep runtime execution and startup compilation separate and easy to review. +- Use one module-level wrapper instance for registration and runtime calls. + +## 2. Search-Space Reference + +### How Tracing Works + +`_trace_dispatch(...)` expands the inputs declared by `get_warmup_keys(...)`. Each concrete combination becomes a `dispatch_values` mapping from input names to selected values. `_when` may reject that mapping; otherwise the tracer evaluates `dispatch(...)` to construct one `CompileKey`. Equal keys are deduplicated after all combinations are evaluated. + +One call to `dispatch(...)` returns one key, but many input points may map to the same key. Prefer this traced mapping over manually reconstructing keys in warmup code; `dispatch(...)` should express the same specialization logic used by the runtime path. + +### Define Input Spaces + +Use ranges and alternatives for independent axes, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. + +#### Integer Ranges + +Use `WarmupIntRange` for integer ranges: + +```python +return self._trace_dispatch(self.dispatch)( + num_prefills=WarmupIntRange(1, max_prefills + 1), +) +``` + +`WarmupIntRange(start, stop, step)` follows Python `range(...)` semantics: `start` is inclusive, `stop` is exclusive, and `step` defaults to 1. + +For non-linear integer sequences, use `advance` to provide the action that computes each next value: + +```python +return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange( + 1, + max_tokens + 1, + advance=lambda value: next_power_of_2(value) + 1, + ), +) +``` + +This is useful for traversing specialization boundaries without enumerating every integer. `advance` cannot be combined with a non-default `step`, and it must return a value greater than its input so expansion always makes forward progress. + +#### Independent Alternatives + +Use tuples or lists for independent alternatives. Multiple expanded inputs form a cartesian product: + +```python +return self._trace_dispatch(self.dispatch)( + query_slice_start=WarmupIntRange(0, 2), + query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), + COMPRESS_RATIO=list(compress_ratios), +) +``` + +#### Coupled Inputs + +Use `zip_inputs(...)` when values must vary together row-by-row: + +```python +WARMUP_INPUTS = zip_inputs( + dict(compress_ratio=1, topk=0, topk_width=512), + dict(compress_ratio=4, topk=512, topk_width=512), +) + + +return self._trace_dispatch(self.dispatch)( + WARMUP_INPUTS, + WINDOW_SIZE=window_size, +) +``` + +Multiple `zip_inputs(...)` groups may be passed as positional arguments. The tracer forms the cartesian product across groups while preserving row-wise coupling inside each group. + +Every row in a `zip_inputs(...)` group must use the same string keys. A `zip_inputs(...)` group cannot specify a field that is also specified as a keyword input to `_trace_dispatch(...)`. + +#### Conditional Filtering + +Use `_when=...` to filter generated input points before they are passed to `dispatch(...)`. This is useful when independent ranges contain invalid combinations, but the validity rule belongs with the kernel warmup definition. + +```python +def _is_valid_warmup_input( + self, + *, + query_len: int, + num_reqs: int, + max_num_batched_tokens: int, +) -> bool: + return query_len + num_reqs - 1 <= max_num_batched_tokens + + +return self._trace_dispatch(self.dispatch)( + query_len=WarmupIntRange(1, max_tokens + 1), + num_reqs=WarmupIntRange(1, max_reqs + 1), + max_num_batched_tokens=max_tokens, + _when=self._is_valid_warmup_input, +) +``` + +`_when` accepts a function, bound method, or lambda and supports the same AST subset as `dispatch(...)`, including local assignments in function predicates. + +The predicate is evaluated on the expanded warmup inputs. If it returns `False`, that input point is skipped and no `CompileKey` is produced for it. + +### Write Dispatch Rules + +#### Local Assignments + +The traced body may contain local assignments, optionally annotated, followed by one `return self.CompileKey(...)` call. Local assignments let a kernel name intermediate specialization choices once and reuse them across fields: + +```python +def dispatch( + self, + *, + num_tokens: int, + vectorized: bool, +) -> CompileKey: + block_size = next_power_of_2(num_tokens) + return self.CompileKey( + BLOCK_SIZE=block_size, + VECTOR_WIDTH=4 if vectorized and block_size >= 4 else 1, + ) +``` + +#### Supported Expressions + +The evaluator supports these expressions inside local assignments and `CompileKey(...)` fields: + +| Feature | What It Allows | +| --- | --- | +| Names | Read dispatch inputs, local assignments, defaults, and module globals. | +| Constants | Use literals such as integers, strings, booleans, and `None`. | +| Attributes | Read structured values such as `cfg.block_size` or `mla_dims.v_head_dim`. | +| Subscriptions | Read sequence positions or mapping values such as `config[0]` and `config["block_size"]`. | +| Tuple/list literals | Build shapes, strides, and other small structured fields. | +| Conditional expressions | Select a field with `x if condition else y`. | +| Boolean expressions | Combine predicates with `and`, `or`, and `not`. | +| Comparisons | Use `==`, `!=`, `<`, `<=`, `>`, `>=`, `in`, `not in`, `is`, and `is not`. | +| Arithmetic | Use `+`, `-`, `*`, `//`, `%`, and `**`. | +| Unary minus | Build negative sentinel values or signed descriptors. | +| Helper calls | Call helpers with positional and explicit keyword arguments. | + +Python builtins such as `min(...)`, `max(...)`, and `len(...)` are resolved unless the name is overridden locally or globally. + +#### Helper Calls + +Helpers are useful for small specialization rules: + +```python +def dispatch(self, *, num_tokens: int, block_size: int) -> CompileKey: + return self.CompileKey( + PADDED_TOKENS=round_up(num_tokens, multiple=block_size), + ) +``` + +`_trace_dispatch(...)` does not inspect helper bodies. It evaluates the call arguments and invokes the helper as ordinary Python, so control flow inside that helper is outside the AST interpreter's scope. Keep helpers deterministic and side-effect free. + +#### Direct Keyword Forwarding + +For many direct pass-through fields, the dispatch `**kwargs` parameter may be unpacked into `CompileKey(...)`: + +```python +def dispatch( + self, + *, + num_tokens: int, + **compile_key_fields: int, +) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + block_size=next_power_of_2(num_tokens), + ) +``` + +Unmatched dispatch arguments become compile-key fields and warmup inputs. Keep transformed inputs named and explicit. The unpacking must use the dispatch method's own `**kwargs` parameter directly and exactly once; arbitrary mappings, repeated unpacking, and helper-call `**kwargs` are rejected. The fully explicit form remains supported and is often clearer for non-trivial mappings. + +#### Unsupported Syntax + +Conditional expressions (`x if condition else y`) are supported, but statement-level `if` blocks are not supported directly inside traced `dispatch(...)` or `_when` bodies. The tracer expects a straight-line sequence of local assignments followed by one return expression. Small, pure helpers called by traced expressions execute as normal Python with concrete values and may use ordinary control flow, including `if` blocks. Do not put loops, mutation, side effects, or backend imports directly inside traced functions. Put environment and model gating in `get_warmup_keys(...)` or the outer warmup entry point. + +### Compile-Key Deduplication + +`_trace_dispatch(...)` deduplicates the resulting keys while preserving order. This is important when many runtime-like inputs map to the same static bucket. + +For example, this warmup range expands every token count, but the compile key only depends on the power-of-two bucket: + +```python +def dispatch( + self, + *, + num_tokens: int, +) -> CompileKey: + return self.CompileKey( + BLOCK_SIZE=next_power_of_2(num_tokens), + ) + + +def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + return self._trace_dispatch(self.dispatch)( + num_tokens=WarmupIntRange(1, max_tokens + 1), + ) +``` + +For `max_tokens == 8`, the expanded inputs are `1, 2, 3, 4, 5, 6, 7, 8`, but the returned keys are: + +```python +[ + CompileKey(BLOCK_SIZE=1), + CompileKey(BLOCK_SIZE=2), + CompileKey(BLOCK_SIZE=4), + CompileKey(BLOCK_SIZE=8), +] +``` + +Deduplication happens after `dispatch(...)` is evaluated, so the warmup system removes duplicate compile keys, not duplicate input values. `CompileKey` must be hashable for this to work; using `@dataclass(frozen=True)` is the standard pattern. diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index f171c42e808b..0fbb6801e44e 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -7,6 +7,59 @@ For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tili !!! note Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md). +## Compatibility Matrix + +!!! note + The symbols used below have the following meanings: + + - ✅ = Full compatibility + - 🟠 = Partial compatibility + - ❌ = No compatibility + - ❔ = Unknown or TBD + +### Model x Feature + +| Architecture | Models | CG for Image | CG for Video | Multi-Path Graph | +| ------------ | ------ | ------------ | ------------ | --------------- | +| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | +| `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | +| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | +| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | +| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | +| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | +| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | +| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | +| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | +| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ | + +### Model x Hardware + +| Architecture | NV Blackwell | NV Ampere | AMD MI300X | AMD MI350X / MI355X | +| ------------ | ---------------- | ------------- | -------------- | --------------------- | +| `DeepseekOCRForCausalLM` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Ernie4_5_VLMoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Gemma3ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Glm4vForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Gemma4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `InternVLChatModel` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `KimiVLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Llama4ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen2VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen2_5_VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen3_5ForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Qwen3_5MoeForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | +| `Step3VLForConditionalGeneration` | ✅︎ | ✅︎ | ❔ | ✅︎ | + +!!! note + Encoder CUDA Graph has currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. + For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested. + Encoder CUDA Graph has also been tested with AMD MI350X (gfx950) used `--mm-encoder-attn-backend=FLASH_ATTN` (the ROCm default). + ## Motivation Vision encoder inference incurs CUDA kernel launch overhead on the host side. The overhead is more significant when the batch size is small or image size is small. @@ -113,29 +166,6 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra !!! note The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager. -**Supported models:** - -| Architecture | Models | CG for Image | CG for Video | Dual-Path Graph | -| ------------ | ------ | ------------ | ------------ | --------------- | -| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ | -| `Ernie4_5_VLMoeForConditionalGeneration` | `ERNIE-4.5-VL` | ✅︎ | ❌︎ | ❌︎ | -| `Gemma3ForConditionalGeneration` | `Gemma3` | ✅︎ | ❌︎ | ❌︎ | -| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ | -| `Gemma4ForConditionalGeneration` | `Gemma-4` | ✅︎ | ✅︎ | ❌︎ | -| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ | -| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ | -| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ | -| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | -| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ | - -!!! note - Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. - For Qwen2-VL and Qwen2.5-VL only FA2 and FA3 has been tested. - ## Configuration Four fields in `CompilationConfig` control encoder CUDA Graphs: @@ -228,7 +258,7 @@ model = vllm.LLM( ) ``` -## About the Performance +## Benchmark Results The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details. diff --git a/docs/features/nixl_connector_compatibility.md b/docs/features/nixl_connector_compatibility.md index 7ad839b88aab..36b8876c694f 100644 --- a/docs/features/nixl_connector_compatibility.md +++ b/docs/features/nixl_connector_compatibility.md @@ -82,6 +82,8 @@ By default, a **compatibility hash** is checked during handshake. P and D instan - Attention backend - KV cache dtype (`cache_dtype`) - EAGLE/MTP-style speculative method and draft-model configuration +- NIXL transfer mode (push vs pull) — a push (WRITE) connector and a pull (READ) + connector use incompatible transfer protocols and must never be paired !!! warning Disable the hash check with `--kv-transfer-config '{"kv_connector_extra_config": {"enforce_handshake_compat": false}}'` at your own risk. diff --git a/docs/features/quantization/b12x.md b/docs/features/quantization/b12x.md new file mode 100644 index 000000000000..534bc6eb608b --- /dev/null +++ b/docs/features/quantization/b12x.md @@ -0,0 +1,24 @@ +# B12X Linear Backend + +[B12X](https://pypi.org/project/b12x/) provides optional CUDA kernels for +NVIDIA SM120 and SM121 GPUs. Install the dependency with: + +```bash +uv pip install "vllm[b12x]" +``` + +B12X participates in automatic kernel selection after established optimized +backends and before emulation. Select it explicitly with: + +```bash +vllm serve --linear-backend b12x +``` + +## Supported Configurations + +| Backend | Supported configurations | +| ------- | ------------------------ | +| Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, and MXFP4 | + +Dense W4A16 layers are not handled by B12X and continue to use another +compatible backend such as Marlin. diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 8f3ac9fa13be..7858fff2d4a1 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -86,8 +86,9 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and | `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. | | `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. | | `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. | -| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. | -| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. | +| `rejection_sample_method` | `string` | `standard` | `standard`, `synthetic`, or `block`. | +| `synthetic_acceptance_rates` | `list[float]` | `None` | Per-position unconditional acceptance rates for `synthetic` rejection sampling. Each entry in `[0, 1]`; length must equal `num_speculative_tokens`; must be non-increasing. | +| `synthetic_acceptance_length` | `float` | `None` | Target mean acceptance length for `synthetic`; in `[1, num_speculative_tokens + 1]`. Mutually exclusive with `synthetic_acceptance_rates`. | | `use_heterogeneous_vocab` | `boolean` | `false` | Allow draft and target models with different vocabularies. Builds a token-level intersection at initialisation and constrains draft logits to shared tokens only. Only compatible with `method=draft_model`. Probabilistic draft sampling (`draft_sample_method='probabilistic'`) is not yet supported when this option is enabled. | !!! note diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 4b82178a4dca..a417bf4af7e1 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -156,6 +156,9 @@ cd vllm uv pip install -e . --torch-backend=auto ``` +!!! note "CUDA Architecture & PTX Flags" + vLLM normalizes CUDA architectures on a per-source basis to optimize build times and wheel sizes. Global `+PTX` requests in `TORCH_CUDA_ARCH_LIST` (e.g., `TORCH_CUDA_ARCH_LIST="8.0+PTX"`) are ignored for general extension targets; vLLM generates PTX only for specific internal kernels that require it. + !!! tip Building from source requires a lot of compilation. If you are building from source repeatedly, it's more efficient to cache the compilation results. diff --git a/docs/getting_started/installation/gpu.rocm.inc.md b/docs/getting_started/installation/gpu.rocm.inc.md index 59c9723e666b..ad5af7994054 100644 --- a/docs/getting_started/installation/gpu.rocm.inc.md +++ b/docs/getting_started/installation/gpu.rocm.inc.md @@ -154,6 +154,77 @@ uv pip install vllm==${VLLM_VERSION} \ --8<-- [end:pre-built-wheels] --8<-- [start:build-wheel-from-source] +#### Set up using Python-only build (without compilation) {#python-only-build} + +If you only need to change Python code, you can build and install vLLM without +compilation. Changes you make to the code will be reflected when you run vLLM: + +```bash +git clone https://github.com/vllm-project/vllm.git +cd vllm +VLLM_USE_PRECOMPILED=1 python3 setup.py develop +``` + +This command will do the following: + +1. Look for the current branch in your vLLM clone. +1. Identify the corresponding base commit in the main branch. +1. Detect the ROCm version in your environment and select the matching wheel + variant. +1. Download the pre-built wheel of the base commit. +1. Use its compiled libraries and `vllm-rs` binary in the installation. + +!!! note + 1. If you change C++, HIP, or kernel code, you cannot use Python-only build; + otherwise you may see an import error about a library not being found or + an undefined symbol. + 2. If you rebase your development branch, it is recommended to uninstall + vLLM and re-run the above command to make sure your libraries are up to + date. + +!!! tip "Rebuilding the Rust frontend" +If you need to recompile the `vllm-rs` Rust frontend binary, you can rebuild and +install it without re-running the full installation: + + ```bash + ./build_rust.sh # release build + ./build_rust.sh --debug # faster build for development + ``` + + This will install the required Rust toolchain if needed, build the binary, + and place it in `vllm/vllm-rs`. + +If you see an error about a wheel not being found, the wheel for your base +commit and ROCm patch version might not be available. Check the available +variants under `https://wheels.vllm.ai/rocm//`. For example, ROCm 7.2.1 +uses the `rocm721` variant. + +There are more environment variables to control the behavior of Python-only +build: + +- `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file + path of a pre-compiled wheel to use. All other logic to find the wheel will be + skipped. +- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the full commit hash used to + download the pre-compiled wheel. +- `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the ROCm variant subdirectory, e.g., + `rocm700` or `rocm721`. If not specified, the variant is auto-detected based + on your system's ROCm version. An explicitly specified variant must match the + detected environment. + +You can find more information about vLLM's wheels in +[Install the latest code](#install-the-latest-code). + +!!! note + There is a possibility that your source code may have a different commit ID + compared to the vLLM wheel, which could potentially lead to unknown errors. + It is recommended to use the same commit ID for the source code as the vLLM + wheel you have installed. Please refer to + [Install the latest code](#install-the-latest-code) for instructions on how + to install a specified wheel. + +#### Full build (with compilation) {#full-build} + !!! tip - If you found that the following installation step does not work for you, please refer to [docker/Dockerfile.rocm_base](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.rocm_base). Dockerfile is a form of installation steps. diff --git a/docs/getting_started/installation/gpu.xpu.inc.md b/docs/getting_started/installation/gpu.xpu.inc.md index ef207c8d83d0..84dc69ca84c5 100644 --- a/docs/getting_started/installation/gpu.xpu.inc.md +++ b/docs/getting_started/installation/gpu.xpu.inc.md @@ -20,7 +20,26 @@ There is no extra information on creating a new Python environment for this devi --8<-- [end:set-up-using-python] --8<-- [start:pre-built-wheels] -Currently, there are no pre-built XPU wheels. +Pre-built vLLM XPU wheels are published to `wheels.vllm.ai`. Each XPU wheel +index also contains the `triton==3.7.2+xpu` shim described below. PyTorch XPU +packages are served from the PyTorch XPU index, so both index URLs are needed. + +#### Install the latest code + +To install the wheel built from the latest main branch: + +```bash +uv pip install vllm --extra-index-url https://wheels.vllm.ai/nightly/xpu --extra-index-url https://download.pytorch.org/whl/xpu --index-strategy unsafe-best-match +``` + +#### Install specific revisions + +If you want to access the wheels for previous commits (e.g. to bisect the behavior change, performance regression), you can specify the commit hash in the URL: + +```bash +export VLLM_COMMIT=730bd35378bf2a5b56b6d3a45be28b3092d26519 # use full commit hash from the main branch +uv pip install vllm --extra-index-url https://wheels.vllm.ai/${VLLM_COMMIT}/xpu --extra-index-url https://download.pytorch.org/whl/xpu --index-strategy unsafe-best-match +``` --8<-- [end:pre-built-wheels] --8<-- [start:build-wheel-from-source] @@ -36,25 +55,24 @@ pip install --upgrade pip pip install -v -r requirements/xpu.txt ``` -- Then, install the correct Triton package for Intel XPU. - - The default `triton` package (for NVIDIA GPUs) may be installed as a transitive dependency (e.g., via `xgrammar`). For Intel XPU, you must replace it with `triton-xpu`: - - ```bash - pip uninstall -y triton triton-xpu - pip install triton-xpu==3.7.2 --extra-index-url https://download.pytorch.org/whl/xpu - ``` - - !!! note - - `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues. - - For torch 2.13 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.2`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu). - -- Finally, build and install vLLM XPU backend: +- Then, install vLLM XPU backend: ```bash VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v ``` +!!! note + `requirements/xpu.txt` pins `triton==3.7.2+xpu`, a compatibility shim + hosted on `https://wheels.vllm.ai/xpu` that transparently resolves to + the real Intel XPU implementation (`triton-xpu`). This exists because + some transitive dependencies (e.g. `xgrammar`) unconditionally + require a distribution literally named `triton`, which otherwise + resolves to the NVIDIA-only PyPI `triton` package on XPU and can + cause correctness or runtime issues. No manual uninstall/reinstall of + `triton`/`triton-xpu` is needed; both `pip install` and `uv pip + install --index-strategy unsafe-best-match` resolve the correct + package automatically. + --8<-- [end:build-wheel-from-source] --8<-- [start:pre-built-images] diff --git a/docs/models/pooling_models/specific_models.md b/docs/models/pooling_models/specific_models.md index 8753f1fd07c3..0c7b0a1bf3fc 100644 --- a/docs/models/pooling_models/specific_models.md +++ b/docs/models/pooling_models/specific_models.md @@ -365,36 +365,99 @@ curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{ ## BAAI/bge-m3 -The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json` -the architecture is declared as `XLMRobertaModel`, which makes `vLLM` load it as a vanilla ROBERTA model without the -extra weights. To load the full model weights, override its architecture like this: +`BAAI/bge-m3` supports dense retrieval, lexical matching, and ColBERT-style +multi-vector retrieval. Its `config.json` declares `XLMRobertaModel`, so vLLM +otherwise loads it as a vanilla RoBERTa model without the extra sparse and +ColBERT weights. The examples below therefore override the architecture with +`BgeM3EmbeddingModel`. + +The three retrieval modes map to concrete pooling tasks as follows: + +| Retrieval mode | Pooling task | Output | +| -------------- | ------------ | ------ | +| Dense | `embed` | One embedding vector per input | +| Lexical/sparse | `token_classify` | One scalar weight per non-special token | +| ColBERT multi-vector | `token_embed` | One embedding vector per non-special token | + +Serve one concrete mode by selecting its task at load time: + +```shell +vllm serve BAAI/bge-m3 \ + --runner pooling \ + --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \ + --pooler-config.task +``` + +For dense embeddings, replace `` with `embed` and use the Embeddings API: ```shell -vllm serve BAAI/bge-m3 --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' +curl -s http://localhost:8000/v1/embeddings \ + -H "Content-Type: application/json" -d '{ + "model": "BAAI/bge-m3", + "input": ["What is BGE M3?", "Definition of BM25"] + }' ``` -Then you obtain the sparse embeddings like this: +For lexical weights, replace `` with `token_classify` and use the +Pooling API: ```shell curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ - "model": "BAAI/bge-m3", - "task": "token_classify", - "input": ["What is BGE M3?", "Definition of BM25"] + "model": "BAAI/bge-m3", + "task": "token_classify", + "input": ["What is BGE M3?", "Definition of BM25"] }' ``` Due to limitations in the output schema, the output consists of a list of -token scores for each token for each input. This means that you'll have to call -`/tokenize` as well to be able to pair tokens with scores. -Refer to the tests in `tests/models/language/pooling/test_bge_m3.py` to see how -to do that. +token scores for each input. Call `/tokenize` as well to pair token IDs with +their scores. See +[`test_bge_m3.py`](../../../tests/models/language/pooling/test_bge_m3.py) for a +complete example that also combines repeated token IDs. + +For ColBERT vectors, replace `` with `token_embed` and use the Pooling API: + +```shell +curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ + "model": "BAAI/bge-m3", + "task": "token_embed", + "input": ["What is BGE M3?", "Definition of BM25"] +}' +``` + +### Dense and sparse output through an IO processor plugin + +The source tree includes a reference +[BGE-M3 IO processor plugin](../../../tests/plugins/bge_m3_sparse_plugin) that +formats dense embeddings, sparse token weights, or both in one response. From +a source checkout, install it in the vLLM environment and load it as follows: -You can obtain the colbert embeddings like this: +```shell +uv pip install ./tests/plugins/bge_m3_sparse_plugin + +vllm serve BAAI/bge-m3 \ + --runner pooling \ + --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \ + --io-processor-plugin bge_m3_sparse_plugin +``` + +The plugin selects the internal `embed&token_classify` task so the model +computes dense and lexical outputs together. Public requests must use task +`plugin` and put the plugin-specific fields under `data`: ```shell curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{ - "model": "BAAI/bge-m3", - "task": "token_embed", - "input": ["What is BGE M3?", "Definition of BM25"] + "model": "BAAI/bge-m3", + "task": "plugin", + "data": { + "input": ["What is BGE M3?", "Definition of BM25"], + "embed_task": "dense&sparse", + "return_tokens": true + } }' ``` + +`embed_task` accepts `dense`, `sparse`, or `dense&sparse`. The combined +`embed&token_classify` task is an internal execution contract for this plugin, +not a generic Pooling API response format. Without the plugin, select one of +the three concrete tasks above. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 64c7c7062ab9..f9d55cb97e7a 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -289,6 +289,21 @@ os.environ["http_proxy"] = "http://your.proxy.server:port" os.environ["https_proxy"] = "http://your.proxy.server:port" ``` +### MatrixHub + +[MatrixHub](https://github.com/matrixhub-ai/matrixhub) is a self-hosted model registry and distribution layer that caches models from upstream hubs and serves them over a Hugging Face-compatible API inside your own network. + +Since the API is Hugging Face-compatible, you only need to point `HF_ENDPOINT` at your MatrixHub instance: + +```shell +export HF_ENDPOINT="http://" +vllm serve Qwen/Qwen3-0.6B +``` + +vLLM then downloads model weights from MatrixHub over the internal network instead of the public Hugging Face Hub, which is useful for air-gapped clusters and for avoiding repeated downloads across nodes. + +See the [MatrixHub guide for vLLM](https://matrixhub.ai/docs/guides/use-with-vllm/) for an end-to-end walkthrough, including Docker and Kubernetes deployment examples. + ### ModelScope To use models from [ModelScope](https://www.modelscope.cn) instead of Hugging Face Hub, set an environment variable: @@ -523,7 +538,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + IE+ | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ | | `Gemma3ForConditionalGeneration` | Gemma 3 | T + IE+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | -| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I+ + V + A* | `google/gemma-4-E2B-it`, etc. | | ✅︎ | +| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I+ + V + A* | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ | | `Gemma4UnifiedForConditionalGeneration` | Gemma 4 Unified | T + I+ + V + A | `google/gemma-4-12B-it`, etc. | | ✅︎ | | `GLM4VForCausalLM`^ | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | @@ -572,6 +587,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MossAudioModel` | MOSS-Audio | T + A+ | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ | | `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` | ✅︎ | ✅︎ | | `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` | ✅︎ | ✅︎ | @@ -619,6 +635,13 @@ Some models are supported only via the [Transformers modeling backend](#transfor * Only specific variants of the model support this modality (see notes below).
Q `Qwen*-VL` officially uses `qwen_vl_utils` for image preprocessing, while vLLM uses `transformers`' `video_processing_qwen*`, which leads to slightly different results compared to the official Hugging Face repository examples. +!!! note + For `Dots3NoteForCausalLM`, the vision and audio towers are only loaded when the + corresponding modality is enabled via `--limit-mm-per-prompt`. Video inputs are + decoded into frames and audio, so they require both towers. The checkpoint also + ships one MTP layer, enabled with + `--speculative-config '{"method":"mtp","num_speculative_tokens":1}'`. + !!! note `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its MobileNet-v5 vision backbone. @@ -656,6 +679,13 @@ Some models are supported only via the [Transformers modeling backend](#transfor coordinate decoding and are not exposed by this vLLM implementation. See [Moondream3 prompt recipes](../features/multimodal_inputs.md#moondream3-prompt-recipes). +!!! note + Both Muse Glimmer architecture names map to the same vLLM implementation: + checkpoints with a vision config accept image and video inputs, while + vision-less checkpoints run as a text-only model. Speculative decoding uses the + `meta-models/Muse-Glimmer-30B-assistant` checkpoint, which vLLM serves through + its [DFlash](../features/speculative_decoding/README.md) path. + !!! 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/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index d8e6c7e4099b..b394e41415aa 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -158,6 +158,8 @@ When running vLLM as an HTTP server, the following endpoints are available for w !!! note The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set. +The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, or sparse-NCCL transport. + ## Extending the System Every piece of the system is replaceable: the weights you send (`WeightSource`), diff --git a/docs/usage/security.md b/docs/usage/security.md index 68c05a68420d..5f7b0ce75329 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -345,7 +345,7 @@ vLLM supports loading out-of-tree HTTP routes via the `vllm.endpoint_plugins` en ## gRPC Interface -vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server. +vLLM provides optional gRPC `Inference` and `Control` services on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server. **Warning:** The gRPC interface is **insecure by default** — it does not implement authentication, authorization, or encryption. It should be considered a private, internal interface intended for use only between co-located services within a trusted network. Do not expose the gRPC port to the public internet or untrusted clients. If you enable the gRPC interface, protect it via network-level access controls such as firewall rules, network segmentation, or deployment on an isolated private network. @@ -354,8 +354,9 @@ vLLM provides an optional gRPC Generate service on a separate TCP port, enabled An attacker who can reach the gRPC port can: 1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials -2. **Consume GPU and compute resources** by submitting unbounded generation requests -3. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM. +2. **Mutate engine state** by pausing generation, sleeping the engine, or initiating configured RL weight updates through the `Control` service +3. **Consume GPU and compute resources** by submitting unbounded generation requests +4. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM. ### Recommendations diff --git a/examples/deployment/chart-helm/Chart.yaml b/examples/deployment/chart-helm/Chart.yaml index fb0f06f6d270..5270fe40b893 100644 --- a/examples/deployment/chart-helm/Chart.yaml +++ b/examples/deployment/chart-helm/Chart.yaml @@ -15,7 +15,7 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 0.0.1 +version: 0.0.2 maintainers: - name: mfournioux diff --git a/examples/deployment/chart-helm/templates/_helpers.tpl b/examples/deployment/chart-helm/templates/_helpers.tpl index 3226c1d79c42..7ddb23c686ef 100644 --- a/examples/deployment/chart-helm/templates/_helpers.tpl +++ b/examples/deployment/chart-helm/templates/_helpers.tpl @@ -12,10 +12,17 @@ Define service name {{- if .Values.serviceName }} {{- .Values.serviceName | lower | trim }} {{- else }} -"{{ .Release.Name }}-service" +{{- printf "%s-service" .Release.Name }} {{- end }} {{- end }} +{{/* +Define deployment name +*/}} +{{- define "chart.deployment-name" -}} +{{- printf "%s-deployment-vllm" .Release.Name }} +{{- end }} + {{/* Define service port */}} @@ -162,4 +169,4 @@ runAsUser: {{- with .Values.labels -}} {{ toYaml . }} {{- end }} -{{- end }} \ No newline at end of file +{{- end }} diff --git a/examples/deployment/chart-helm/templates/deployment.yaml b/examples/deployment/chart-helm/templates/deployment.yaml index a0a3c4b9ee52..2eada91b9a4c 100644 --- a/examples/deployment/chart-helm/templates/deployment.yaml +++ b/examples/deployment/chart-helm/templates/deployment.yaml @@ -1,23 +1,21 @@ apiVersion: apps/v1 kind: Deployment metadata: - name: "{{ .Release.Name }}-deployment-vllm" + name: {{ include "chart.deployment-name" . | quote }} namespace: {{ .Release.Namespace }} labels: {{- include "chart.labels" . | nindent 4 }} spec: replicas: {{ .Values.replicaCount }} {{- include "chart.strategy" . | nindent 2 }} - selector: + selector: matchLabels: - environment: "test" - release: "test" + {{- include "chart.labels" . | nindent 6 }} progressDeadlineSeconds: 1200 template: metadata: labels: - environment: "test" - release: "test" + {{- include "chart.labels" . | nindent 8 }} spec: containers: - name: "vllm" @@ -128,4 +126,4 @@ spec: values: {{- toYaml . | nindent 20 }} {{- end }} - {{- end }} \ No newline at end of file + {{- end }} diff --git a/examples/deployment/chart-helm/templates/hpa.yaml b/examples/deployment/chart-helm/templates/hpa.yaml index 5ca94c821354..fa1f862a1bb1 100644 --- a/examples/deployment/chart-helm/templates/hpa.yaml +++ b/examples/deployment/chart-helm/templates/hpa.yaml @@ -8,7 +8,7 @@ spec: scaleTargetRef: apiVersion: apps/v1 kind: Deployment - name: vllm + name: {{ include "chart.deployment-name" . | quote }} minReplicas: {{ .Values.autoscaling.minReplicas }} maxReplicas: {{ .Values.autoscaling.maxReplicas }} metrics: @@ -28,4 +28,4 @@ spec: type: Utilization averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }} {{- end }} -{{- end }} \ No newline at end of file +{{- end }} diff --git a/examples/deployment/chart-helm/templates/service.yaml b/examples/deployment/chart-helm/templates/service.yaml index 12d0f68b03a3..a2ee8e2887db 100644 --- a/examples/deployment/chart-helm/templates/service.yaml +++ b/examples/deployment/chart-helm/templates/service.yaml @@ -1,7 +1,7 @@ apiVersion: v1 kind: Service metadata: - name: "{{ .Release.Name }}-service" + name: {{ include "chart.service-name" . | quote }} namespace: {{ .Release.Namespace }} spec: type: ClusterIP @@ -11,4 +11,4 @@ spec: targetPort: {{ include "chart.container-port-name" . }} protocol: TCP selector: - {{- include "chart.labels" . | nindent 4 }} \ No newline at end of file + {{- include "chart.labels" . | nindent 4 }} diff --git a/examples/deployment/chart-helm/tests/deployment_test.yaml b/examples/deployment/chart-helm/tests/deployment_test.yaml index 9b7472cf0fd4..90583d54c02e 100644 --- a/examples/deployment/chart-helm/tests/deployment_test.yaml +++ b/examples/deployment/chart-helm/tests/deployment_test.yaml @@ -2,6 +2,23 @@ suite: test deployment templates: - deployment.yaml tests: + - it: should use configured labels for the deployment selector and pods + set: + labels: + environment: production + release: qwen-serving + asserts: + - equal: + path: spec.selector.matchLabels + value: + environment: production + release: qwen-serving + - equal: + path: spec.template.metadata.labels + value: + environment: production + release: qwen-serving + - it: should create wait-download-model init container when modelDownload is enabled set: extraInit: @@ -132,4 +149,4 @@ tests: value: ghcr.io/llm-d/llm-d-routing-sidecar:v0.2.0 - equal: path: spec.template.spec.initContainers[1].ports[0].containerPort - value: 8080 \ No newline at end of file + value: 8080 diff --git a/examples/deployment/chart-helm/tests/hpa_test.yaml b/examples/deployment/chart-helm/tests/hpa_test.yaml new file mode 100644 index 000000000000..49e664bc7b16 --- /dev/null +++ b/examples/deployment/chart-helm/tests/hpa_test.yaml @@ -0,0 +1,17 @@ +suite: test horizontal pod autoscaler +templates: + - hpa.yaml +release: + name: demo +tests: + - it: should target the deployment created by this release + set: + autoscaling: + enabled: true + minReplicas: 1 + maxReplicas: 3 + targetCPUUtilizationPercentage: 80 + asserts: + - equal: + path: spec.scaleTargetRef.name + value: demo-deployment-vllm diff --git a/examples/deployment/chart-helm/tests/service_test.yaml b/examples/deployment/chart-helm/tests/service_test.yaml new file mode 100644 index 000000000000..67ddc4f894a7 --- /dev/null +++ b/examples/deployment/chart-helm/tests/service_test.yaml @@ -0,0 +1,25 @@ +suite: test service +templates: + - service.yaml +release: + name: demo +tests: + - it: should honor the configured service name + set: + serviceName: vllm-api + asserts: + - equal: + path: metadata.name + value: vllm-api + + - it: should select pods using the configured labels + set: + labels: + environment: production + release: qwen-serving + asserts: + - equal: + path: spec.selector + value: + environment: production + release: qwen-serving diff --git a/examples/deployment/chart-helm/values.schema.json b/examples/deployment/chart-helm/values.schema.json index 0d0e0098bc19..f0276ff05141 100644 --- a/examples/deployment/chart-helm/values.schema.json +++ b/examples/deployment/chart-helm/values.schema.json @@ -28,7 +28,10 @@ "type": "integer" }, "serviceName": { - "type": "null" + "type": [ + "null", + "string" + ] }, "servicePort": { "type": "integer" @@ -326,4 +329,4 @@ "secrets", "servicePort" ] -} \ No newline at end of file +} diff --git a/examples/tool_chat_template_muse_glimmer.jinja b/examples/tool_chat_template_muse_glimmer.jinja new file mode 100644 index 000000000000..0ea2208934f9 --- /dev/null +++ b/examples/tool_chat_template_muse_glimmer.jinja @@ -0,0 +1 @@ +{%- macro render_content(content) -%}{%- if content is string -%}{{- content -}}{%- elif content is not none -%}{%- for part in content -%}{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}{%- elif part['type'] == 'text' -%}{{- part['text'] -}}{%- endif -%}{%- endfor -%}{%- endif -%}{%- endmacro -%}{%- macro render_atem(tc) -%}{%- set args = tc.function.arguments -%}{%- if args is not mapping -%}{{- raise_exception('MuseGlimmer ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.') -}}{%- endif -%}{{- '\n\n' -}}{%- for k, v in args.items() -%}{{- '' -}}{%- if v is boolean -%}{%- if v -%}true{%- else -%}false{%- endif -%}{%- elif v is none -%}null{%- elif v is mapping or (v is iterable and v is not string) -%}{{- v | tojson -}}{%- else -%}{{- v -}}{%- endif -%}{{- '\n' -}}{%- endfor -%}{{- '\n' -}}{%- endmacro -%}{%- macro render_tool_defs(tools) -%}{{- 'In this environment you have access to a set of tools you can use to answer the user\'s question.\n\n' -}}{{- 'You can invoke a function by writing a "" block like the following:\n' -}}{{- '\n\n$PARAMETER_VALUE\n...\n\n\n\n' -}}{{- 'String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n' -}}{{- 'Here are the functions available in JSONSchema format:\n' -}}{{- '// Tool metadata\n' -}}{%- set nsns = namespace(seen=[]) -%}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{%- set tns = fn.name.split('.')[0] -%}{%- if tns not in nsns.seen -%}{%- set nsns.seen = nsns.seen + [tns] -%}{%- endif -%}{%- endfor -%}{%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%}{%- for tns in nsns.seen -%}{{- '{"name": ' + (tns | tojson) + ', "description": ' + ((nd[tns] if tns in nd else '') | tojson) + '}\n' -}}{%- endfor -%}{{- '// Function schemas' -}}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{{- '\n{"name": ' + (fn.name | tojson) + ', "description": ' + (fn.description | tojson) + ', "parameters": ' + (fn.parameters | tojson) + '}' -}}{%- endfor -%}{{- '\n\nHere\'s an example of how to call a function in the tool set. To make parallel tool calls, emit each call as its own message (one tool call per message) in consecutive messages:\n' -}}{{- '(If the tool namespace is not specified, invoke the function directly as `example_function_name` rather than `example_tool_name.example_function_name`)\n\n' -}}{{- 'to=example_tool_name.example_function_name\n\n' -}}{{- '\n\n' -}}{{- 'value_1\n' -}}{{- 'This is the value for the second parameter\nthat can span\n"multiple" lines\n\n' -}}{{- '\n' -}}{%- endmacro -%}{%- macro render_reasoning() -%}{%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else 'high' -%}{{- 'Reasoning strength: ' + rs + '.' -}}{%- endmacro -%}{%- macro render_system_meta(tools) -%}{%- set rns = namespace(recipients=['"self"'], nslist=[]) -%}{%- if tools -%}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{%- set tns = fn.name.split('.')[0] -%}{%- if tns not in rns.nslist -%}{%- set rns.nslist = rns.nslist + [tns] -%}{%- endif -%}{%- endfor -%}{%- for tns in rns.nslist -%}{%- set rns.recipients = rns.recipients + ['"' + tns + '.*"'] -%}{%- endfor -%}{%- endif -%}{%- set rns.recipients = rns.recipients + ['"user"'] -%}{{- '# Valid recipients: ' + rns.recipients | join(', ') + '.' -}}{%- endmacro -%}{{- bos_token -}}{%- set ns = namespace(has_system=false) -%}{%- for m in messages -%}{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}{%- endfor -%}{%- if not ns.has_system and (add_generation_prompt or tools) -%}{{- '<|start|>system<|message|>You are a helpful AI assistant.' -}}{%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else '2026-01-04' -%}{{- '\nKnowledge cutoff: ' + kc + '.' -}}{%- if current_date is defined and current_date -%}{{- '\nCurrent date: ' + current_date + '.' -}}{%- elif strftime_now is defined -%}{{- '\nCurrent date: ' + strftime_now('%Y-%m-%d') + '.' -}}{%- endif -%}{{- '\n\n' -}}{{- render_reasoning() -}}{%- if tools -%}{{- '\n\n' -}}{{- render_tool_defs(tools) -}}{%- endif -%}{{- '\n\n' -}}{{- render_system_meta(tools) -}}{{- '<|eot|>' -}}{%- endif -%}{%- for message in messages -%}{%- set role = message['role'] -%}{%- set end_token = '<|eom|>' if (not loop.last and messages[loop.index0 + 1]['role'] == role) else '<|eot|>' -%}{%- if role == 'system' -%}{{- '<|start|>system<|message|>' -}}{{- render_content(message['content']) -}}{{- '\n\n' -}}{{- render_reasoning() -}}{%- if tools -%}{{- '\n\n' -}}{{- render_tool_defs(tools) -}}{%- endif -%}{{- '\n\n' -}}{{- render_system_meta(tools) -}}{{- '<|eot|>' -}}{%- elif role == 'user' -%}{{- '<|start|>user<|message|>' -}}{{- render_content(message['content']) -}}{{- '<|eot|>' -}}{%- elif role == 'tool' -%}{%- set tname = message.get('name') -%}{%- if not tname -%}{%- set tcid = message.get('tool_call_id') -%}{%- set rns = namespace(name=tcid if tcid else '') -%}{%- for m in messages -%}{%- if m.get('tool_calls') -%}{%- for tc in m['tool_calls'] -%}{%- if tcid is not none and tc.id is defined and tc.id == tcid -%}{%- set rns.name = tc.function.name -%}{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}{%- set tname = rns.name -%}{%- endif -%}{{- '<|start|>tool ' + tname + '<|message|>\n' -}}{{- render_content(message['content']) -}}{{- '\n<|eot|>' -}}{%- elif role == 'assistant' -%}{%- if message.get('reasoning_content') -%}{{- '<|start|>assistant to=self<|message|>' + message['reasoning_content'] + '<|eom|>' -}}{%- endif -%}{%- if message.get('tool_calls') -%}{%- for tc in message['tool_calls'] -%}{{- '<|start|>assistant to=' + tc.function.name + '<|message|>' -}}{{- render_atem(tc) -}}{%- if loop.last -%}{{- end_token -}}{%- else -%}{{- '<|eom|>' -}}{%- endif -%}{%- endfor -%}{%- else -%}{%- set recipient = message.get('recipient') or 'user' -%}{%- set end_turn = message.get('end_turn') -%}{%- if end_turn is none -%}{%- set end_turn = not (recipient and recipient != 'user') -%}{%- endif -%}{{- '<|start|>assistant' -}}{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}{{- '<|message|>' -}}{{- render_content(message['content']) -}}{{- ('<|eot|>' if end_turn else '<|eom|>') -}}{%- endif -%}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%} diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index d7a396bfa0ca..1760ed651734 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -1,7 +1,7 @@ lmcache >= 0.3.9 -# CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0 -# until a fixed newer release is verified for runtime images. -cupy-cuda13x < 14.1.0 +# CuPy 14.1.0 imports pytest from cupy.testing._random, which breaks runtime +# images. 14.1.1 fixes the root cause, so only 14.1.0 is excluded. +cupy-cuda13x != 14.1.0 nixl == 1.3.2 # CUDA 12 build. On the CUDA 13 image install-kv-connectors.sh swaps this for # the mooncake-transfer-engine-cuda13 variant of the same version. diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index cb45f29ed59a..741e33dc0245 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -126,7 +126,7 @@ cloudpickle==3.1.2 # via -r requirements/test/../common.txt cohere==7.0.8 # via -r requirements/test/cuda.in -cohere-melody==0.9.0 +cohere-melody==0.11.1 # via -r requirements/test/cuda.in colorama==0.4.6 # via diff --git a/requirements/test/cuda.in b/requirements/test/cuda.in index 8a56743b1c62..d11da082ad16 100644 --- a/requirements/test/cuda.in +++ b/requirements/test/cuda.in @@ -71,7 +71,7 @@ gpt-oss >= 0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test -cohere_melody>=0.9.0 # required for cohere command reasoning parser test +cohere_melody>=0.11.1 # required for cohere command reasoning/tool parser tests cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types) # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index f0bea6e5ba6d..fa3597b77017 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -131,7 +131,7 @@ cloudpickle==3.1.2 # via -r requirements/test/../common.txt cohere==7.0.8 # via -r requirements/test/cuda.in -cohere-melody==0.9.0 +cohere-melody==0.11.1 # via -r requirements/test/cuda.in colorama==0.4.6 # via diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 5858a028ca97..a7986b6e04c1 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -69,7 +69,7 @@ gpt-oss>=0.0.7; python_version > '3.11' perceptron # required for isaac test kaldi-native-fbank>=1.18.7 # required for fireredasr2 test -cohere_melody>=0.9.0 # required for cohere command reasoning parser test +cohere_melody>=0.11.1 # required for cohere command reasoning/tool parser tests cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types) # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library. diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 83f491a257fa..a7b7172ac22c 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -131,7 +131,7 @@ cloudpickle==3.1.2 # tilelang cohere==7.0.8 # via -r requirements/test/rocm.in -cohere-melody==0.9.0 +cohere-melody==0.11.1 # via -r requirements/test/rocm.in colorama==0.4.6 # via diff --git a/requirements/xpu.txt b/requirements/xpu.txt index aa33fe3f0ecf..2997c9aba0f6 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -20,4 +20,4 @@ torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend auto_round_lib==0.14.2 -vllm_xpu_kernels==0.1.12.3 \ No newline at end of file +vllm_xpu_kernels==0.1.13.2 \ No newline at end of file diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 0e25aea26474..428ad6460499 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -9,6 +9,21 @@ service Control { rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} rpc Abort (AbortRequest) returns (AbortResponse) {} rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} + + // Reinforcement-learning lifecycle and weight updates. + rpc PauseGeneration (PauseGenerationRequest) returns (PauseGenerationResponse) {} + rpc ResumeGeneration (ResumeGenerationRequest) returns (ResumeGenerationResponse) {} + rpc IsPaused (IsPausedRequest) returns (IsPausedResponse) {} + rpc Sleep (SleepRequest) returns (SleepResponse) {} + rpc WakeUp (WakeUpRequest) returns (WakeUpResponse) {} + rpc IsSleeping (IsSleepingRequest) returns (IsSleepingResponse) {} + rpc InitWeightTransferEngine (InitWeightTransferEngineRequest) returns (InitWeightTransferEngineResponse) {} + rpc StartWeightUpdate (StartWeightUpdateRequest) returns (StartWeightUpdateResponse) {} + rpc StartDraftWeightUpdate (StartDraftWeightUpdateRequest) returns (StartDraftWeightUpdateResponse) {} + rpc UpdateWeights (UpdateWeightsRequest) returns (UpdateWeightsResponse) {} + rpc FinishWeightUpdate (FinishWeightUpdateRequest) returns (FinishWeightUpdateResponse) {} + rpc UpdateWeightVersion (UpdateWeightVersionRequest) returns (UpdateWeightVersionResponse) {} + rpc GetWeightVersion (GetWeightVersionRequest) returns (GetWeightVersionResponse) {} } message GetServerInfoRequest {} @@ -23,6 +38,14 @@ message ServerInfo { uint64 total_kv_blocks = 7; uint64 max_running_requests = 8; uint64 max_batched_tokens = 9; + RlCapabilities rl_capabilities = 11; +} + +message RlCapabilities { + bool weight_transfer_enabled = 1; + string weight_transfer_backend = 2; + bool sleep_mode_enabled = 3; + bool draft_weight_updates_enabled = 4; } message ParallelismInfo { @@ -53,6 +76,64 @@ message AbortRequest { message AbortResponse {} +// ====================================================================================== +// Reinforcement-learning control +// ====================================================================================== + +enum PauseMode { + PAUSE_MODE_UNSPECIFIED = 0; + PAUSE_MODE_ABORT = 1; + PAUSE_MODE_WAIT = 2; + PAUSE_MODE_KEEP = 3; +} + +message PauseGenerationRequest { + PauseMode mode = 1; + optional bool clear_cache = 2; +} +message PauseGenerationResponse {} + +message ResumeGenerationRequest {} +message ResumeGenerationResponse {} + +message IsPausedRequest {} +message IsPausedResponse { bool paused = 1; } + +message SleepRequest { + optional uint32 level = 1; + PauseMode mode = 2; +} +message SleepResponse {} + +message WakeUpRequest { repeated string tags = 1; } +message WakeUpResponse {} + +message IsSleepingRequest {} +message IsSleepingResponse { bool sleeping = 1; } + +// The payloads are backend-specific JSON objects. Tensor data remains on the +// configured NCCL, IPC, or sparse-NCCL transport rather than crossing gRPC. +message InitWeightTransferEngineRequest { bytes init_info_json = 1; } +message InitWeightTransferEngineResponse {} + +message StartWeightUpdateRequest {} +message StartWeightUpdateResponse {} + +message StartDraftWeightUpdateRequest {} +message StartDraftWeightUpdateResponse {} + +message UpdateWeightsRequest { bytes update_info_json = 1; } +message UpdateWeightsResponse {} + +message FinishWeightUpdateRequest { optional string weight_version = 1; } +message FinishWeightUpdateResponse {} + +message UpdateWeightVersionRequest { string weight_version = 1; } +message UpdateWeightVersionResponse {} + +message GetWeightVersionRequest {} +message GetWeightVersionResponse { string weight_version = 1; } + // ====================================================================================== // KV discovery // ====================================================================================== diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 021d93a1f7fb..897319ca9c07 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -107,6 +107,8 @@ message ResponseOptions { bool output_token_ids = 5; bool output_logprobs = 6; optional CandidateTokens output_candidates = 7; + // Defaults to true when omitted. + optional bool skip_special_tokens = 8; } message KVCacheParameters { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index d1c20ad58220..7f8663f85268 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -1,12 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; use futures::future::{join_all, try_join_all}; use itertools::Itertools; use serde::Serialize; +use serde_json::Value as JsonValue; use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, info, trace}; @@ -674,6 +676,77 @@ impl EngineCoreClient { .collect()) } + /// Initialize the configured RL weight-transfer backend. + pub async fn init_weight_transfer_engine(&self, init_info: JsonValue) -> Result<()> { + self.collective_rpc( + "init_weight_transfer_engine", + None, + Vec::::new(), + BTreeMap::from([("init_info".to_string(), init_info)]), + ) + .await?; + Ok(()) + } + + /// Start a weight update for the base model. + pub async fn start_weight_update(&self) -> Result<()> { + self.collective_rpc( + "start_weight_update", + None, + Vec::::new(), + BTreeMap::::new(), + ) + .await?; + Ok(()) + } + + /// Start a weight update for the speculative draft model. + pub async fn start_draft_weight_update(&self) -> Result<()> { + self.collective_rpc( + "start_draft_weight_update", + None, + Vec::::new(), + BTreeMap::::new(), + ) + .await?; + Ok(()) + } + + /// Apply one backend-specific weight metadata chunk. + pub async fn update_weights(&self, update_info: JsonValue) -> Result<()> { + self.collective_rpc( + "update_weights", + None, + Vec::::new(), + BTreeMap::from([("update_info".to_string(), update_info)]), + ) + .await?; + Ok(()) + } + + /// Finish the current weight update. + pub async fn finish_weight_update(&self) -> Result<()> { + self.collective_rpc( + "finish_weight_update", + None, + Vec::::new(), + BTreeMap::::new(), + ) + .await?; + Ok(()) + } + + /// Set the committed weight version on every connected engine. + pub async fn set_weight_version(&self, weight_version: &str) -> Result<()> { + self.call_utility::<(), _>("set_weight_version", (weight_version,)).await?; + Ok(()) + } + + /// Return the committed weight version agreed on by every connected engine. + pub async fn get_weight_version(&self) -> Result { + self.call_utility_consensus("get_weight_version", ()).await + } + /// Return whether the engine is currently sleeping at any level. pub async fn is_sleeping(&self) -> Result { self.call_utility_consensus("is_sleeping", ()).await diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index ee066635ff0c..28fb4d84fe05 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -67,6 +67,9 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { kv_cache_size_tokens: None, kv_cache_max_concurrency: None, kv_events_config: None, + weight_transfer_backend: None, + enable_sleep_mode: false, + supports_draft_weight_updates: false, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index a4e1cebde67d..b4214a06f27f 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -87,6 +87,15 @@ pub struct EngineCoreReadyResponse { /// KV-event publisher configuration, if configured. #[serde(default)] pub kv_events_config: Option, + /// Configured RL weight-transfer backend, if weight transfer is enabled. + #[serde(default)] + pub weight_transfer_backend: Option, + /// Whether the engine was started with sleep mode enabled. + #[serde(default)] + pub enable_sleep_mode: bool, + /// Whether the engine has a speculative draft model that can be updated. + #[serde(default)] + pub supports_draft_weight_updates: bool, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 1d8d6e37cce3..8f07a8059e92 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2686,6 +2686,12 @@ fn python_msgpack_fixtures_match_rust_encoding() { let ready_response: EngineCoreReadyResponse = rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); + assert_eq!( + ready_response.weight_transfer_backend.as_deref(), + Some("nccl") + ); + assert!(ready_response.enable_sleep_mode); + assert!(ready_response.supports_draft_weight_updates); let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode"); assert!(kv_events_config.enable_kv_cache_events); assert_eq!(kv_events_config.publisher, "zmq"); diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index 7883c7150cab..389e97cd9996 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -412,6 +412,9 @@ class EngineCoreReadyResponse: kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None kv_events_config: KVEventsConfig | None = None + weight_transfer_backend: str | None = None + enable_sleep_mode: bool = False + supports_draft_weight_updates: bool = False ready_response = EngineCoreReadyResponse( @@ -430,6 +433,9 @@ class EngineCoreReadyResponse: max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + weight_transfer_backend="nccl", + enable_sleep_mode=True, + supports_draft_weight_updates=True, kv_events_config=KVEventsConfig( enable_kv_cache_events=True, publisher="zmq", diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index 201d8ace95ee..c7d9af88c82f 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -3,9 +3,13 @@ use std::sync::Arc; +use serde_json::Value as JsonValue; use thiserror_ext::AsReport as _; +use tokio::sync::Mutex; use tonic::{Request, Response, Status}; +use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; +use vllm_engine_core_client::protocol::utility::PauseMode as EnginePauseMode; use super::{ControlServer, pb}; use crate::state::AppState; @@ -15,17 +19,25 @@ pub(crate) type ControlGrpcService = ControlServer; /// gRPC control service backed by the shared application state. pub struct ControlServiceImpl { state: Arc, + rl_lock: Mutex<()>, } impl ControlServiceImpl { pub fn new(state: Arc) -> Self { - Self { state } + Self { + state, + rl_lock: Mutex::new(()), + } } fn ready(&self) -> &EngineCoreReadyResponse { self.state.engine_core_client().ready_response() } + fn client(&self) -> &EngineCoreClient { + self.state.engine_core_client() + } + fn parallelism_info(&self) -> pb::ParallelismInfo { let ready = self.ready(); pb::ParallelismInfo { @@ -36,10 +48,91 @@ impl ControlServiceImpl { decode_context_parallel_size: ready.decode_context_parallel_size, } } + + fn weight_transfer_backend(&self) -> Option<&str> { + let responses = self.client().ready_responses(); + let backend = responses.first()?.weight_transfer_backend.as_deref()?; + responses + .iter() + .all(|ready| ready.weight_transfer_backend.as_deref() == Some(backend)) + .then_some(backend) + } + + fn sleep_mode_enabled(&self) -> bool { + self.client().ready_responses().iter().all(|ready| ready.enable_sleep_mode) + } + + fn draft_weight_updates_enabled(&self) -> bool { + self.client() + .ready_responses() + .iter() + .all(|ready| ready.supports_draft_weight_updates) + } + + fn rl_capabilities(&self) -> pb::RlCapabilities { + let backend = self.weight_transfer_backend(); + pb::RlCapabilities { + weight_transfer_enabled: backend.is_some(), + weight_transfer_backend: backend.unwrap_or_default().to_string(), + sleep_mode_enabled: self.sleep_mode_enabled(), + draft_weight_updates_enabled: self.draft_weight_updates_enabled(), + } + } + + fn require_weight_transfer(&self) -> Result<(), Status> { + self.weight_transfer_backend().map(|_| ()).ok_or_else(|| { + Status::failed_precondition( + "weight transfer is not configured; start vLLM with --weight-transfer-config", + ) + }) + } + + fn require_sleep_mode(&self) -> Result<(), Status> { + self.sleep_mode_enabled().then_some(()).ok_or_else(|| { + Status::failed_precondition( + "sleep mode is not configured; start vLLM with --enable-sleep-mode", + ) + }) + } } const GRPC_API_VERSION: &str = "vllm"; +fn utility_status(method: &'static str, error: vllm_engine_core_client::Error) -> Status { + Status::internal(format!("{method} failed: {}", error.to_report_string())) +} + +fn pause_mode(mode: i32) -> Result { + match pb::PauseMode::try_from(mode) { + Ok(pb::PauseMode::Unspecified | pb::PauseMode::Abort) => Ok(EnginePauseMode::Abort), + Ok(pb::PauseMode::Wait) => Ok(EnginePauseMode::Wait), + Ok(pb::PauseMode::Keep) => Ok(EnginePauseMode::Keep), + Err(_) => Err(Status::invalid_argument("invalid pause mode")), + } +} + +fn json_object(bytes: &[u8], field: &'static str) -> Result { + let value = serde_json::from_slice::(bytes).map_err(|error| { + Status::invalid_argument(format!( + "{field} must contain valid JSON: {}", + error.to_report_string() + )) + })?; + if !value.is_object() { + return Err(Status::invalid_argument(format!( + "{field} must contain a JSON object" + ))); + } + Ok(value) +} + +fn weight_version(value: String) -> Result { + if value.trim().is_empty() { + return Err(Status::invalid_argument("weight_version must not be empty")); + } + Ok(value) +} + #[tonic::async_trait] impl pb::control_server::Control for ControlServiceImpl { async fn get_server_info( @@ -57,6 +150,7 @@ impl pb::control_server::Control for ControlServiceImpl { total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(), max_running_requests: ready.max_num_seqs, max_batched_tokens: ready.max_num_batched_tokens, + rl_capabilities: Some(self.rl_capabilities()), })) } @@ -112,6 +206,194 @@ impl pb::control_server::Control for ControlServiceImpl { let sources = client.ready_responses().into_iter().filter_map(kv_event_source).collect(); Ok(Response::new(pb::GetKvEventSourcesResponse { sources })) } + + async fn pause_generation( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let mode = pause_mode(request.mode)?; + let clear_cache = request.clear_cache.unwrap_or(true); + let _guard = self.rl_lock.lock().await; + self.client() + .pause_scheduler(mode, clear_cache) + .await + .map_err(|error| utility_status("pause_generation", error))?; + Ok(Response::new(pb::PauseGenerationResponse {})) + } + + async fn resume_generation( + &self, + _request: Request, + ) -> Result, Status> { + let _guard = self.rl_lock.lock().await; + self.client() + .resume_scheduler() + .await + .map_err(|error| utility_status("resume_generation", error))?; + Ok(Response::new(pb::ResumeGenerationResponse {})) + } + + async fn is_paused( + &self, + _request: Request, + ) -> Result, Status> { + let paused = self + .client() + .is_scheduler_paused() + .await + .map_err(|error| utility_status("is_paused", error))?; + Ok(Response::new(pb::IsPausedResponse { paused })) + } + + async fn sleep( + &self, + request: Request, + ) -> Result, Status> { + self.require_sleep_mode()?; + let request = request.into_inner(); + let mode = pause_mode(request.mode)?; + let level = request.level.unwrap_or(1); + let _guard = self.rl_lock.lock().await; + self.client() + .sleep(level, mode) + .await + .map_err(|error| utility_status("sleep", error))?; + Ok(Response::new(pb::SleepResponse {})) + } + + async fn wake_up( + &self, + request: Request, + ) -> Result, Status> { + self.require_sleep_mode()?; + let tags = request.into_inner().tags; + let tags = (!tags.is_empty()).then_some(tags); + let _guard = self.rl_lock.lock().await; + self.client() + .wake_up(tags) + .await + .map_err(|error| utility_status("wake_up", error))?; + Ok(Response::new(pb::WakeUpResponse {})) + } + + async fn is_sleeping( + &self, + _request: Request, + ) -> Result, Status> { + let sleeping = self + .client() + .is_sleeping() + .await + .map_err(|error| utility_status("is_sleeping", error))?; + Ok(Response::new(pb::IsSleepingResponse { sleeping })) + } + + async fn init_weight_transfer_engine( + &self, + request: Request, + ) -> Result, Status> { + self.require_weight_transfer()?; + let init_info = json_object(&request.into_inner().init_info_json, "init_info_json")?; + let _guard = self.rl_lock.lock().await; + self.client() + .init_weight_transfer_engine(init_info) + .await + .map_err(|error| utility_status("init_weight_transfer_engine", error))?; + Ok(Response::new(pb::InitWeightTransferEngineResponse {})) + } + + async fn start_weight_update( + &self, + _request: Request, + ) -> Result, Status> { + self.require_weight_transfer()?; + let _guard = self.rl_lock.lock().await; + self.client() + .start_weight_update() + .await + .map_err(|error| utility_status("start_weight_update", error))?; + Ok(Response::new(pb::StartWeightUpdateResponse {})) + } + + async fn start_draft_weight_update( + &self, + _request: Request, + ) -> Result, Status> { + self.require_weight_transfer()?; + if !self.draft_weight_updates_enabled() { + return Err(Status::failed_precondition( + "draft weight updates require a configured speculative draft model", + )); + } + let _guard = self.rl_lock.lock().await; + self.client() + .start_draft_weight_update() + .await + .map_err(|error| utility_status("start_draft_weight_update", error))?; + Ok(Response::new(pb::StartDraftWeightUpdateResponse {})) + } + + async fn update_weights( + &self, + request: Request, + ) -> Result, Status> { + self.require_weight_transfer()?; + let update_info = json_object(&request.into_inner().update_info_json, "update_info_json")?; + let _guard = self.rl_lock.lock().await; + self.client() + .update_weights(update_info) + .await + .map_err(|error| utility_status("update_weights", error))?; + Ok(Response::new(pb::UpdateWeightsResponse {})) + } + + async fn finish_weight_update( + &self, + request: Request, + ) -> Result, Status> { + self.require_weight_transfer()?; + let version = request.into_inner().weight_version.map(weight_version).transpose()?; + let _guard = self.rl_lock.lock().await; + self.client() + .finish_weight_update() + .await + .map_err(|error| utility_status("finish_weight_update", error))?; + if let Some(version) = version { + self.client() + .set_weight_version(&version) + .await + .map_err(|error| utility_status("update_weight_version", error))?; + } + Ok(Response::new(pb::FinishWeightUpdateResponse {})) + } + + async fn update_weight_version( + &self, + request: Request, + ) -> Result, Status> { + let version = weight_version(request.into_inner().weight_version)?; + let _guard = self.rl_lock.lock().await; + self.client() + .set_weight_version(&version) + .await + .map_err(|error| utility_status("update_weight_version", error))?; + Ok(Response::new(pb::UpdateWeightVersionResponse {})) + } + + async fn get_weight_version( + &self, + _request: Request, + ) -> Result, Status> { + let weight_version = self + .client() + .get_weight_version() + .await + .map_err(|error| utility_status("get_weight_version", error))?; + Ok(Response::new(pb::GetWeightVersionResponse { + weight_version, + })) + } } pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 6ac1abdcd894..4460f6b22249 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -158,7 +158,9 @@ pub fn to_text_request( } let decode_options = TextDecodeOptions { - skip_special_tokens: true, + skip_special_tokens: response + .and_then(|options| options.skip_special_tokens) + .unwrap_or(true), include_stop_str_in_output: stopping.is_some_and(|s| s.include_stop_strings), stop_strings: stopping.map(|s| &s.stop_strings).filter(|ss| !ss.is_empty()).cloned(), min_tokens: stopping.map_or(0, |s| s.min_new_tokens), diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 7eabd2701443..d5ee87a87f28 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -12,6 +12,7 @@ use std::time::Duration; use futures::StreamExt as _; use hyper_util::rt::TokioIo; use openssl::ssl::{SslConnector, SslFiletype, SslMethod}; +use rmpv::Value; use serial_test::serial; use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; use tokio::net::TcpStream; @@ -30,14 +31,15 @@ use vllm_engine_core_client::mock_engine::{ DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, default_ready_response, }; -use vllm_engine_core_client::protocol::handshake::KvEventsConfig; +use vllm_engine_core_client::protocol::decode_value; +use vllm_engine_core_client::protocol::handshake::{EngineCoreReadyResponse, KvEventsConfig}; use vllm_engine_core_client::protocol::output::{ EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, + UtilityCallOutput, }; use vllm_engine_core_client::protocol::request::EngineCoreRequest; -use vllm_engine_core_client::test_utils::{ - IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready, -}; +use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; +use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task_with_ready}; use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode}; use vllm_llm::Llm; use vllm_text::tokenizer::DynTokenizer; @@ -291,13 +293,10 @@ async fn setup_grpc_service_with_backend( where F: FnOnce(&EngineCoreRequest) + Send + 'static, { - let ipc = IpcNamespace::new().expect("create ipc namespace"); - let handshake_address = ipc.handshake_endpoint(); - let engine_id = engine_id.into(); - - let engine_task = MockEngineTask::new(spawn_mock_engine_task( - handshake_address.clone(), - engine_id.clone(), + setup_grpc_service_with_engine_script( + engine_id, + default_ready_response(), + backend, move |dealer, push| { boxed_test_future(async move { let add = recv_engine_message(dealer).await; @@ -311,6 +310,33 @@ where .await; }) }, + ) + .await +} + +async fn setup_grpc_service_with_engine_script( + engine_id: impl Into, + ready: EngineCoreReadyResponse, + backend: Arc, + script: F, +) -> ( + InferenceServer, + ControlServer, + tokio::sync::watch::Receiver, + MockEngineTask, +) +where + F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static, +{ + let ipc = IpcNamespace::new().expect("create ipc namespace"); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = engine_id.into(); + + let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready( + handshake_address.clone(), + engine_id.clone(), + ready, + script, )); let client = EngineCoreClient::connect( @@ -1509,6 +1535,11 @@ async fn control_reports_server_and_model_info() { assert_eq!(parallelism.data_parallel_size, 1); assert_eq!(parallelism.data_parallel_rank, 0); assert_eq!(parallelism.decode_context_parallel_size, 1); + let rl = server.rl_capabilities.expect("RL capabilities"); + assert!(!rl.weight_transfer_enabled); + assert!(rl.weight_transfer_backend.is_empty()); + assert!(!rl.sleep_mode_enabled); + assert!(!rl.draft_weight_updates_enabled); let model = client .get_model_info(pb::GetModelInfoRequest {}) @@ -1527,6 +1558,65 @@ async fn control_reports_server_and_model_info() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_forwards_weight_update_without_pause_guard() { + let mut ready = default_ready_response(); + ready.weight_transfer_backend = Some("nccl".to_string()); + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-rl".to_vec(), + ready, + Arc::new(FakeTextBackend), + |dealer, push| { + boxed_test_future(async move { + let frames = recv_engine_message(dealer).await; + assert_eq!(frames[0].as_ref(), &[0x03]); + let payload = decode_value(&frames[1]).expect("decode utility payload"); + let fields = payload.as_array().expect("utility payload array"); + let call_id = fields[1].as_u64().expect("utility call id"); + assert_eq!(fields[2].as_str(), Some("collective_rpc")); + let args = fields[3].as_array().expect("collective_rpc arguments"); + assert_eq!(args[0].as_str(), Some("update_weights")); + send_outputs( + push, + UtilityCallOutput { + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(UtilityResultEnvelope::without_type_info( + Value::Array(vec![Value::Nil]), + )), + }, + ..Default::default() + } + .into(), + ) + .await; + }) + }, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = ControlClient::new(channel); + + client + .update_weights(pb::UpdateWeightsRequest { + update_info_json: br#"{"names":["model.weight"]}"#.to_vec(), + }) + .await + .expect("forward weight update without a pause probe"); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn control_aggregates_multi_engine_capacity() { let ipc = IpcNamespace::new().expect("create ipc namespace"); @@ -1536,6 +1626,9 @@ async fn control_aggregates_multi_engine_capacity() { ready_0.max_model_len = 8_192; ready_0.num_gpu_blocks = 10; ready_0.data_parallel_size = 2; + ready_0.weight_transfer_backend = Some("nccl".to_string()); + ready_0.enable_sleep_mode = true; + ready_0.supports_draft_weight_updates = true; let mut ready_1 = default_ready_response(); ready_1.max_model_len = 4_096; @@ -1586,6 +1679,11 @@ async fn control_aggregates_multi_engine_capacity() { .into_inner(); assert_eq!(server.max_model_len, 4_096); assert_eq!(server.total_kv_blocks, 30); + let rl = server.rl_capabilities.expect("RL capabilities"); + assert!(!rl.weight_transfer_enabled); + assert!(rl.weight_transfer_backend.is_empty()); + assert!(!rl.sleep_mode_enabled); + assert!(!rl.draft_weight_updates_enabled); assert_eq!(server.parallelism.unwrap().data_parallel_size, 4); drop(engine_tasks); diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index ca19d0ab9e9d..4423287e5175 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -217,7 +217,8 @@ where health_reporter.set_serving::().await; health_reporter.set_serving::().await; let control_service = - grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone())); + grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone())) + .max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES); let inference_service = grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone())) .max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES); diff --git a/setup.py b/setup.py index 680376d5124e..87d3dacd4f89 100644 --- a/setup.py +++ b/setup.py @@ -507,20 +507,69 @@ class precompiled_wheel_utils: @staticmethod def fetch_metadata_for_variant( - commit: str, variant: str | None + commit: str, + variant: str | None, + *, + rocm: bool = False, ) -> tuple[list[dict], str]: """ Fetches metadata for a specific variant of the precompiled wheel. + + For non-ROCm, fetches vllm metadata. + + For ROCm, discovers all first-level packages and combines their + metadata into a single list. """ - variant_dir = f"{variant}/" if variant is not None else "" - repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/" - meta_url = repo_url + "metadata.json" - print(f"Trying to fetch nightly build metadata from {meta_url}") + import json + from html.parser import HTMLParser from urllib.request import urlopen - with urlopen(meta_url) as resp: - # urlopen raises HTTPError on unexpected status code - wheels = json.loads(resp.read().decode("utf-8")) + variant_dir = f"{variant}/" if variant is not None else "" + + if not rocm: + # Keep original behavior + repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/" + meta_url = repo_url + "metadata.json" + print(f"Trying to fetch nightly build metadata from {meta_url}") + with urlopen(meta_url) as resp: + wheels = json.loads(resp.read().decode("utf-8")) + + return wheels, repo_url + + # ROCm: discover all packages under the variant directory. + repo_url = f"https://wheels.vllm.ai/rocm/{commit}/{variant_dir}" + + class LinkParser(HTMLParser): + def __init__(self): + super().__init__() + self.links: list[str] = [] + + def handle_starttag(self, tag, attrs): + if tag == "a": + href = dict(attrs).get("href") + if href: + self.links.append(href) + + with urlopen(repo_url) as resp: + parser = LinkParser() + parser.feed(resp.read().decode("utf-8")) + + packages = [ + href.rstrip("/") + for href in parser.links + if href.endswith("/") + and href not in ("../", "./", "/") + and "/" not in href.rstrip("/") + ] + + wheels: list[dict] = [] + + for package in packages: + meta_url = f"{repo_url}{package}/metadata.json" + print(f"Trying to fetch nightly build metadata from {meta_url}") + with urlopen(meta_url) as resp: + package_wheels = json.loads(resp.read().decode("utf-8")) + wheels.extend(package_wheels) return wheels, repo_url @staticmethod @@ -582,6 +631,141 @@ def detect_system_cuda_variant() -> str: print(f"Detected CUDA {cuda_version}, using variant {variant}") return variant + @staticmethod + def rocm_version_to_variant(rocm_version: str) -> str: + """Convert a ROCm version string to a wheel variant, e.g. 7.2.3 -> rocm723.""" + return "rocm" + rocm_version.replace(".", "") + + @staticmethod + def detect_system_rocm_variant() -> str | None: + """Auto-detect the ROCm wheel variant from the installed ROCm stack.""" + rocm_version = get_rocm_version() + if not rocm_version: + try: + import torch + + rocm_version = torch.version.hip + except Exception: + pass + if not rocm_version: + return None + variant = precompiled_wheel_utils.rocm_version_to_variant(rocm_version) + print(f"Detected ROCm {rocm_version}, using variant {variant}") + return variant + + @staticmethod + def fetch_available_rocm_variants(commit: str) -> list[str]: + """List ROCm wheel variants published for a commit on wheels.vllm.ai.""" + from urllib.request import urlopen + + index_url = f"https://wheels.vllm.ai/rocm/{commit}/" + print(f"Fetching available ROCm variants from {index_url}") + try: + with urlopen(index_url) as resp: + html = resp.read().decode("utf-8") + except Exception as e: + logger.warning( + "Failed to fetch ROCm variant index for commit %s: %s", commit, e + ) + return [] + variants = sorted(set(re.findall(r"rocm\d+", html))) + print(f"Available ROCm variants for commit {commit}: {variants}") + return variants + + @staticmethod + def resolve_rocm_wheel_variant( + commit: str, variant_override: str | None + ) -> str | None: + """Resolve a ROCm wheel variant for a commit from wheels.vllm.ai.""" + env_variant = precompiled_wheel_utils.detect_system_rocm_variant() + available = precompiled_wheel_utils.fetch_available_rocm_variants(commit) + + if variant_override is not None: + if env_variant and variant_override != env_variant: + logger.warning( + "VLLM_PRECOMPILED_WHEEL_VARIANT=%s does not match the " + "detected environment ROCm variant %s; refusing to use a " + "different ROCm patch wheel", + variant_override, + env_variant, + ) + return None + if available and variant_override not in available: + logger.warning( + "Requested ROCm variant %s is not available for commit %s " + "(available: %s)", + variant_override, + commit, + available, + ) + return None + return variant_override + + if env_variant is None: + if available: + print( + "Could not detect ROCm version from the environment; " + f"available variants for commit {commit}: {available}" + ) + return None + + if env_variant in available: + return env_variant + + logger.warning( + "Environment ROCm variant %s is not available for commit %s " + "(available: %s). Precompiled wheels may not be compatible.", + env_variant, + commit, + available, + ) + return None + + @staticmethod + def warn_if_rocm_torch_version_mismatch( + wheels: list[dict], repo_url: str, arch: str + ) -> None: + """Warn if installed torch differs from the custom ROCm build on + wheels.vllm.ai and suggest the correct install command.""" + try: + installed = torch.__version__ + except Exception: + return + + def _wheel_version(pkg: str) -> str | None: + for w in wheels: + if w.get("package_name") == pkg and arch in w.get("platform_tag", ""): + v = w["version"] + if w.get("variant") and "+" not in v: + v = f"{v}+{w['variant']}" + return v + return None + + expected = _wheel_version("torch") + if expected is None or installed == expected: + return + + pkgs = f"torch=={expected}" + triton_ver = _wheel_version("triton") + if triton_ver: + pkgs += f" triton=={triton_ver}" + + logger.warning( + "Installed PyTorch %s does not match the custom build %s " + "shipped with vLLM ROCm wheels. The ABI may differ from " + "official releases. If you hit extension load errors, " + "reinstall from the vLLM index:\n" + " pip install %s --extra-index-url %s\n" + " uv pip install %s --extra-index-url %s " + "--index-strategy unsafe-best-match", + installed, + expected, + pkgs, + repo_url, + pkgs, + repo_url, + ) + @staticmethod def find_local_rocm_wheel() -> str | None: """Search for a local vllm wheel in common locations.""" @@ -639,6 +823,57 @@ def determine_wheel_url_rocm() -> tuple[str, str | None]: print(f"Found local ROCm wheel: {local_wheel}") return local_wheel, None + import platform + + arch = platform.machine() + commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower() + if not commit or len(commit) != 40: + print( + f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}" + ", trying to fetch base commit in main branch" + ) + commit = precompiled_wheel_utils.get_base_commit_in_main_branch() + variant = precompiled_wheel_utils.resolve_rocm_wheel_variant( + commit, os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None) + ) + print(f"Using precompiled ROCm wheel commit {commit} with variant {variant}") + wheels, repo_url = None, None + if variant is not None: + try: + wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant( + commit, variant, rocm=True + ) + precompiled_wheel_utils.warn_if_rocm_torch_version_mismatch( + wheels, repo_url, arch + ) + except Exception as e: + logger.warning( + "Failed to fetch ROCm wheel metadata for variant %s: %s", + variant, + e, + ) + if wheels is not None and repo_url is not None: + from urllib.parse import urljoin + + for wheel in wheels: + if wheel.get("package_name") == "vllm" and arch in wheel.get( + "platform_tag", "" + ): + print(f"Found precompiled wheel metadata: {wheel}") + if "path" not in wheel: + raise ValueError(f"Wheel metadata missing path: {wheel}") + wheel_url = urljoin(f"{repo_url}vllm/", wheel["path"]) + download_filename = wheel.get("filename") + print(f"Using precompiled wheel URL: {wheel_url}") + return wheel_url, download_filename + logger.warning( + "No precompiled vllm wheel found for architecture %s " + "from repo %s. All available wheels: %s", + arch, + repo_url, + wheels, + ) + # Fall back to AMD's PyPI index index_url = os.getenv( "VLLM_ROCM_WHEEL_INDEX", "https://pypi.amd.com/vllm-rocm/simple" @@ -669,8 +904,7 @@ def determine_wheel_url() -> tuple[str, str | None]: print(f"Using user-specified precompiled wheel location: {wheel_location}") return wheel_location, None else: - # ROCm: use local wheel or AMD's PyPI index - # TODO: When we have ROCm nightly wheels, we can update this logic. + # ROCm: resolve wheels from wheels.vllm.ai with environment-matched variant if precompiled_wheel_utils.is_rocm_system(): return precompiled_wheel_utils.determine_wheel_url_rocm() @@ -1291,6 +1525,7 @@ def add_vllm_package_data(filename: str) -> None: # only; also needs system GStreamer + libv4l (see docs). "deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"], "flashinfer": [], # Kept for backwards compatibility + "b12x": ["b12x==1.2.4"], # Optional deps for Helion kernel development # NOTE: When updating helion version, also update CI files: # - .buildkite/test_areas/kernels.yaml diff --git a/tests/basic_correctness/test_mem.py b/tests/basic_correctness/test_mem.py index 0618562b3ec5..696d2f66bc5b 100644 --- a/tests/basic_correctness/test_mem.py +++ b/tests/basic_correctness/test_mem.py @@ -75,6 +75,41 @@ def test_basic_cumem(): assert torch.allclose(output, torch.ones_like(output) * 3) +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +def test_discard_tags(): + """Test that discard(tags) selectively frees GPU memory for specific + tags while keeping other tags mapped and usable.""" + allocator = get_mem_allocator_instance() + + with allocator.use_memory_pool("weights"): + weights = torch.ones(1024, 1024, device=DEVICE_TYPE) + + with allocator.use_memory_pool("kv_cache"): + kv = torch.ones(512, 512, device=DEVICE_TYPE) + + free_bytes = torch.accelerator.get_memory_info()[0] + + # Discard kv_cache only — weights should remain valid + allocator.discard("kv_cache") + + free_bytes_after_discard = torch.accelerator.get_memory_info()[0] + assert free_bytes_after_discard > free_bytes + + # Weights are still usable + assert torch.allclose(weights, torch.ones_like(weights)) + + # Wake up and verify kv_cache is remapped (zeroed content) + allocator.wake_up() + # After wake_up the VA is remapped; content is not preserved + # but the allocation is valid + assert kv.shape == (512, 512) + + # Full sleep/wake cycle still works after discard + allocator.sleep(offload_tags="weights") + allocator.wake_up() + assert torch.allclose(weights, torch.ones_like(weights)) + + @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU") def test_cumem_with_cudagraph(): diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index 84a7733f1816..39b991204137 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -252,6 +252,11 @@ def test_splitting_ops_dynamic(): # populated when the engine decides to use piecewise compilation. assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE assert config.compilation_config.splitting_ops_contain_attention() + splitting_ops = config.compilation_config.splitting_ops + assert splitting_ops is not None + assert { + "vllm::qwen_gdn_attention_core_fused_norm_packed", + } <= set(splitting_ops) # When use_inductor_graph_partition=True config = VllmConfig( diff --git a/tests/config/test_config_utils.py b/tests/config/test_config_utils.py index 35bc1e167b52..24ef1b52a95d 100644 --- a/tests/config/test_config_utils.py +++ b/tests/config/test_config_utils.py @@ -216,6 +216,11 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs(): assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash +def test_cache_config_hash_ignores_prefix_cache_retention_interval(): + base_hash = CacheConfig().compute_hash() + assert CacheConfig(prefix_cache_retention_interval=64).compute_hash() == base_hash + + def test_envs_compile_factors_relocation_invariant(tmp_path): """Relocating HOME or the XDG roots must not change the compile-cache env hash. diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index 7e425d68eecb..75b4bd2963e0 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -86,7 +86,7 @@ def record(hf_config: PretrainedConfig) -> PretrainedConfig: @pytest.mark.cpu_test -def test_inkling_override_exposes_only_first_mtp_depth(): +def test_inkling_override_exposes_all_mtp_depths(): text_config = _make_hf_config( architectures=["InklingForCausalLM"], model_type="inkling_model", @@ -107,7 +107,9 @@ def test_inkling_override_exposes_only_first_mtp_depth(): assert out is text_config assert out.model_type == "inkling_mtp" assert out.architectures == ["InklingMTPModel"] - assert out.n_predict == 1 + # Multi-module MTP: every checkpoint depth is exposed (module i drafts + # speculative token i), no longer clamped to the first depth. + assert out.n_predict == 8 assert out.num_nextn_predict_layers == 8 assert out.chain_hidden_post_norm is False assert out.local_layer_ids == [0, 2, 4] diff --git a/tests/distributed/test_multiproc_executor.py b/tests/distributed/test_multiproc_executor.py index 042349e4fbb1..599025d1c8ec 100644 --- a/tests/distributed/test_multiproc_executor.py +++ b/tests/distributed/test_multiproc_executor.py @@ -11,15 +11,62 @@ import os import socket +import pytest + from tests.utils import multi_gpu_test from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.executor import multiproc_executor from vllm.v1.executor.multiproc_executor import MultiprocExecutor MODEL = "facebook/opt-125m" +@pytest.mark.parametrize( + ("local_world_size", "data_parallel_size_local", "expected_num_local_procs"), + [(1, 4, 4), (4, 1, 4), (2, 0, 2)], +) +def test_multiproc_executor_counts_all_local_dp_workers( + monkeypatch: pytest.MonkeyPatch, + local_world_size: int, + data_parallel_size_local: int, + expected_num_local_procs: int, +): + """All colocated DP workers share the node's startup CPU budget.""" + executor = object.__new__(MultiprocExecutor) + executor.world_size = local_world_size + executor.local_world_size = local_world_size + executor.parallel_config = type( + "ParallelConfig", + (), + {"data_parallel_size_local": data_parallel_size_local}, + )() + + monkeypatch.setattr( + executor, + "_get_parallel_sizes", + lambda: (local_world_size, 1, 1), + ) + + class StopExecutorInit(Exception): + pass + + def capture_num_local_procs(num_local_procs: int): + assert num_local_procs == expected_num_local_procs + raise StopExecutorInit + + monkeypatch.setattr( + multiproc_executor, + "set_multiprocessing_worker_envs", + capture_num_local_procs, + ) + + with pytest.raises(StopExecutorInit): + executor._init_executor() + executor._finalizer.detach() + + def create_vllm_config( tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 2feb9f7a039d..6a799ee912d4 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -476,6 +476,7 @@ def test_prefix_cache_default(): # should be None by default (depends on model). engine_args = EngineArgs.from_cli_args(args=args) assert engine_args.enable_prefix_caching is None + assert engine_args.prefix_cache_retention_interval == 0 # with flag to turn it on. args = parser.parse_args(["--enable-prefix-caching"]) @@ -487,6 +488,28 @@ def test_prefix_cache_default(): engine_args = EngineArgs.from_cli_args(args=args) assert not engine_args.enable_prefix_caching + args = parser.parse_args(["--prefix-cache-retention-interval", "64"]) + engine_args = EngineArgs.from_cli_args(args=args) + assert engine_args.prefix_cache_retention_interval == 64 + + +def test_prefix_cache_retention_interval_from_deprecated_env( + monkeypatch, caplog, disable_log_dedup +): + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") + + engine_args = EngineArgs() + + assert engine_args.prefix_cache_retention_interval == 64 + assert "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in caplog.text + assert "deprecated" in caplog.text + assert "prefix_cache_retention_interval" in caplog.text + + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = parser.parse_args(["--prefix-cache-retention-interval", "32"]) + engine_args = EngineArgs.from_cli_args(args) + assert engine_args.prefix_cache_retention_interval == 32 + @pytest.mark.parametrize( ("arg", "expected", "option"), diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index db792f8c7327..1cd7227bdc99 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -15,12 +15,14 @@ import json from argparse import Namespace from http import HTTPStatus +from typing import Annotated from unittest.mock import MagicMock import pytest from fastapi import FastAPI from fastapi.exceptions import RequestValidationError from fastapi.testclient import TestClient +from pydantic import BaseModel, Field, ValidationError from vllm.entrypoints.anthropic.api_router import attach_router from vllm.entrypoints.anthropic.protocol import ( @@ -44,7 +46,10 @@ PromptTokenUsageInfo, UsageInfo, ) -from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler +from vllm.entrypoints.serve.exception_handling.handlers.validation import ( + validation_exception_handler, +) +from vllm.exceptions import VLLMValidationError _convert = AnthropicServingMessages._convert_anthropic_to_openai_request _img_url = AnthropicServingMessages._convert_image_source_to_url @@ -1456,3 +1461,95 @@ def test_empty_cache_salt_returns_bad_request(self): assert response.status_code == HTTPStatus.BAD_REQUEST handler.create_messages.assert_not_awaited() + + +# ====================================================================== +# Client-caused errors are 4xx, not 500 (Issue #52088) +# ====================================================================== + + +class TestClientErrorResponses: + @staticmethod + def _make_api_app(handler: MagicMock): + app = FastAPI() + attach_router(app) + app.state.args = Namespace(log_error_stack=False) + app.exception_handler(RequestValidationError)(validation_exception_handler) + app.state.anthropic_serving_messages = handler + return app + + @staticmethod + def _request_body() -> dict: + return { + "model": "test-model", + "max_tokens": 1, + "messages": [{"role": "user", "content": "Hello"}], + } + + @staticmethod + def _conversion_error() -> ValidationError: + """A real pydantic ValidationError like the one ChatCompletionRequest + construction raises when Anthropic input violates the OpenAI schema.""" + + class _StubRequest(BaseModel): + stop: Annotated[list[str], Field(max_length=4)] | None = None + + with pytest.raises(ValidationError) as exc_info: + _StubRequest(stop=["a"] * 6) + return exc_info.value + + def test_validation_error_returns_bad_request(self): + """A pydantic ValidationError during Anthropic->OpenAI conversion is + surfaced as a 400 BadRequestError, not a 500.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = self._conversion_error() + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/v1/messages", json=self._request_body()) + + assert response.status_code == HTTPStatus.BAD_REQUEST + body = response.json() + assert body["type"] == "error" + assert body["error"]["type"] == "BadRequestError" + assert "at most 4 items" in body["error"]["message"] + + def test_vllm_client_error_returns_bad_request(self): + """VLLMClientError raised by the serving layer maps to 400.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = VLLMValidationError( + "Invalid value for stop", parameter="stop" + ) + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/v1/messages", json=self._request_body()) + + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json()["error"]["type"] == "BadRequestError" + + def test_generic_error_still_returns_internal_server_error(self): + """Non-client errors keep the existing 500 behaviour.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.create_messages.side_effect = RuntimeError("boom") + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post("/v1/messages", json=self._request_body()) + + assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert response.json()["error"]["type"] == "InternalServerError" + + def test_count_tokens_validation_error_returns_bad_request(self): + """The count_tokens route maps conversion errors to 400 as well.""" + handler = MagicMock(spec=AnthropicServingMessages) + handler.count_tokens.side_effect = self._conversion_error() + + app = self._make_api_app(handler) + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post( + "/v1/messages/count_tokens", json=self._request_body() + ) + + assert response.status_code == HTTPStatus.BAD_REQUEST + assert response.json()["error"]["type"] == "BadRequestError" diff --git a/tests/entrypoints/cohere/test_api_router.py b/tests/entrypoints/cohere/test_api_router.py index 8380faa29ebf..45e2ec5cb0df 100644 --- a/tests/entrypoints/cohere/test_api_router.py +++ b/tests/entrypoints/cohere/test_api_router.py @@ -30,8 +30,10 @@ CohereChatV2Response, ) from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse -from vllm.entrypoints.serve.utils.server_utils import ( +from vllm.entrypoints.serve.exception_handling.handlers.http import ( http_exception_handler, +) +from vllm.entrypoints.serve.exception_handling.handlers.validation import ( validation_exception_handler, ) diff --git a/tests/entrypoints/llm/offline_mode/test_offline_mode.py b/tests/entrypoints/llm/offline_mode/test_offline_mode.py index 0708597079fc..555e8b9b7d07 100644 --- a/tests/entrypoints/llm/offline_mode/test_offline_mode.py +++ b/tests/entrypoints/llm/offline_mode/test_offline_mode.py @@ -9,9 +9,6 @@ import regex as re import urllib3 -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory - MODEL_CONFIGS = [ { "model": "facebook/opt-125m", @@ -55,19 +52,26 @@ ] +def _create_runner(vllm_runner, model_config): + runner_config = model_config.copy() + model = runner_config.pop("model") + tokenizer_name = runner_config.pop("tokenizer", None) + return vllm_runner(model, tokenizer_name=tokenizer_name, **runner_config) + + @pytest.fixture(scope="module") -def cache_models(): +def cache_models(vllm_runner): # Cache model files first for model_config in MODEL_CONFIGS: - LLM(**model_config) - cleanup_dist_env_and_memory() + with _create_runner(vllm_runner, model_config): + pass yield @pytest.mark.skip_global_cleanup @pytest.mark.usefixtures("cache_models") -def test_offline_mode(monkeypatch: pytest.MonkeyPatch): +def test_offline_mode(monkeypatch: pytest.MonkeyPatch, vllm_runner): # Set HF to offline mode and ensure we can still construct an LLM with monkeypatch.context() as m: try: @@ -93,7 +97,8 @@ def disable_connect(*args, **kwargs): _re_import_modules() # Cached model files should be used in offline mode for model_config in MODEL_CONFIGS: - LLM(**model_config) + with _create_runner(vllm_runner, model_config): + pass finally: # Reset the environment after the test # NB: Assuming tests are run in online mode @@ -136,7 +141,7 @@ def _re_import_modules(): @pytest.mark.skip_global_cleanup @pytest.mark.usefixtures("cache_models") -def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch): +def test_model_from_huggingface_offline(monkeypatch: pytest.MonkeyPatch, vllm_runner): # Set HF to offline mode and ensure we can still construct an LLM with monkeypatch.context() as m: try: @@ -159,7 +164,8 @@ def disable_connect(*args, **kwargs): # Need to re-import huggingface_hub # and friends to set up offline mode _re_import_modules() - LLM(model="facebook/opt-125m") + with vllm_runner("facebook/opt-125m"): + pass finally: # Reset the environment after the test # NB: Assuming tests are run in online mode diff --git a/tests/entrypoints/llm/test_chat.py b/tests/entrypoints/llm/test_chat.py index cbc57b80da48..6ee8d34ce138 100644 --- a/tests/entrypoints/llm/test_chat.py +++ b/tests/entrypoints/llm/test_chat.py @@ -4,46 +4,36 @@ import pytest -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams @pytest.fixture(scope="function") -def text_llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() +def text_llm(vllm_runner): + with vllm_runner( + "meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0 + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="function") -def llm_for_failure_test(): +def llm_for_failure_test(vllm_runner): """ Fixture for testing issue #26081. Uses a small max_model_len to easily trigger length errors. """ - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model="meta-llama/Llama-3.2-1B-Instruct", + with vllm_runner( + "meta-llama/Llama-3.2-1B-Instruct", enforce_eager=True, seed=0, max_model_len=128, disable_log_stats=True, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) def test_chat(text_llm): @@ -99,21 +89,16 @@ def test_llm_chat_tokenization_no_double_bos(text_llm): @pytest.fixture(scope="function") -def thinking_llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model="Qwen/Qwen3-0.6B", +def thinking_llm(vllm_runner): + with vllm_runner( + "Qwen/Qwen3-0.6B", max_model_len=4096, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.parametrize("enable_thinking", [True, False]) diff --git a/tests/entrypoints/llm/test_collective_rpc.py b/tests/entrypoints/llm/test_collective_rpc.py index d66455889368..36f1436e327f 100644 --- a/tests/entrypoints/llm/test_collective_rpc.py +++ b/tests/entrypoints/llm/test_collective_rpc.py @@ -4,15 +4,13 @@ import pytest import torch -from vllm import LLM - from ...utils import create_new_process_for_each_test @pytest.mark.parametrize("tp_size", [1, 2]) @pytest.mark.parametrize("backend", ["mp", "ray"]) @create_new_process_for_each_test() -def test_collective_rpc(tp_size, backend, monkeypatch): +def test_collective_rpc(tp_size, backend, monkeypatch, vllm_runner): if torch.accelerator.device_count() < tp_size: pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}") if tp_size == 1 and backend == "ray": @@ -26,11 +24,11 @@ def echo_rank(self): return self.rank monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") - llm = LLM( - model="hmellor/tiny-random-LlamaForCausalLM", + with vllm_runner( + "hmellor/tiny-random-LlamaForCausalLM", enforce_eager=True, load_format="dummy", tensor_parallel_size=tp_size, distributed_executor_backend=backend, - ) - assert llm.collective_rpc(echo_rank) == list(range(tp_size)) + ) as runner: + assert runner.llm.collective_rpc(echo_rank) == list(range(tp_size)) diff --git a/tests/entrypoints/llm/test_generate.py b/tests/entrypoints/llm/test_generate.py index 82f38adfb772..e79e62b3619d 100644 --- a/tests/entrypoints/llm/test_generate.py +++ b/tests/entrypoints/llm/test_generate.py @@ -6,7 +6,6 @@ import pytest from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "distilbert/distilgpt2" @@ -26,22 +25,17 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, max_num_batched_tokens=4096, tensor_parallel_size=1, gpu_memory_utilization=0.10, enforce_eager=True, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup @@ -97,34 +91,34 @@ def test_single_prompt_priority(llm: LLM): assert len(outputs) == 1 -def test_max_model_len(): +def test_max_model_len(vllm_runner): max_model_len = 20 - llm = LLM( - model=MODEL_NAME, + with vllm_runner( + MODEL_NAME, max_model_len=max_model_len, gpu_memory_utilization=0.10, enforce_eager=True, # reduce test time - ) - sampling_params = SamplingParams(max_tokens=max_model_len + 10) - outputs = llm.generate(PROMPTS, sampling_params) - for output in outputs: - num_total_tokens = len(output.prompt_token_ids) + len( - output.outputs[0].token_ids - ) - # Total tokens must not exceed max_model_len. - # It can be less if generation finishes due to other reasons (e.g., EOS) - # before reaching the absolute model length limit. - assert num_total_tokens <= max_model_len - - -def test_log_stats(): - llm = LLM( - model=MODEL_NAME, + ) as runner: + sampling_params = SamplingParams(max_tokens=max_model_len + 10) + outputs = runner.llm.generate(PROMPTS, sampling_params) + for output in outputs: + num_total_tokens = len(output.prompt_token_ids) + len( + output.outputs[0].token_ids + ) + # Total tokens must not exceed max_model_len. + # It can be less if generation finishes due to other reasons (e.g., EOS) + # before reaching the absolute model length limit. + assert num_total_tokens <= max_model_len + + +def test_log_stats(vllm_runner): + with vllm_runner( + MODEL_NAME, disable_log_stats=False, gpu_memory_utilization=0.10, enforce_eager=True, # reduce test time - ) - outputs = llm.generate(PROMPTS, sampling_params=None) + ) as runner: + outputs = runner.llm.generate(PROMPTS, sampling_params=None) - # disable_log_stats is False, every output should have metrics - assert all(output.metrics is not None for output in outputs) + # disable_log_stats is False, every output should have metrics + assert all(output.metrics is not None for output in outputs) diff --git a/tests/entrypoints/llm/test_gpu_utilization.py b/tests/entrypoints/llm/test_gpu_utilization.py index 896091533ad2..fc24a73dd410 100644 --- a/tests/entrypoints/llm/test_gpu_utilization.py +++ b/tests/entrypoints/llm/test_gpu_utilization.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm import LLM, SamplingParams +from vllm import SamplingParams -def test_gpu_memory_utilization(): +def test_gpu_memory_utilization(vllm_runner): prompts = [ "Hello, my name is", "The president of the United States is", @@ -15,13 +15,26 @@ def test_gpu_memory_utilization(): # makes sure gpu_memory_utilization is per-instance limit, # not a global limit - llms = [ - LLM(model="facebook/opt-125m", gpu_memory_utilization=0.3, enforce_eager=True) - for i in range(3) - ] - for llm in llms: - outputs = llm.generate(prompts, sampling_params) - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + with ( + vllm_runner( + "facebook/opt-125m", + gpu_memory_utilization=0.3, + enforce_eager=True, + ) as runner_0, + vllm_runner( + "facebook/opt-125m", + gpu_memory_utilization=0.3, + enforce_eager=True, + ) as runner_1, + vllm_runner( + "facebook/opt-125m", + gpu_memory_utilization=0.3, + enforce_eager=True, + ) as runner_2, + ): + for runner in (runner_0, runner_1, runner_2): + outputs = runner.llm.generate(prompts, sampling_params) + for output in outputs: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") diff --git a/tests/entrypoints/llm/test_prompt_validation.py b/tests/entrypoints/llm/test_prompt_validation.py index 8dd55c6b10e0..6000c03da973 100644 --- a/tests/entrypoints/llm/test_prompt_validation.py +++ b/tests/entrypoints/llm/test_prompt_validation.py @@ -4,30 +4,35 @@ import pytest import torch -from vllm import LLM from vllm.exceptions import VLLMValidationError -def test_empty_prompt(): - llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"): - llm.generate([""]) - - -def test_out_of_vocab_token(): - llm = LLM(model="openai-community/gpt2", enforce_eager=True) - with pytest.raises(VLLMValidationError, match="out of vocabulary"): - llm.generate({"prompt_token_ids": [999999]}) - - -def test_require_mm_embeds(): - llm = LLM( - model="llava-hf/llava-1.5-7b-hf", - enforce_eager=True, - enable_mm_embeds=False, - ) - with pytest.raises(ValueError, match="--enable-mm-embeds"): - llm.generate( +def test_empty_prompt(vllm_runner): + with ( + vllm_runner("openai-community/gpt2", enforce_eager=True) as runner, + pytest.raises(VLLMValidationError, match="decoder prompt cannot be empty"), + ): + runner.llm.generate([""]) + + +def test_out_of_vocab_token(vllm_runner): + with ( + vllm_runner("openai-community/gpt2", enforce_eager=True) as runner, + pytest.raises(VLLMValidationError, match="out of vocabulary"), + ): + runner.llm.generate({"prompt_token_ids": [999999]}) + + +def test_require_mm_embeds(vllm_runner): + with ( + vllm_runner( + "llava-hf/llava-1.5-7b-hf", + enforce_eager=True, + enable_mm_embeds=False, + ) as runner, + pytest.raises(ValueError, match="--enable-mm-embeds"), + ): + runner.llm.generate( { "prompt": "", "multi_modal_data": {"image": torch.empty(1, 1, 1)}, diff --git a/tests/entrypoints/llm/test_struct_output_generate.py b/tests/entrypoints/llm/test_struct_output_generate.py index 219ee7cd3875..0f55ad7adc3e 100644 --- a/tests/entrypoints/llm/test_struct_output_generate.py +++ b/tests/entrypoints/llm/test_struct_output_generate.py @@ -9,13 +9,11 @@ import jsonschema import pytest import regex as re -import torch from pydantic import BaseModel from tests.reasoning.utils import run_reasoning_extraction from vllm.config import StructuredOutputsConfig -from vllm.distributed import cleanup_dist_env_and_memory -from vllm.entrypoints.llm import LLM +from vllm.exceptions import VLLMValidationError from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager @@ -212,11 +210,11 @@ class CarDescription(BaseModel): PARAMS_MODELS_BACKENDS_TOKENIZER_MODE, ) def test_structured_output( - request: pytest.FixtureRequest, backend: str, tokenizer_mode: str, model_name: str, speculative_config: dict[str, Any], + vllm_runner, ): sample_json_schema = SAMPLE_JSON_SCHEMA unsupported_json_schema = UNSUPPORTED_JSON_SCHEMA @@ -229,8 +227,8 @@ def test_structured_output( # Use a single LLM instance for several scenarios to # speed up the test suite. - llm = LLM( - model=model_name, + with vllm_runner( + model_name, enforce_eager=True, max_model_len=1024, structured_outputs_config=dict( @@ -242,181 +240,307 @@ def test_structured_output( config_format="auto" if not model_name.startswith("mistralai/") else "hf", speculative_config=speculative_config, **platform_args, - ) - request.addfinalizer(llm.llm_engine.engine_core.shutdown) - - # - # Test 1: Generate JSON output based on a provided schema - # - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=4096, - structured_outputs=StructuredOutputsParams(json=sample_json_schema), - ) - - prompt = ( - "Give an example JSON for an employee profile that fits this " - "schema. Make the response as short as possible. Schema: " - f"{sample_json_schema}" - ) - outputs = llm.generate( - [prompt] * 2, - sampling_params=sampling_params, - use_tqdm=True, - ) - - assert outputs is not None - - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt - - generated_text = output.outputs[0].text - assert generated_text is not None - if backend != "lm-format-enforcer": - assert "\n" not in generated_text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - try: - output_json = json.loads(generated_text) - except json.JSONDecodeError as e: - pytest.fail( - f"Invalid JSON from backend={backend}: {generated_text!r}\n" - f"Schema: {sample_json_schema}\nError: {e}" - ) - jsonschema.validate(instance=output_json, schema=sample_json_schema) - - # - # Test 2: Generate JSON object without a schema - # - if backend != "outlines": + ) as runner: + # + # Test 1: Generate JSON output based on a provided schema + # sampling_params = SamplingParams( temperature=1.0, max_tokens=4096, - n=2, - structured_outputs=StructuredOutputsParams(json_object=True), + structured_outputs=StructuredOutputsParams(json=sample_json_schema), ) - outputs = llm.generate( - prompts=( - "Generate a JSON object with curly braces for a person with " - "name and age fields for John Smith who is 31 years old. " - "Make the response as short as possible." - ), + prompt = ( + "Give an example JSON for an employee profile that fits this " + "schema. Make the response as short as possible. Schema: " + f"{sample_json_schema}" + ) + outputs = runner.llm.generate( + [prompt] * 2, sampling_params=sampling_params, use_tqdm=True, ) assert outputs is not None + for output in outputs: assert output is not None assert isinstance(output, RequestOutput) + prompt = output.prompt - for i in range(2): - generated_text = output.outputs[i].text - print(generated_text) - assert generated_text is not None + generated_text = output.outputs[0].text + assert generated_text is not None + if backend != "lm-format-enforcer": + assert "\n" not in generated_text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + try: + output_json = json.loads(generated_text) + except json.JSONDecodeError as e: + pytest.fail( + f"Invalid JSON from backend={backend}: {generated_text!r}\n" + f"Schema: {sample_json_schema}\nError: {e}" + ) + jsonschema.validate(instance=output_json, schema=sample_json_schema) - # Parse to verify it is a valid JSON object - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) + # + # Test 2: Generate JSON object without a schema + # + if backend != "outlines": + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=4096, + n=2, + structured_outputs=StructuredOutputsParams(json_object=True), + ) + + outputs = runner.llm.generate( + prompts=( + "Generate a JSON object with curly braces for a person with " + "name and age fields for John Smith who is 31 years old. " + "Make the response as short as possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + + for i in range(2): + generated_text = output.outputs[i].text + print(generated_text) + assert generated_text is not None - # - # Test 3: test a jsonschema incompatible with xgrammar - # - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=4096, - structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), - ) - if backend.startswith("xgrammar"): - with pytest.raises( - ValueError, - match="The provided JSON schema contains features " - "not supported by xgrammar.", - ): + # Parse to verify it is a valid JSON object + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) + + # + # Test 3: test a jsonschema incompatible with xgrammar + # + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=4096, + structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), + ) + if backend.startswith("xgrammar"): + with pytest.raises( + VLLMValidationError, + match="The provided JSON schema contains features " + "not supported by xgrammar.", + ): + prompt = ( + f"Give an example JSON for an employee profile that " + f"fits this schema: {unsupported_json_schema}. " + f"Make the response as short as possible." + ) + runner.llm.generate( + [prompt] * 2, + sampling_params=sampling_params, + use_tqdm=True, + ) + else: prompt = ( - f"Give an example JSON for an employee profile that " + f"Give an example JSON object for a grade that " f"fits this schema: {unsupported_json_schema}. " f"Make the response as short as possible." ) - llm.generate( - [prompt] * 2, + outputs = runner.llm.generate( + prompt, + sampling_params=sampling_params, + use_tqdm=True, + ) + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + generated_text = output.outputs[0].text + assert generated_text is not None + print(generated_text) + + # Parse to verify it is valid JSON + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) + + if backend not in ["outlines", "lm-format-enforcer"]: + # + # Test 4: Generate SQL statement using EBNF grammar + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf), + ) + outputs = runner.llm.generate( + ( + "Generate a sql statement that selects col_1 from " + "table_1 where it is equal to 1. Make the response as short as " + "possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + prompt = output.prompt + + generated_text = output.outputs[0].text + assert generated_text is not None + + # remove spaces for comparison b/c we removed them in the grammar + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace( + " ", "" + ) + + assert generated_text.strip() == ground_truth + + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # + # Test 5: Generate SQL statement using Lark grammar + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark), + ) + outputs = runner.llm.generate( + ( + "Generate a sql statement that selects col_1 from " + "table_1 where it is equal to 1. Make the response as short as " + "possible." + ), sampling_params=sampling_params, use_tqdm=True, ) - else: + + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + prompt = output.prompt + + generated_text = output.outputs[0].text + assert generated_text is not None + + # use Lark to parse the output, and make sure it's a valid parse tree + from lark import Lark + + parser = Lark(sample_sql_lark) + parser.parse(generated_text) + + # remove spaces for comparison b/c we removed them in the grammar + ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace( + " ", "" + ) + + assert generated_text.strip() == ground_truth + + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + + # + # Test 6: Test invalid grammar input + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(grammar="not a grammar"), + ) + with pytest.raises( + VLLMValidationError, match="Failed to convert the grammar " + ): + runner.llm.generate( + ( + "Generate a sql statement that selects col_1 from " + "table_1 where it is equal to 1. Make the response as short " + "as possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) + + # + # Test 7: Generate text based on a regex pattern + # + sampling_params = SamplingParams( + temperature=0.8, + top_p=0.95, + structured_outputs=StructuredOutputsParams(regex=sample_regex), + ) + prompt = ( - f"Give an example JSON object for a grade that " - f"fits this schema: {unsupported_json_schema}. " + f"Give an example IPv4 address with this regex: {sample_regex}. " f"Make the response as short as possible." ) - outputs = llm.generate( - prompt, + outputs = runner.llm.generate( + [prompt] * 2, sampling_params=sampling_params, use_tqdm=True, ) + assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) + prompt = output.prompt generated_text = output.outputs[0].text - assert generated_text is not None print(generated_text) + assert generated_text is not None + assert re.fullmatch(sample_regex, generated_text) is not None + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - # Parse to verify it is valid JSON - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) - - if backend not in ["outlines", "lm-format-enforcer"]: # - # Test 4: Generate SQL statement using EBNF grammar + # Test 8: Generate text based on a choices # sampling_params = SamplingParams( temperature=0.8, top_p=0.95, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(grammar=sample_sql_ebnf), + structured_outputs=StructuredOutputsParams( + choice=sample_structured_outputs_choices + ), ) - outputs = llm.generate( + + outputs = runner.llm.generate( ( - "Generate a sql statement that selects col_1 from " - "table_1 where it is equal to 1. Make the response as short as " - "possible." + "The best language for type-safe systems programming is " + "(Make the response as short as possible.) " ), sampling_params=sampling_params, use_tqdm=True, ) - assert outputs is not None for output in outputs: assert output is not None assert isinstance(output, RequestOutput) prompt = output.prompt - generated_text = output.outputs[0].text + print(generated_text) assert generated_text is not None - - # remove spaces for comparison b/c we removed them in the grammar - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "") - - assert generated_text.strip() == ground_truth - + assert generated_text in sample_structured_outputs_choices print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") # - # Test 5: Generate SQL statement using Lark grammar + # Test 9: Generate structured output using a Pydantic model with an enum # + json_schema = CarDescription.model_json_schema() sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, + temperature=1.0, max_tokens=1000, - structured_outputs=StructuredOutputsParams(grammar=sample_sql_lark), + structured_outputs=StructuredOutputsParams(json=json_schema), ) - outputs = llm.generate( + + outputs = runner.llm.generate( ( - "Generate a sql statement that selects col_1 from " - "table_1 where it is equal to 1. Make the response as short as " + "Generate a JSON with the brand, model and car_type of the most " + "iconic car from the 90's. Make the response as short as " "possible." ), sampling_params=sampling_params, @@ -424,6 +548,7 @@ def test_structured_output( ) assert outputs is not None + for output in outputs: assert output is not None assert isinstance(output, RequestOutput) @@ -431,220 +556,97 @@ def test_structured_output( generated_text = output.outputs[0].text assert generated_text is not None - - # use Lark to parse the output, and make sure it's a valid parse tree - from lark import Lark - - parser = Lark(sample_sql_lark) - parser.parse(generated_text) - - # remove spaces for comparison b/c we removed them in the grammar - ground_truth = "SELECT col_1 from table_1 where col_1 = 1".replace(" ", "") - - assert generated_text.strip() == ground_truth - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + try: + output_json = json.loads(generated_text) + except json.JSONDecodeError as e: + pytest.fail( + f"Invalid JSON from backend={backend}: {generated_text!r}\n" + f"Schema: {json_schema}\nError: {e}" + ) + jsonschema.validate(instance=output_json, schema=json_schema) # - # Test 6: Test invalid grammar input + # Test 10: Generate structured with minLength and maxLength # + min_length = 50 + max_length = 50 + json_schema = { + "type": "object", + "properties": { + "description": { + "type": "string", + "maxLength": max_length, + "minLength": min_length, + } + }, + "required": ["description"], + "additionalProperties": False, + } + sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(grammar="not a grammar"), + temperature=1.0, + max_tokens=4096, + structured_outputs=StructuredOutputsParams(json=json_schema), ) - with pytest.raises(ValueError, match="Failed to convert the grammar "): - llm.generate( - ( - "Generate a sql statement that selects col_1 from " - "table_1 where it is equal to 1. Make the response as short " - "as possible." - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - - # - # Test 7: Generate text based on a regex pattern - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - structured_outputs=StructuredOutputsParams(regex=sample_regex), - ) - - prompt = ( - f"Give an example IPv4 address with this regex: {sample_regex}. " - f"Make the response as short as possible." - ) - outputs = llm.generate( - [prompt] * 2, - sampling_params=sampling_params, - use_tqdm=True, - ) - - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - print(generated_text) - assert generated_text is not None - assert re.fullmatch(sample_regex, generated_text) is not None - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - # - # Test 8: Generate text based on a choices - # - sampling_params = SamplingParams( - temperature=0.8, - top_p=0.95, - structured_outputs=StructuredOutputsParams( - choice=sample_structured_outputs_choices - ), - ) - outputs = llm.generate( - ( - "The best language for type-safe systems programming is " - "(Make the response as short as possible.) " - ), - sampling_params=sampling_params, - use_tqdm=True, - ) - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - print(generated_text) - assert generated_text is not None - assert generated_text in sample_structured_outputs_choices - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - # - # Test 9: Generate structured output using a Pydantic model with an enum - # - json_schema = CarDescription.model_json_schema() - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(json=json_schema), - ) - - outputs = llm.generate( - ( - "Generate a JSON with the brand, model and car_type of the most " - "iconic car from the 90's. Make the response as short as " - "possible." - ), - sampling_params=sampling_params, - use_tqdm=True, - ) + outputs = runner.llm.generate( + ( + "Generate a description of a frog using 50 characters. " + "Make the response as short as possible." + ), + sampling_params=sampling_params, + use_tqdm=True, + ) - assert outputs is not None + assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + prompt = output.prompt - generated_text = output.outputs[0].text - assert generated_text is not None - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - try: - output_json = json.loads(generated_text) - except json.JSONDecodeError as e: - pytest.fail( - f"Invalid JSON from backend={backend}: {generated_text!r}\n" - f"Schema: {json_schema}\nError: {e}" - ) - jsonschema.validate(instance=output_json, schema=json_schema) - - # - # Test 10: Generate structured with minLength and maxLength - # - min_length = 50 - max_length = 50 - json_schema = { - "type": "object", - "properties": { - "description": { - "type": "string", - "maxLength": max_length, - "minLength": min_length, + generated_text = output.outputs[0].text + assert generated_text is not None + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + try: + output_json = json.loads(generated_text) + except json.JSONDecodeError as e: + pytest.fail( + f"Invalid JSON from backend={backend}: {generated_text!r}\n" + f"Schema: {json_schema}\nError: {e}" + ) + jsonschema.validate(instance=output_json, schema=json_schema) + + if backend not in ["outlines", "lm-format-enforcer"]: + # + # Test 11: Generate structured output using structural_tag format + # + structural_tag_config = { + "type": "structural_tag", + "structures": [ + { + "begin": "", + "schema": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "additionalProperties": False, + }, + "end": "", + } + ], + "triggers": ["", - "schema": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "additionalProperties": False, - }, - "end": "", - } - ], - "triggers": ["(.*?)" - matches = re.findall(function_call_pattern, generated_text) + # Change this once other backends support structural_tag + outputs = runner.llm.generate( + prompt, sampling_params=sampling_params, use_tqdm=True + ) + assert outputs is not None - if not matches: - print( - f"Warning: No function calls found in response: {generated_text!r}" - ) - continue + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + generated_text = output.outputs[0].text + assert generated_text is not None - # Take the first function call if multiple are found - json_str = matches[0] - try: - json_content = json.loads(json_str) - assert "city" in json_content - assert isinstance(json_content["city"], str) - print(f"Found valid function call: {generated_text!r}") - except (json.JSONDecodeError, AssertionError) as e: - pytest.fail( - f"Invalid function call format: {generated_text!r}\nError: {str(e)}" - ) + # Search for function call pattern in the response + function_call_pattern = r"(.*?)" + matches = re.findall(function_call_pattern, generated_text) + + if not matches: + print( + f"Warning: No function calls found in response: {generated_text!r}" + ) + continue + + # Take the first function call if multiple are found + json_str = matches[0] + try: + json_content = json.loads(json_str) + assert "city" in json_content + assert isinstance(json_content["city"], str) + print(f"Found valid function call: {generated_text!r}") + except (json.JSONDecodeError, AssertionError) as e: + pytest.fail( + f"Invalid function call format: {generated_text!r}\nError: {str(e)}" + ) @pytest.mark.parametrize( @@ -738,14 +742,15 @@ def test_structured_output_with_reasoning_matrices( model_name: str, speculative_config: dict[str, Any] | None, async_scheduling: bool, + vllm_runner, ): if current_platform.is_tpu() and speculative_config: pytest.skip("TPU does not support speculative decoding") # Use a single LLM instance for several scenarios to # speed up the test suite. - llm = LLM( - model=model_name, + with vllm_runner( + model_name, # Don't use eager execution on TPUs because we want to test for no # recompilation at runtime enforce_eager=bool(not current_platform.is_tpu()), @@ -759,271 +764,273 @@ def test_structured_output_with_reasoning_matrices( tokenizer_mode=tokenizer_mode, speculative_config=speculative_config, async_scheduling=async_scheduling, - ) - tokenizer = llm.get_tokenizer() - reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)( - tokenizer=tokenizer - ) - - reasoning_prompt = "Solve the following math problem step-by-step, then provide the final answer as JSON object with a single key 'result'. Make sure to correct your reasoning if there are any issue should it arise.\nProblem: What is 5 * 8 + 2?" # noqa: E501 - reasoning_schema = { - "type": "object", - "properties": {"result": {"type": "integer"}}, - "required": ["result"], - "additionalProperties": False, - } - if "Qwen3" in model_name: - reasoning_prompt += "\n" - - sampling_params = SamplingParams( - temperature=0.1, - max_tokens=8192, - structured_outputs=StructuredOutputsParams(json=reasoning_schema), - ) - outputs = llm.generate( - [reasoning_prompt], - sampling_params=sampling_params, - use_tqdm=True, - ) - - assert outputs is not None - output = outputs[0] - assert output is not None and isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - reasoning, content = run_reasoning_extraction(reasoner, [generated_text]) - print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}") - - if "Qwen3" in model_name: - assert content is not None - - assert reasoning is not None - - if content is not None: - output_json = json.loads(content) - jsonschema.validate(instance=output_json, schema=reasoning_schema) + ) as runner: + tokenizer = runner.llm.get_tokenizer() + reasoner = ReasoningParserManager.get_reasoning_parser(reasoning_parser)( + tokenizer=tokenizer + ) + + reasoning_prompt = "Solve the following math problem step-by-step, then provide the final answer as JSON object with a single key 'result'. Make sure to correct your reasoning if there are any issue should it arise.\nProblem: What is 5 * 8 + 2?" # noqa: E501 + reasoning_schema = { + "type": "object", + "properties": {"result": {"type": "integer"}}, + "required": ["result"], + "additionalProperties": False, + } + if "Qwen3" in model_name: + reasoning_prompt += "\n" + + sampling_params = SamplingParams( + temperature=0.1, + max_tokens=8192, + structured_outputs=StructuredOutputsParams(json=reasoning_schema), + ) + outputs = runner.llm.generate( + [reasoning_prompt], + sampling_params=sampling_params, + use_tqdm=True, + ) + + assert outputs is not None + output = outputs[0] + assert output is not None and isinstance(output, RequestOutput) + prompt = output.prompt + generated_text = output.outputs[0].text + reasoning, content = run_reasoning_extraction(reasoner, [generated_text]) + print(f"Prompt: {prompt!r}\nReasoning: {reasoning!r}\nContent: {content!r}") + + if "Qwen3" in model_name: + assert content is not None + + assert reasoning is not None + + if content is not None: + output_json = json.loads(content) + jsonschema.validate(instance=output_json, schema=reasoning_schema) @pytest.mark.parametrize("model_name, tokenizer_mode", PARAMS_MODELS_TOKENIZER_MODE) def test_structured_output_auto_mode( model_name: str, tokenizer_mode: str, + vllm_runner, ): unsupported_json_schema = UNSUPPORTED_JSON_SCHEMA - llm = LLM( - model=model_name, + with vllm_runner( + model_name, max_model_len=1024, structured_outputs_config=dict(backend="auto"), tokenizer_mode=tokenizer_mode, load_format="auto", config_format="auto", - ) - - sampling_params = SamplingParams( - temperature=1.0, - max_tokens=1000, - structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), - ) - - prompts = ( - "Give an example JSON object for a grade " - "that fits this schema: " - f"{unsupported_json_schema}. Make the response as short as possible." - ) - # This would fail with the default of "xgrammar", but in "auto" - # we will handle fallback automatically. - outputs = llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True) - # Make sure `auto` backend handling doesn't mess up sampling_params - # and that we can reuse it without error. - outputs.extend( - llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True) - ) - - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - generated_text = output.outputs[0].text - assert generated_text is not None - print(generated_text) + ) as runner: + sampling_params = SamplingParams( + temperature=1.0, + max_tokens=1000, + structured_outputs=StructuredOutputsParams(json=unsupported_json_schema), + ) + + prompts = ( + "Give an example JSON object for a grade " + "that fits this schema: " + f"{unsupported_json_schema}. Make the response as short as possible." + ) + # This would fail with the default of "xgrammar", but in "auto" + # we will handle fallback automatically. + outputs = runner.llm.generate( + prompts, sampling_params=sampling_params, use_tqdm=True + ) + # Make sure `auto` backend handling doesn't mess up sampling_params + # and that we can reuse it without error. + outputs.extend( + runner.llm.generate(prompts, sampling_params=sampling_params, use_tqdm=True) + ) + + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + generated_text = output.outputs[0].text + assert generated_text is not None + print(generated_text) - # Parse to verify it is valid JSON - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) + # Parse to verify it is valid JSON + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) -def test_guidance_no_additional_properties(): - llm = LLM( - model="Qwen/Qwen2.5-1.5B-Instruct", +def test_guidance_no_additional_properties(vllm_runner): + with vllm_runner( + "Qwen/Qwen2.5-1.5B-Instruct", max_model_len=1024, structured_outputs_config=dict( backend="guidance", disable_any_whitespace=True, disable_additional_properties=True, ), - ) - - schema = { - "type": "object", - "properties": { - "a1": {"type": "string"}, - "a2": {"type": "string"}, - "a3": {"type": "string"}, - }, - "required": ["a1", "a2", "a3"], - } + ) as runner: + schema = { + "type": "object", + "properties": { + "a1": {"type": "string"}, + "a2": {"type": "string"}, + "a3": {"type": "string"}, + }, + "required": ["a1", "a2", "a3"], + } - prompt = ( - "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a " - "helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a " - "large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. " - "Make the response as short as possible." - "<|im_end|>\n<|im_start|>assistant\n" - ) - - def generate_with_backend(backend): - structured_outputs_params = StructuredOutputsParams( - json=schema, - backend=backend, - disable_any_whitespace=True, - disable_additional_properties=True, - ) - sampling_params = SamplingParams( - temperature=0, max_tokens=256, structured_outputs=structured_outputs_params + prompt = ( + "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a " + "helpful assistant.<|im_end|>\n<|im_start|>user\nPlease generate a " + "large JSON object with key-value pairs a1=b1, a2=b2, ..., a20=b20. " + "Make the response as short as possible." + "<|im_end|>\n<|im_start|>assistant\n" ) - outputs = llm.generate(prompt, sampling_params=sampling_params) - assert outputs is not None - generated_text = outputs[0].outputs[0].text - assert generated_text is not None - parsed_json = json.loads(generated_text) - assert isinstance(parsed_json, dict) - jsonschema.validate(instance=parsed_json, schema=schema) - return parsed_json - - generated = generate_with_backend("guidance") - assert "a1" in generated - assert "a2" in generated - assert "a3" in generated - assert "a4" not in generated - assert "a5" not in generated - assert "a6" not in generated + def generate_with_backend(backend): + structured_outputs_params = StructuredOutputsParams( + json=schema, + backend=backend, + disable_any_whitespace=True, + disable_additional_properties=True, + ) + sampling_params = SamplingParams( + temperature=0, + max_tokens=256, + structured_outputs=structured_outputs_params, + ) + + outputs = runner.llm.generate(prompt, sampling_params=sampling_params) + assert outputs is not None + generated_text = outputs[0].outputs[0].text + assert generated_text is not None + parsed_json = json.loads(generated_text) + assert isinstance(parsed_json, dict) + jsonschema.validate(instance=parsed_json, schema=schema) + return parsed_json + + generated = generate_with_backend("guidance") + assert "a1" in generated + assert "a2" in generated + assert "a3" in generated + assert "a4" not in generated + assert "a5" not in generated + assert "a6" not in generated @pytest.mark.parametrize("backend", ["guidance", "xgrammar", "outlines"]) def test_structured_output_batched_with_non_structured_outputs_requests( backend: str, + vllm_runner, ): sample_json_schema = SAMPLE_JSON_SCHEMA # Don't use eager execution on TPUs because we want to test for no # recompilation at runtime enforce_eager = bool(not current_platform.is_tpu()) - llm = LLM( - model="meta-llama/Meta-Llama-3.1-8B-Instruct", + with vllm_runner( + "meta-llama/Meta-Llama-3.1-8B-Instruct", enforce_eager=enforce_eager, max_model_len=1024, structured_outputs_config=StructuredOutputsConfig( backend=backend, disable_any_whitespace=backend in {"xgrammar", "guidance"}, ), - ) - - structured_outputs_prompt = ( - "Give an example JSON for an employee profile that fits this " - "schema. Make the response as short as possible. Schema: " - f"{sample_json_schema}" - ) - - non_structured_outputs_prompt = "The diameter of the Earth in kilometers is " - - prompts = [structured_outputs_prompt, non_structured_outputs_prompt] - sampling_params = [ - SamplingParams( - temperature=1.0, - max_tokens=400, - structured_outputs=StructuredOutputsParams(json=sample_json_schema), - ), - # No max tokens, temp=0 to assert on contents - SamplingParams( - seed=42, - temperature=0, - top_p=1.0, - ), - ] + ) as runner: + structured_outputs_prompt = ( + "Give an example JSON for an employee profile that fits this " + "schema. Make the response as short as possible. Schema: " + f"{sample_json_schema}" + ) - outputs = llm.generate( - prompts=prompts, sampling_params=sampling_params, use_tqdm=True - ) + non_structured_outputs_prompt = "The diameter of the Earth in kilometers is " - assert outputs is not None + prompts = [structured_outputs_prompt, non_structured_outputs_prompt] + sampling_params = [ + SamplingParams( + temperature=1.0, + max_tokens=400, + structured_outputs=StructuredOutputsParams(json=sample_json_schema), + ), + # No max tokens, temp=0 to assert on contents + SamplingParams( + seed=42, + temperature=0, + top_p=1.0, + ), + ] - # Free memory as soon as possible as failed assertions - # will short circuit and not free up memory - del llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + outputs = runner.llm.generate( + prompts=prompts, sampling_params=sampling_params, use_tqdm=True + ) - for index, output in enumerate(outputs): - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt + assert outputs is not None - generated_text = output.outputs[0].text - assert generated_text is not None - print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}") + for index, output in enumerate(outputs): + assert output is not None + assert isinstance(output, RequestOutput) + prompt = output.prompt - if index == 0: - # First prompt is structured outputs, expect valid JSON - assert "\n" not in generated_text - output_json = json.loads(generated_text) - jsonschema.validate(instance=output_json, schema=sample_json_schema) - else: - # Second prompt is not structured outputs, expect valid output - # Cannot assert on exact output, but we can expect it to be factual - assert "12,742" in generated_text + generated_text = output.outputs[0].text + assert generated_text is not None + print(f"Prompt:\n{prompt!r}\nGenerated text:\n{generated_text!r}") - # non-structured outputs requests should not return a valid JSON here - with pytest.raises(ValueError): + if index == 0: + # First prompt is structured outputs, expect valid JSON + assert "\n" not in generated_text output_json = json.loads(generated_text) + jsonschema.validate(instance=output_json, schema=sample_json_schema) + else: + # Second prompt is not structured outputs, expect valid output + # Cannot assert on exact output, but we can expect it to be factual + assert "12,742" in generated_text + + # non-structured outputs requests should not return a valid JSON here + with pytest.raises(ValueError): + output_json = json.loads(generated_text) @pytest.mark.parametrize("backend", ["xgrammar"]) -def test_structured_output_with_structural_tag(backend: str): - llm = LLM( - model="Qwen/Qwen2.5-1.5B-Instruct", +def test_structured_output_with_structural_tag(backend: str, vllm_runner): + with vllm_runner( + "Qwen/Qwen2.5-1.5B-Instruct", structured_outputs_config=StructuredOutputsConfig(backend=backend), - ) - - structural_tag_config = { - "type": "structural_tag", - "format": { - "type": "triggered_tags", - "tags": [ - {"begin": "hello_flag", "content": {"type": "any_text"}, "end": "hello"} - ], - "triggers": ["hello"], - "stop_after_first": False, - }, - } + ) as runner: + structural_tag_config = { + "type": "structural_tag", + "format": { + "type": "triggered_tags", + "tags": [ + { + "begin": "hello_flag", + "content": {"type": "any_text"}, + "end": "hello", + } + ], + "triggers": ["hello"], + "stop_after_first": False, + }, + } - sampling_params = SamplingParams( - temperature=0.0, - max_tokens=500, - structured_outputs=StructuredOutputsParams( - structural_tag=json.dumps(structural_tag_config) - ), - ) - - prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start" - outputs = llm.generate(prompt, sampling_params=sampling_params, use_tqdm=True) - assert outputs is not None - for output in outputs: - assert output is not None - assert isinstance(output, RequestOutput) - prompt = output.prompt - generated_text = output.outputs[0].text - assert generated_text is not None - assert "hello_flag" in generated_text, ( - f"Expected 'hello_flag' to be in generated text, but got: {generated_text}" + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=500, + structured_outputs=StructuredOutputsParams( + structural_tag=json.dumps(structural_tag_config) + ), ) + + prompt = "Hello and repeat hello 10 times, do not say anything else. Only say hello hello hello, now start" + outputs = runner.llm.generate( + prompt, sampling_params=sampling_params, use_tqdm=True + ) + assert outputs is not None + for output in outputs: + assert output is not None + assert isinstance(output, RequestOutput) + prompt = output.prompt + generated_text = output.outputs[0].text + assert generated_text is not None + assert "hello_flag" in generated_text, ( + f"Expected 'hello_flag' to be in generated text, but got: {generated_text}" + ) diff --git a/tests/entrypoints/multimodal/conftest.py b/tests/entrypoints/multimodal/conftest.py index 8003f1bf7dcb..f27c042d21e1 100644 --- a/tests/entrypoints/multimodal/conftest.py +++ b/tests/entrypoints/multimodal/conftest.py @@ -1,11 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Callable, Iterator -from contextlib import contextmanager -from typing import Any - -import pytest - # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA) TEST_IMAGE_ASSETS = [ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" @@ -13,70 +7,3 @@ "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png", "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png", ] - - -def _shutdown_llm(llm: Any, gpu_memory_utilization: float) -> None: - from vllm.distributed import cleanup_dist_env_and_memory - from vllm.platforms import current_platform - - try: - shutdown_timeout = 60.0 if current_platform.is_rocm() else None - llm.llm_engine.engine_core.shutdown(timeout=shutdown_timeout) - except Exception: - pass - - del llm - - try: - import torch - - torch._dynamo.reset() - except Exception: - pass - - cleanup_dist_env_and_memory() - - if current_platform.is_rocm(): - from tests.utils import wait_for_rocm_memory_to_settle - - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) - - -@contextmanager -def managed_llm(*args: Any, **kwargs: Any) -> Iterator[Any]: - from vllm import LLM - - llm = LLM(*args, **kwargs) - gpu_memory_utilization = ( - llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization - ) - try: - yield llm - finally: - _shutdown_llm(llm, gpu_memory_utilization) - - -def _make_managed_llm_factory() -> Iterator[Callable[..., Any]]: - from vllm import LLM - - llms: list[tuple[Any, float]] = [] - - def make_llm(*args: Any, **kwargs: Any) -> Any: - llm = LLM(*args, **kwargs) - gpu_memory_utilization = ( - llm.llm_engine.vllm_config.cache_config.gpu_memory_utilization - ) - llms.append((llm, gpu_memory_utilization)) - return llm - - try: - yield make_llm - finally: - while llms: - llm, gpu_memory_utilization = llms.pop() - _shutdown_llm(llm, gpu_memory_utilization) - - -@pytest.fixture -def multimodal_llm_factory() -> Iterator[Callable[..., Any]]: - yield from _make_managed_llm_factory() diff --git a/tests/entrypoints/multimodal/llm/test_chat.py b/tests/entrypoints/multimodal/llm/test_chat.py index 4de1f5cb80a0..1e19900ce838 100644 --- a/tests/entrypoints/multimodal/llm/test_chat.py +++ b/tests/entrypoints/multimodal/llm/test_chat.py @@ -1,21 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import weakref + import pytest from tests.entrypoints.multimodal.conftest import TEST_IMAGE_ASSETS @pytest.fixture(scope="function") -def vision_llm(multimodal_llm_factory): - return multimodal_llm_factory( - model="microsoft/Phi-3.5-vision-instruct", +def vision_llm(vllm_runner): + with vllm_runner( + "microsoft/Phi-3.5-vision-instruct", max_model_len=4096, max_num_seqs=5, enforce_eager=True, trust_remote_code=True, limit_mm_per_prompt={"image": 2}, seed=0, - ) + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.parametrize( diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py index 076a381f6cd4..5c378de35dbe 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_external_injection.py @@ -69,7 +69,7 @@ def test_inject_into_mm_cache( image_urls, mm_processor_cache_type, caplog_vllm, - multimodal_llm_factory, + vllm_runner, ): """Test that inject_into_mm_cache() injects pre-processed mm_kwargs into the processor cache and MM cache hit metrics are updated correctly. @@ -79,117 +79,115 @@ def test_inject_into_mm_cache( 2. Extract cached kwargs, call inject_into_mm_cache with a new hash, then generate with a pre-rendered input -> verifies injection works """ - llm = multimodal_llm_factory( - model="llava-hf/llava-1.5-7b-hf", + with vllm_runner( + "llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, enforce_eager=True, disable_log_stats=False, limit_mm_per_prompt={"image": 2}, mm_processor_cache_type=mm_processor_cache_type, - ) - - # Step 1: Normal requests to populate the cache - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0) - - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (2, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(50.0) - - # Step 2: Use a second image to get valid expanded tokens and - # placeholder positions via the renderer. - llm.chat(_make_messages(image_urls[1])) - queries_before = _get_mm_cache_stats(llm.get_metrics())[0] # 3 - - renderer = llm.llm_engine.renderer - cache = renderer.mm_processor_cache - assert cache is not None, "Processor cache should be enabled" - - _, eng_prompts = renderer.render_chat( - [_make_messages(image_urls[1])], - ChatParams(), - ) - eng_input = eng_prompts[0] - - # Inject pre-processed mm_kwargs with a NEW hash via public API - new_mm_hash = "deadbeef" * 8 - mm_hashes = {"image": [new_mm_hash]} - mm_kwargs = eng_input["mm_kwargs"] - - llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) - - # Build pre-rendered input (no externally_processed flag needed) - pre_rendered_input = { - "type": "multimodal", - "prompt_token_ids": eng_input["prompt_token_ids"], - "mm_kwargs": mm_kwargs, - "mm_hashes": mm_hashes, - "mm_placeholders": eng_input["mm_placeholders"], - } - - llm.generate( - pre_rendered_input, - sampling_params=SamplingParams(max_tokens=1), - ) - - # Verify cache was queried and injection happened - queries_after = _get_mm_cache_stats(llm.get_metrics())[0] - assert queries_after > queries_before, ( - "Cache should have been queried for the injected item" - ) - mm_rate = _get_mm_cache_log(llm, caplog_vllm) - assert mm_rate >= 0.0, "MM cache hit rate should be reported" + ) as runner: + # Step 1: Normal requests to populate the cache + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (1, 0) + + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (2, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(50.0) + + # Step 2: Use a second image to get valid expanded tokens and + # placeholder positions via the renderer. + runner.llm.chat(_make_messages(image_urls[1])) + queries_before = _get_mm_cache_stats(runner.llm.get_metrics())[0] # 3 + + renderer = runner.llm.llm_engine.renderer + cache = renderer.mm_processor_cache + assert cache is not None, "Processor cache should be enabled" + + _, eng_prompts = renderer.render_chat( + [_make_messages(image_urls[1])], + ChatParams(), + ) + eng_input = eng_prompts[0] + + # Inject pre-processed mm_kwargs with a NEW hash via public API + new_mm_hash = "deadbeef" * 8 + mm_hashes = {"image": [new_mm_hash]} + mm_kwargs = eng_input["mm_kwargs"] + + runner.llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) + + # Build pre-rendered input (no externally_processed flag needed) + pre_rendered_input = { + "type": "multimodal", + "prompt_token_ids": eng_input["prompt_token_ids"], + "mm_kwargs": mm_kwargs, + "mm_hashes": mm_hashes, + "mm_placeholders": eng_input["mm_placeholders"], + } + + runner.llm.generate( + pre_rendered_input, + sampling_params=SamplingParams(max_tokens=1), + ) + + # Verify cache was queried and injection happened + queries_after = _get_mm_cache_stats(runner.llm.get_metrics())[0] + assert queries_after > queries_before, ( + "Cache should have been queried for the injected item" + ) + mm_rate = _get_mm_cache_log(runner.llm, caplog_vllm) + assert mm_rate >= 0.0, "MM cache hit rate should be reported" @pytest.mark.parametrize("image_urls", [TEST_IMAGE_ASSETS[:1]], indirect=True) def test_inject_into_mm_cache_without_cache( num_gpus_available, image_urls, - multimodal_llm_factory, + vllm_runner, ): """Test that inject_into_mm_cache works gracefully when processor cache is disabled (mm_processor_cache_gb=0). Should not crash. """ - llm = multimodal_llm_factory( - model="llava-hf/llava-1.5-7b-hf", + with vllm_runner( + "llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, enforce_eager=True, disable_log_stats=False, limit_mm_per_prompt={"image": 2}, mm_processor_cache_gb=0, - ) - - # Run a normal chat request first to warm up the model. - llm.chat(_make_messages(image_urls[0])) - - # Use the renderer to get a proper EngineInput with expanded tokens - renderer = llm.llm_engine.renderer - _, eng_prompts = renderer.render_chat( - [_make_messages(image_urls[0])], - ChatParams(), - ) - eng_input = eng_prompts[0] - - mm_hashes = {"image": ["abcd1234" * 8]} - mm_kwargs = eng_input["mm_kwargs"] - - # inject_into_mm_cache should not crash even without cache - llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) - - # Build and generate with pre-rendered input - pre_rendered_input = { - "type": "multimodal", - "prompt_token_ids": eng_input["prompt_token_ids"], - "mm_kwargs": mm_kwargs, - "mm_hashes": mm_hashes, - "mm_placeholders": eng_input["mm_placeholders"], - } - - result = llm.generate( - pre_rendered_input, - sampling_params=SamplingParams(max_tokens=1), - ) - assert len(result) == 1, "Should produce one output" - assert len(result[0].outputs) >= 1, "Should have at least one output sequence" + ) as runner: + # Run a normal chat request first to warm up the model. + runner.llm.chat(_make_messages(image_urls[0])) + + # Use the renderer to get a proper EngineInput with expanded tokens + renderer = runner.llm.llm_engine.renderer + _, eng_prompts = renderer.render_chat( + [_make_messages(image_urls[0])], + ChatParams(), + ) + eng_input = eng_prompts[0] + + mm_hashes = {"image": ["abcd1234" * 8]} + mm_kwargs = eng_input["mm_kwargs"] + + # inject_into_mm_cache should not crash even without cache + runner.llm.llm_engine.input_processor.inject_into_mm_cache(mm_hashes, mm_kwargs) + + # Build and generate with pre-rendered input + pre_rendered_input = { + "type": "multimodal", + "prompt_token_ids": eng_input["prompt_token_ids"], + "mm_kwargs": mm_kwargs, + "mm_hashes": mm_hashes, + "mm_placeholders": eng_input["mm_placeholders"], + } + + result = runner.llm.generate( + pre_rendered_input, + sampling_params=SamplingParams(max_tokens=1), + ) + assert len(result) == 1, "Should produce one output" + assert len(result[0].outputs) >= 1, "Should have at least one output sequence" diff --git a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py index dbea37f64eea..ac86788f69a6 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py +++ b/tests/entrypoints/multimodal/llm/test_mm_cache_stats.py @@ -61,38 +61,37 @@ def test_mm_cache_stats( image_urls, mm_processor_cache_type, caplog_vllm, - multimodal_llm_factory, + vllm_runner, ): - llm = multimodal_llm_factory( - model="llava-hf/llava-1.5-7b-hf", + with vllm_runner( + "llava-hf/llava-1.5-7b-hf", max_model_len=4096, max_num_seqs=5, enforce_eager=True, mm_processor_cache_type=mm_processor_cache_type, disable_log_stats=False, limit_mm_per_prompt={"image": 2}, - ) - - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) - - llm.chat(_make_messages(image_urls[1])) - assert _get_mm_cache_stats(llm.get_metrics()) == (2, 0) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) - - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (3, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(33.3) - - # NOTE: This only resets hit rate stats in CachingMetrics - # The raw queries and hits counts remain unaffected - llm.reset_mm_cache() - - llm.chat(_make_messages(image_urls[0])) - assert _get_mm_cache_stats(llm.get_metrics()) == (4, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) - - llm.chat(_make_messages(image_urls[1])) - assert _get_mm_cache_stats(llm.get_metrics()) == (5, 1) - assert _get_mm_cache_log(llm, caplog_vllm) == pytest.approx(0.0) + ) as runner: + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (1, 0) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) + + runner.llm.chat(_make_messages(image_urls[1])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (2, 0) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) + + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (3, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(33.3) + + # NOTE: This only resets hit rate stats in CachingMetrics + # The raw queries and hits counts remain unaffected + runner.llm.reset_mm_cache() + + runner.llm.chat(_make_messages(image_urls[0])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (4, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) + + runner.llm.chat(_make_messages(image_urls[1])) + assert _get_mm_cache_stats(runner.llm.get_metrics()) == (5, 1) + assert _get_mm_cache_log(runner.llm, caplog_vllm) == pytest.approx(0.0) diff --git a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py index 0ea180712469..f63e9463bc7f 100644 --- a/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py +++ b/tests/entrypoints/multimodal/llm/test_mm_embeds_only.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import weakref + import pytest -from tests.entrypoints.multimodal.conftest import managed_llm from vllm import LLM, SamplingParams from vllm.assets.image import ImageAsset from vllm.exceptions import VLLMValidationError @@ -14,17 +15,19 @@ @pytest.fixture(scope="module") -def llm(): +def llm(vllm_runner): """LLM with enable_mm_embeds=True and all modality limits zeroed out.""" - with managed_llm( - model=MODEL, + with vllm_runner( + MODEL, max_model_len=2048, enforce_eager=True, gpu_memory_utilization=0.8, enable_mm_embeds=True, limit_mm_per_prompt={"image": 0}, - ) as llm: - yield llm + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 1d6919e9c89f..debf547593fa 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -28,9 +28,11 @@ from vllm.entrypoints.openai.chat_completion.serving import ( OpenAIServingChat, _get_mm_token_counts, + _make_completion_tokens_details, _make_prompt_tokens_details, ) from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, ErrorResponse, RequestResponseMetadata, ) @@ -626,6 +628,7 @@ def _build_minimal_metrics_serving_chat( serving.response_role = "assistant" serving.parser_cls = None serving.enable_auto_tools = False + serving._include_reasoning_tokens_details = False serving.enable_prompt_tokens_details = False serving.enable_log_outputs = False serving.enable_log_deltas = False @@ -666,6 +669,13 @@ async def _single_request_output( yield request_output +async def _stream_request_outputs( + *request_outputs: RequestOutput, +) -> AsyncIterator[RequestOutput]: + for request_output in request_outputs: + yield request_output + + async def _collect_metrics_stream_chunks( serving: OpenAIServingChat, request: ChatCompletionRequest, @@ -724,6 +734,7 @@ async def test_chat_per_request_metrics_follow_server_flag(): request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), ) assert disabled_response.metrics is None + assert disabled_response.usage.completion_tokens_details is None enabled_serving = _build_minimal_metrics_serving_chat( enable_per_request_metrics=True @@ -781,6 +792,55 @@ async def test_chat_streaming_metrics_ride_on_usage_chunk(): assert usage_chunks[-1]["metrics"]["time_to_first_token_ms"] == pytest.approx(500.0) +@pytest.mark.asyncio +async def test_streaming_reasoning_usage_counts_across_deltas(): + serving = _build_minimal_metrics_serving_chat(enable_per_request_metrics=False) + serving._include_reasoning_tokens_details = True + serving.model_config = None + + parser = MagicMock() + parser.parse_delta.side_effect = [ + DeltaMessage(reasoning="reasoning"), + DeltaMessage(content="answer"), + ] + parser.count_reasoning_tokens.side_effect = lambda token_ids: sum( + token_id == 20 for token_id in token_ids + ) + serving.parser_cls = MagicMock(return_value=parser) + + first = _make_metrics_request_output(metrics=None, token_ids=(10, 20)) + first.outputs[0].text = "reasoning" + first.outputs[0].finish_reason = None + second = _make_metrics_request_output(metrics=None, token_ids=(11, 30)) + second.outputs[0].text = "answer" + + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test prompt"}], + max_tokens=10, + stream=True, + stream_options={"include_usage": True}, + ) + chunks: list[dict[str, Any]] = [] + async for line in serving.chat_completion_stream_generator( + request, + _stream_request_outputs(first, second), + "chatcmpl-test-id", + "test-model", + conversation=[{"role": "user", "content": "Test"}], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"), + ): + payload = line.removeprefix("data: ").strip() + if payload != "[DONE]": + chunks.append(json.loads(payload)) + + usage_chunks = [chunk for chunk in chunks if chunk.get("usage")] + assert usage_chunks[-1]["usage"]["completion_tokens_details"] == { + "reasoning_tokens": 1 + } + + @dataclass class MockEngine: model_config: MockModelConfig = field(default_factory=MockModelConfig) @@ -810,6 +870,7 @@ async def _async_serving_chat_init(): def test_async_serving_chat_init(): serving_completion = asyncio.run(_async_serving_chat_init()) assert serving_completion.chat_template == CHAT_TEMPLATE + assert serving_completion._include_reasoning_tokens_details is False def test_mm_prompt_tokens_details(): @@ -850,6 +911,10 @@ def test_mm_prompt_tokens_details(): assert details.multimodal_tokens == {"image": 600, "video": 1200} +def test_completion_tokens_details(): + assert _make_completion_tokens_details(7).reasoning_tokens == 7 + + @pytest.mark.asyncio async def test_serving_chat_returns_correct_model_name(): mock_engine = MagicMock(spec=AsyncLLM) diff --git a/tests/entrypoints/openai/completion/test_completion.py b/tests/entrypoints/openai/completion/test_completion.py index 2e7fd090694b..af9f1def3614 100644 --- a/tests/entrypoints/openai/completion/test_completion.py +++ b/tests/entrypoints/openai/completion/test_completion.py @@ -822,3 +822,19 @@ def test_completion_request_bad_words_default_empty(): default_sampling_params={}, ) assert sampling_params.bad_words == [] + + +def test_completion_request_forwards_routed_experts_prompt_start(): + request = CompletionRequest( + model="test-model", + prompt="Hello", + max_tokens=10, + routed_experts_prompt_start=3, + ) + + sampling_params = request.to_sampling_params( + max_tokens=10, + default_sampling_params={}, + ) + + assert sampling_params.routed_experts_prompt_start == 3 diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 48b68e96d8b2..51fea4662363 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -360,6 +360,10 @@ def __init__(self): def get_vocab(self): return self._vocab + def decode(self, token_ids): + id_to_token = {v: k for k, v in self._vocab.items()} + return "".join(id_to_token.get(token_id, "x") for token_id in token_ids) + # Force non-harmony, SimpleContext path monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) diff --git a/tests/entrypoints/openai/test_return_routed_experts.py b/tests/entrypoints/openai/test_return_routed_experts.py index b69c4c49158e..a64e80c39868 100644 --- a/tests/entrypoints/openai/test_return_routed_experts.py +++ b/tests/entrypoints/openai/test_return_routed_experts.py @@ -20,6 +20,18 @@ NUM_HIDDEN_LAYERS = 2 +def assert_valid_routed_experts(encoded: str | None) -> None: + assert encoded is not None + routed_experts = np.load(io.BytesIO(base64.b64decode(encoded))) + assert routed_experts.ndim == 3 + num_tokens, num_layers, topk = routed_experts.shape + assert num_tokens > 0 + assert num_layers == NUM_HIDDEN_LAYERS + assert topk == NUM_EXPERTS_PER_TOK + assert (routed_experts >= 0).all() + assert (routed_experts < NUM_LOCAL_EXPERTS).all() + + @pytest.fixture(scope="module") def server(): args = [ @@ -50,15 +62,5 @@ async def test_routed_experts(server): choice = result.model_dump()["choices"][0] - assert choice["routed_experts"] is not None assert choice["token_ids"] is not None - - # routed_experts is base64-encoded .npy bytes; decode to ndarray. - routed_experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"]))) - assert routed_experts.ndim == 3 - num_tokens, num_layers, topk = routed_experts.shape - assert num_tokens > 0 - assert num_layers == NUM_HIDDEN_LAYERS - assert topk == NUM_EXPERTS_PER_TOK - assert (routed_experts >= 0).all() - assert (routed_experts < NUM_LOCAL_EXPERTS).all() + assert_valid_routed_experts(choice["routed_experts"]) diff --git a/tests/entrypoints/openai/test_stop_token_ids.py b/tests/entrypoints/openai/test_stop_token_ids.py index 74eba026ed99..5e9d79f68649 100644 --- a/tests/entrypoints/openai/test_stop_token_ids.py +++ b/tests/entrypoints/openai/test_stop_token_ids.py @@ -87,6 +87,20 @@ def test_only_client_stop_token_ids(self): assert set(sampling_params.stop_token_ids) == {42, 43} + def test_routed_experts_prompt_start_is_forwarded(self): + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + routed_experts_prompt_start=3, + ) + + sampling_params = request.to_sampling_params( + max_tokens=100, + default_sampling_params={}, + ) + + assert sampling_params.routed_experts_prompt_start == 3 + def test_duplicate_stop_token_ids_deduplicated(self): """Overlapping stop_token_ids between client and server are deduplicated.""" request = ChatCompletionRequest( diff --git a/tests/entrypoints/pooling/basic/test_encode.py b/tests/entrypoints/pooling/basic/test_encode.py index 61222eb78082..b385fc86975c 100644 --- a/tests/entrypoints/pooling/basic/test_encode.py +++ b/tests/entrypoints/pooling/basic/test_encode.py @@ -6,7 +6,6 @@ import pytest from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory from vllm.exceptions import VLLMValidationError MODEL_NAME = "intfloat/multilingual-e5-small" @@ -29,23 +28,20 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/basic/test_tiling_engine.py b/tests/entrypoints/pooling/basic/test_tiling_engine.py index 329f91debb1e..88c9b434b38f 100644 --- a/tests/entrypoints/pooling/basic/test_tiling_engine.py +++ b/tests/entrypoints/pooling/basic/test_tiling_engine.py @@ -6,28 +6,26 @@ import pytest -from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import PoolingParams MODEL_NAME = "intfloat/multilingual-e5-small" @pytest.fixture(scope="module") -def llm(): - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_seqs=2, # small to trigger tiling tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/classify/test_offline.py b/tests/entrypoints/pooling/classify/test_offline.py index 2f6a9f1db4a9..9db7b486426c 100644 --- a/tests/entrypoints/pooling/classify/test_offline.py +++ b/tests/entrypoints/pooling/classify/test_offline.py @@ -7,7 +7,6 @@ from tests.models.utils import softmax from vllm import LLM, ClassificationRequestOutput, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach" @@ -18,23 +17,20 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/embed/test_offline.py b/tests/entrypoints/pooling/embed/test_offline.py index 6dd2955ead2e..641f8b714ef9 100644 --- a/tests/entrypoints/pooling/embed/test_offline.py +++ b/tests/entrypoints/pooling/embed/test_offline.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from vllm import LLM, EmbeddingRequestOutput, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -18,23 +17,21 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - assert embedding_size == llm.model_config.embedding_size - - yield weakref.proxy(llm) - - del llm - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + assert embedding_size == runner.llm.model_config.embedding_size + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index be600cdc8859..10c9bfd3f5fa 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -89,6 +89,17 @@ def hf_model(hf_runner): yield hf_model +@pytest.mark.asyncio +@pytest.mark.parametrize("model_name", [MODEL_NAME]) +async def test_negative_token_ids(client: openai.AsyncOpenAI, model_name: str): + # A negative token id is out of vocabulary just like an over-large one, but + # is not caught by the upper-bound check. Unlike the OpenAI completion + # schema, the pooling schema does not constrain token ids to be + # non-negative, so the request reaches the shared engine-level validation. + with pytest.raises(openai.BadRequestError, match=".*out of vocabulary.*"): + await client.embeddings.create(model=model_name, input=[-1]) + + @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) async def test_basic( diff --git a/tests/entrypoints/pooling/embed/test_protocol.py b/tests/entrypoints/pooling/embed/test_protocol.py index 9d3416b772d1..91c853f91098 100644 --- a/tests/entrypoints/pooling/embed/test_protocol.py +++ b/tests/entrypoints/pooling/embed/test_protocol.py @@ -75,6 +75,41 @@ def test_zero_treated_as_positive(self): # 0.0 >= 0 is True, so bit=1 for all => 127 (signed) assert result.binary[0] == [127] + def test_negative_zero_treated_as_positive(self): + """Packing follows ``value >= 0``, not the IEEE-754 sign bit. + + ``-0.0 >= 0`` is True, so these bits must be set. An implementation + reading the sign bit instead would pack -128 here. + """ + result = build_typed_embeddings([[-0.0] * 8], ["binary"]) + assert result.binary is not None + assert result.binary[0] == [127] + + def test_negative_denormal_treated_as_negative(self): + """Sign is decided at float64 precision. + + The smallest float64 denormal underflows to -0.0 in float32, where + ``>= 0`` is True; narrowing the comparison would wrongly pack 127. + """ + result = build_typed_embeddings([[-5e-324] * 8], ["binary"]) + assert result.binary is not None + assert result.binary[0] == [-128] + + def test_empty_input(self): + result = build_typed_embeddings([], ["binary", "ubinary"]) + assert result.binary == [] + assert result.ubinary == [] + + @pytest.mark.parametrize("embs", [[0.1] * 8, [[[0.1] * 8] * 2]], ids=["1d", "3d"]) + def test_non_2d_input_raises(self, embs): + """A malformed batch must raise rather than pack along the last axis. + + Packing a 1D or 3D input would silently emit wrongly shaped + ``binary``/``ubinary`` values instead of reporting an error. + """ + with pytest.raises(ValueError, match="2D batch"): + build_typed_embeddings(embs, ["binary"]) + def test_non_multiple_of_8_raises(self): embs = [[0.1] * 7] with pytest.raises(ValueError, match="multiple of 8"): diff --git a/tests/entrypoints/pooling/reward/test_token_reward_offline.py b/tests/entrypoints/pooling/reward/test_token_reward_offline.py index 50a4b54682b0..135861e70055 100644 --- a/tests/entrypoints/pooling/reward/test_token_reward_offline.py +++ b/tests/entrypoints/pooling/reward/test_token_reward_offline.py @@ -8,7 +8,6 @@ from tests.models.utils import softmax from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "internlm/internlm2-1_8b-reward" @@ -16,24 +15,21 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, trust_remote_code=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py index e68fb61a99d6..5f0ba9c60e1c 100644 --- a/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_bi_encoder_offline.py @@ -6,8 +6,6 @@ import pytest from tests.entrypoints.pooling.scoring.util import EncoderScoringHfRunner -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "intfloat/multilingual-e5-small" PROMPT = "The chef prepared a delicious meal." @@ -27,23 +25,20 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py index 3ad044987a54..c9bc291e588c 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py @@ -8,7 +8,6 @@ from tests.models.utils import softmax from vllm import LLM, PoolingParams -from vllm.distributed import cleanup_dist_env_and_memory MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls" PROMPT = "The chef prepared a delicious meal." @@ -24,23 +23,20 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py b/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py index 9a4db3d3f84d..b39aa7abb9b3 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_offline.py @@ -5,9 +5,6 @@ import pytest -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory - from .util import ColBERTScoringHfRunner MODEL_NAME = "answerdotai/answerai-colbert-small-v1" @@ -30,23 +27,20 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.fixture(scope="module") diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py b/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py index 4812bd2ab3db..ca7843f8cff8 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_offline_vision.py @@ -5,32 +5,26 @@ import pytest -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory - from .util import make_base64_image, make_image_mm_param MODEL_NAME = "vidore/colpali-v1.3-hf" @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/test_factories.py b/tests/entrypoints/pooling/test_factories.py new file mode 100644 index 000000000000..8972ac49c927 --- /dev/null +++ b/tests/entrypoints/pooling/test_factories.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest + +from vllm.entrypoints.pooling import factories +from vllm.entrypoints.pooling.factories import init_pooling_io_processors +from vllm.entrypoints.pooling.pooling.io_processor import ( + PluginWithIOProcessorPlugins, + UnsupportedCombinedTaskIOProcessor, +) + + +def _bge_m3_config(io_processor_plugin=None): + model_config = MagicMock() + model_config.get_pooling_task.return_value = "embed&token_classify" + model_config.io_processor_plugin = io_processor_plugin + model_config.hf_config.to_dict.return_value = {} + model_config.architecture = "BgeM3EmbeddingModel" + + vllm_config = MagicMock(model_config=model_config) + renderer = MagicMock() + renderer._executor = MagicMock() + chat_template_config = MagicMock( + chat_template=None, + chat_template_content_format="auto", + trust_request_chat_template=False, + ) + return vllm_config, renderer, chat_template_config + + +def test_combined_task_without_plugin_uses_rejection_processor(): + vllm_config, renderer, chat_template_config = _bge_m3_config() + + processors = init_pooling_io_processors( + supported_tasks=("embed", "embed&token_classify"), + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + assert processors.keys() == {"embed&token_classify"} + assert isinstance( + processors["embed&token_classify"], UnsupportedCombinedTaskIOProcessor + ) + + +def test_combined_task_with_plugin_uses_plugin_processor(monkeypatch): + vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin") + monkeypatch.setattr(factories, "has_io_processor", lambda *_: True) + monkeypatch.setattr( + "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor", + lambda *_: MagicMock(), + ) + + processors = init_pooling_io_processors( + supported_tasks=("embed", "embed&token_classify"), + vllm_config=vllm_config, + renderer=renderer, + chat_template_config=chat_template_config, + ) + + assert processors.keys() == {"embed&token_classify", "plugin"} + assert isinstance(processors["plugin"], PluginWithIOProcessorPlugins) + + +def test_combined_task_plain_pooling_request_has_actionable_error(monkeypatch): + from vllm.entrypoints.pooling.pooling.protocol import PoolingCompletionRequest + from vllm.entrypoints.pooling.pooling.serving import ServingPooling + + vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin") + monkeypatch.setattr(factories, "has_io_processor", lambda *_: True) + monkeypatch.setattr( + "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor", + lambda *_: MagicMock(), + ) + + engine_client = MagicMock(renderer=renderer, vllm_config=vllm_config) + models = MagicMock(model_config=vllm_config.model_config) + serving = ServingPooling( + engine_client, + models, + supported_tasks=("embed", "embed&token_classify"), + request_logger=None, + chat_template_config=chat_template_config, + ) + request = PoolingCompletionRequest(model="BAAI/bge-m3", input=["hola"]) + + assert serving.io_processors.keys() == {"embed&token_classify", "plugin"} + io_processor = serving.get_io_processor(request) + with pytest.raises(ValueError, match="plugin request with a 'data' field"): + io_processor.create_pooling_params(request) diff --git a/tests/entrypoints/pooling/token_classify/test_offline.py b/tests/entrypoints/pooling/token_classify/test_offline.py index 3f59b375177f..9715cd63b158 100644 --- a/tests/entrypoints/pooling/token_classify/test_offline.py +++ b/tests/entrypoints/pooling/token_classify/test_offline.py @@ -6,7 +6,6 @@ from vllm import LLM, PoolingRequestOutput from vllm.config import PoolerConfig -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach" @@ -17,24 +16,21 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, pooler_config=PoolerConfig(task="token_classify"), max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - - yield weakref.proxy(llm) - - del llm - - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index 7bc7be486608..e0bda6542284 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -6,7 +6,6 @@ from vllm import LLM, PoolingRequestOutput from vllm.config import PoolerConfig -from vllm.distributed import cleanup_dist_env_and_memory from vllm.tasks import PoolingTask MODEL_NAME = "intfloat/multilingual-e5-small" @@ -17,24 +16,22 @@ @pytest.fixture(scope="module") -def llm(): - # pytest caches the fixture so we use weakref.proxy to - # enable garbage collection - llm = LLM( - model=MODEL_NAME, +def llm(vllm_runner): + with vllm_runner( + MODEL_NAME, + max_model_len=None, pooler_config=PoolerConfig(task="token_embed"), max_num_batched_tokens=32768, tensor_parallel_size=1, gpu_memory_utilization=0.75, enforce_eager=True, seed=0, - ) - assert embedding_size == llm.model_config.embedding_size - - yield weakref.proxy(llm) - - del llm - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + ) as runner: + assert embedding_size == runner.llm.model_config.embedding_size + # pytest caches yielded fixtures until after teardown, so use a proxy to + # avoid retaining the LLM while VllmRunner.__exit__ releases ROCm memory. + yield weakref.proxy(runner.llm) @pytest.mark.skip_global_cleanup diff --git a/tests/entrypoints/serve/exception_handling/__init__.py b/tests/entrypoints/serve/exception_handling/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/utils/test_error_sanitization.py b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py similarity index 56% rename from tests/entrypoints/serve/utils/test_error_sanitization.py rename to tests/entrypoints/serve/exception_handling/test_error_sanitization.py index c871dffb406f..426d1ede745d 100644 --- a/tests/entrypoints/serve/utils/test_error_sanitization.py +++ b/tests/entrypoints/serve/exception_handling/test_error_sanitization.py @@ -10,7 +10,60 @@ import pytest -from vllm.entrypoints.serve.utils.api_utils import sanitize_message +from vllm.entrypoints.serve.exception_handling.utils import sanitize_message + + +def test_sanitize_message(): + assert ( + sanitize_message("<_io.BytesIO object at 0x7a95e299e750>") + == "<_io.BytesIO object>" + ) + + +class TestSanitizeMessageFilePaths: + """sanitize_message should also strip file paths and traceback + frames, not just memory addresses - see #31683.""" + + def test_strips_traceback_style_frame(self): + msg = ( + "1 validation error:\n" + " {'type': 'list_type', 'loc': ('body', 'messages')}\n" + '\n File "/usr/local/lib/python3.12/dist-packages/vllm/' + 'entrypoints/serve/utils/api_utils.py", line 40, ' + "in create_chat_completion\n" + " POST /v1/chat/completions" + ) + result = sanitize_message(msg) + assert "/usr/local/" not in result + assert "api_utils.py" not in result + assert "list_type" in result + + def test_strips_arbitrary_absolute_path(self): + result = sanitize_message("Error in /home/user/project/vllm/server.py") + assert "/home/user" not in result + + def test_strips_single_parent_container_path(self): + """Regression: /app/server.py and /workspace/server.py (common in + container deployments) were missed by the original {2,} quantifier.""" + assert "/app/" not in sanitize_message("Error in /app/server.py") + assert "/workspace/" not in sanitize_message("Error in /workspace/server.py") + + def test_preserves_api_endpoint_paths(self): + msg = "POST /v1/chat/completions failed" + assert "/v1/chat/completions" in sanitize_message(msg) + + def test_preserves_short_field_references(self): + msg = "Invalid value for field 'body.messages'" + assert sanitize_message(msg) == msg + + def test_strips_both_address_and_path(self): + msg = ( + " failed at " + "/usr/local/lib/python3.12/dist-packages/vllm/server.py" + ) + result = sanitize_message(msg) + assert "0x" not in result + assert "/usr/local/" not in result class TestSanitizeMessageCoversLeakPatterns: @@ -63,7 +116,6 @@ class TestAffectedModulesUseSanitize: @pytest.mark.parametrize( "module", [ - "vllm.entrypoints.anthropic.api_router", "vllm.entrypoints.anthropic.serving", "vllm.entrypoints.speech_to_text.realtime.connection", ], diff --git a/tests/entrypoints/serve/instrumentator/test_http_status_metrics.py b/tests/entrypoints/serve/exception_handling/test_http_status_metrics.py similarity index 100% rename from tests/entrypoints/serve/instrumentator/test_http_status_metrics.py rename to tests/entrypoints/serve/exception_handling/test_http_status_metrics.py diff --git a/tests/entrypoints/serve/utils/test_server_utils.py b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py similarity index 98% rename from tests/entrypoints/serve/utils/test_server_utils.py rename to tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py index b135d1c97c70..9d393b414f1a 100644 --- a/tests/entrypoints/serve/utils/test_server_utils.py +++ b/tests/entrypoints/serve/exception_handling/test_validation_exception_handler.py @@ -16,7 +16,7 @@ import pytest from fastapi.exceptions import RequestValidationError -from vllm.entrypoints.serve.utils.server_utils import ( +from vllm.entrypoints.serve.exception_handling.handlers.validation import ( clean_loc_for_param, validation_exception_handler, ) diff --git a/tests/entrypoints/serve/middleware/__init__.py b/tests/entrypoints/serve/middleware/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/middleware/test_authentication_middleware.py b/tests/entrypoints/serve/middleware/test_authentication_middleware.py new file mode 100644 index 000000000000..528261e644db --- /dev/null +++ b/tests/entrypoints/serve/middleware/test_authentication_middleware.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from argparse import Namespace +from typing import get_args + +import pytest +import regex as re +from fastapi import FastAPI +from starlette.responses import JSONResponse +from starlette.routing import Route +from starlette.testclient import TestClient + +from vllm.entrypoints.launchers.api_server.routers import register_api_routers +from vllm.entrypoints.serve.middleware.authenticate import ( + GUARDED_PREFIX, + AuthenticationMiddleware, +) +from vllm.tasks import POOLING_TASKS, SupportedTask + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def get_all_http_routes(app: FastAPI) -> list[tuple[str, list[str]]]: + """Extract all HTTP routes (path, methods) from the FastAPI app.""" + routes = [] + for route in app.routes: + if not isinstance(route, Route): + continue + path = route.path + methods = list(route.methods or {"GET"}) + routes.append((path, methods)) + return routes + + +def generate_test_path(path_template: str) -> str: + """Replace path parameters (e.g. {response_id}) with 'test'.""" + return re.sub(r"\{[^}]+\}", "test", path_template) + + +def _create_app_with_mock_routes(routes: list[tuple[str, list[str]]]) -> FastAPI: + """Create a FastAPI app with AuthenticationMiddleware and mock endpoints.""" + app = FastAPI() + app.add_middleware(AuthenticationMiddleware, tokens=["valid-token"]) + + async def mock_endpoint(): + return JSONResponse({"status": "ok"}) + + for path_template, methods in routes: + allowed_methods = list(set(methods + ["OPTIONS"])) + app.add_api_route( + path_template, + mock_endpoint, + methods=allowed_methods, + include_in_schema=False, + ) + return app + + +class MockModelConfig: + def __init__(self): + self.hf_config = Namespace() + self.hf_config.num_labels = 1 + + def get_pooling_task(self, supported_tasks: tuple["SupportedTask", ...]): + pooling_tasks = [s for s in supported_tasks if s in POOLING_TASKS] + return pooling_tasks[0] if len(pooling_tasks) > 0 else None + + +@pytest.fixture(params=get_args(SupportedTask)) +def task_routes(request, monkeypatch) -> tuple[str, list[tuple[str, list[str]]]]: + """For each supported task, build an app with only that task's routers, + extract all routes, and return the task name and routes.""" + task = request.param + # Enable development mode to register all routes (including dev-only routes). + monkeypatch.setenv("VLLM_SERVER_DEV_MODE", "1") + + app = FastAPI() + args = Namespace() + app.state = Namespace() + app.state.args = args + + # Register routers for this specific task (development mode already enabled). + register_api_routers( + args, app, supported_tasks=(task,), model_config=MockModelConfig() + ) + + routes = get_all_http_routes(app) + return task, routes + + +# --------------------------------------------------------------------------- +# Tests for auto-discovered routes +# --------------------------------------------------------------------------- + + +def test_auto_discovered_protected_routes_require_auth(task_routes): + """For every auto-discovered route that starts with a guarded prefix, + verify that authentication is enforced.""" + task, routes = task_routes + app = _create_app_with_mock_routes(routes) + client = TestClient(app) + + for path_template, methods in routes: + if not path_template.startswith(GUARDED_PREFIX): + continue + + test_path = generate_test_path(path_template) + test_method = methods[0] if methods else "GET" + + resp = client.request(test_method, test_path) + assert resp.status_code == 401, ( + f"[{task}] {test_method} {test_path} should reject missing token" + ) + + resp = client.request( + test_method, test_path, headers={"Authorization": "Bearer wrong"} + ) + assert resp.status_code == 401, ( + f"[{task}] {test_method} {test_path} should reject invalid token" + ) + + resp = client.request( + test_method, test_path, headers={"Authorization": "Bearer valid-token"} + ) + assert resp.status_code == 200, ( + f"[{task}] {test_method} {test_path} should accept valid token" + ) + + +def test_auto_discovered_unprotected_routes_no_auth(task_routes): + """For every auto-discovered route that does NOT start with a guarded + prefix, verify that no authentication is required.""" + task, routes = task_routes + app = _create_app_with_mock_routes(routes) + client = TestClient(app) + + for path_template, methods in routes: + if path_template.startswith(GUARDED_PREFIX): + continue + + test_path = generate_test_path(path_template) + test_method = methods[0] if methods else "GET" + + resp = client.request(test_method, test_path) + assert resp.status_code == 200, ( + f"[{task}] {test_method} {test_path} should be accessible without token" + ) diff --git a/tests/entrypoints/serve/instrumentator/test_optional_middleware.py b/tests/entrypoints/serve/middleware/test_optional_middleware.py similarity index 100% rename from tests/entrypoints/serve/instrumentator/test_optional_middleware.py rename to tests/entrypoints/serve/middleware/test_optional_middleware.py diff --git a/tests/entrypoints/serve/utils/test_api_utils.py b/tests/entrypoints/serve/utils/test_api_utils.py index 2429da27e1e1..c23d0e9e6828 100644 --- a/tests/entrypoints/serve/utils/test_api_utils.py +++ b/tests/entrypoints/serve/utils/test_api_utils.py @@ -6,18 +6,10 @@ from vllm.entrypoints.openai.engine.protocol import StreamOptions from vllm.entrypoints.serve.utils.api_utils import ( get_max_tokens, - sanitize_message, should_include_usage, ) -def test_sanitize_message(): - assert ( - sanitize_message("<_io.BytesIO object at 0x7a95e299e750>") - == "<_io.BytesIO object>" - ) - - @pytest.mark.parametrize( ("stream_options", "expected"), [ @@ -118,49 +110,3 @@ def test_input_length_exceeds_max_model_len(self): input_length=150, default_sampling_params={"max_tokens": 2048}, ) - - -class TestSanitizeMessageFilePaths: - """sanitize_message should also strip file paths and traceback - frames, not just memory addresses - see #31683.""" - - def test_strips_traceback_style_frame(self): - msg = ( - "1 validation error:\n" - " {'type': 'list_type', 'loc': ('body', 'messages')}\n" - '\n File "/usr/local/lib/python3.12/dist-packages/vllm/' - 'entrypoints/serve/utils/api_utils.py", line 40, ' - "in create_chat_completion\n" - " POST /v1/chat/completions" - ) - result = sanitize_message(msg) - assert "/usr/local/" not in result - assert "api_utils.py" not in result - assert "list_type" in result - - def test_strips_arbitrary_absolute_path(self): - result = sanitize_message("Error in /home/user/project/vllm/server.py") - assert "/home/user" not in result - - def test_strips_single_parent_container_path(self): - """Regression: /app/server.py and /workspace/server.py (common in - container deployments) were missed by the original {2,} quantifier.""" - assert "/app/" not in sanitize_message("Error in /app/server.py") - assert "/workspace/" not in sanitize_message("Error in /workspace/server.py") - - def test_preserves_api_endpoint_paths(self): - msg = "POST /v1/chat/completions failed" - assert "/v1/chat/completions" in sanitize_message(msg) - - def test_preserves_short_field_references(self): - msg = "Invalid value for field 'body.messages'" - assert sanitize_message(msg) == msg - - def test_strips_both_address_and_path(self): - msg = ( - " failed at " - "/usr/local/lib/python3.12/dist-packages/vllm/server.py" - ) - result = sanitize_message(msg) - assert "0x" not in result - assert "/usr/local/" not in result diff --git a/tests/entrypoints/serve/utils/test_request_logger.py b/tests/entrypoints/serve/utils/test_request_logger.py index c17f2471e48a..717ff190b3a7 100644 --- a/tests/entrypoints/serve/utils/test_request_logger.py +++ b/tests/entrypoints/serve/utils/test_request_logger.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging from unittest.mock import MagicMock, patch from vllm.entrypoints.serve.utils.request_logger import RequestLogger @@ -27,10 +28,10 @@ def test_request_logger_log_outputs(): mock_logger.info.assert_called_once() call_args = mock_logger.info.call_args.args assert "Generated response %s%s" in call_args[0] + assert "output_token_ids" not in call_args[0] assert call_args[1] == "test-123" assert call_args[3] == "Hello, world!" - assert call_args[4] == [1, 2, 3, 4] - assert call_args[5] == "stop" + assert call_args[4] == "stop" def test_request_logger_log_outputs_streaming_delta(): @@ -56,8 +57,7 @@ def test_request_logger_log_outputs_streaming_delta(): assert call_args[1] == "test-456" assert call_args[2] == " (streaming delta)" assert call_args[3] == "Hello" - assert call_args[4] == [1] - assert call_args[5] is None + assert call_args[4] is None def test_request_logger_log_outputs_streaming_complete(): @@ -83,8 +83,7 @@ def test_request_logger_log_outputs_streaming_complete(): assert call_args[1] == "test-789" assert call_args[2] == " (streaming complete)" assert call_args[3] == "Complete response" - assert call_args[4] == [1, 2, 3] - assert call_args[5] == "length" + assert call_args[4] == "length" def test_request_logger_log_outputs_with_truncation(): @@ -117,11 +116,37 @@ def test_request_logger_log_outputs_with_truncation(): assert len(logged_output) == 10 # Check that token IDs were truncated to first 10 tokens - logged_token_ids = call_args[0][4] + mock_logger.debug.assert_called_once() + logged_token_ids = mock_logger.debug.call_args.args[3] assert logged_token_ids == list(range(10)) assert len(logged_token_ids) == 10 +def test_request_logger_log_output_token_ids_require_debug(): + mock_logger = MagicMock() + mock_logger.isEnabledFor.side_effect = lambda level: level >= logging.INFO + + with patch("vllm.entrypoints.serve.utils.request_logger.logger", mock_logger): + request_logger = RequestLogger(max_log_len=4) + + request_logger.log_outputs( + request_id="test-no-token-ids", + outputs="Test output", + output_token_ids=[1, 2, 3], + finish_reason="stop", + ) + + mock_logger.info.assert_called_once() + assert mock_logger.info.call_args.args == ( + "Generated response %s%s: output: %r, finish_reason: %s", + "test-no-token-ids", + "", + "Test", + "stop", + ) + mock_logger.debug.assert_not_called() + + def test_request_logger_log_outputs_none_values(): """Test log_outputs handles None values correctly.""" mock_logger = MagicMock() @@ -144,8 +169,7 @@ def test_request_logger_log_outputs_none_values(): assert "Generated response %s%s" in call_args[0] assert call_args[1] == "test-none" assert call_args[3] == "Test output" - assert call_args[4] is None - assert call_args[5] == "stop" + assert call_args[4] == "stop" def test_request_logger_log_outputs_empty_output(): @@ -170,8 +194,7 @@ def test_request_logger_log_outputs_empty_output(): assert "Generated response %s%s" in call_args[0] assert call_args[1] == "test-empty" assert call_args[3] == "" - assert call_args[4] == [] - assert call_args[5] == "stop" + assert call_args[4] == "stop" def test_request_logger_log_outputs_integration(): @@ -245,4 +268,4 @@ def test_streaming_complete_logs_full_text_content(): # Verify other parameters assert call_args[1] == "test-streaming-full-text" assert call_args[2] == " (streaming complete)" - assert call_args[5] == "streaming_complete" + assert call_args[4] == "streaming_complete" diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py new file mode 100644 index 000000000000..a40b73e5433c --- /dev/null +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Non-object JSON bodies must fail validation cleanly (4xx), not AttributeError (500). + +mode=before validators that call data.get(...) without an isinstance(data, dict) +guard raise AttributeError for string/list/scalar bodies and surface as HTTP 500. + +This extends the chat completion coverage added in #51654 to the remaining +request models whose before-validators were missing the same guard. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.openai.completion.protocol import CompletionRequest +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.pooling.classify.protocol import ClassificationChatRequest +from vllm.entrypoints.pooling.embed.protocol import EmbeddingChatRequest +from vllm.entrypoints.pooling.pooling.protocol import PoolingChatRequest +from vllm.entrypoints.serve.tokenize.protocol import TokenizeChatRequest +from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest +from vllm.entrypoints.speech_to_text.translation.protocol import TranslationRequest +from vllm.exceptions import VLLMValidationError + +pytestmark = pytest.mark.skip_global_cleanup + +REQUEST_MODELS = [ + CompletionRequest, + ResponsesRequest, + EmbeddingChatRequest, + ClassificationChatRequest, + PoolingChatRequest, + TokenizeChatRequest, + TranscriptionRequest, + TranslationRequest, +] + + +@pytest.mark.parametrize("request_model", REQUEST_MODELS, ids=lambda m: m.__name__) +@pytest.mark.parametrize( + "payload", + [ + "this is not valid json{{{", + ["not", "an", "object"], + 42, + None, + True, + ], +) +def test_request_models_reject_non_object_body(request_model, payload): + with pytest.raises(ValidationError): + request_model.model_validate(payload) + + +def test_completion_request_still_validates_dict_bodies(): + """The guard must not swallow real field-level errors on object bodies.""" + with pytest.raises(VLLMValidationError, match="prompt"): + CompletionRequest.model_validate({"model": "qwen", "prompt": ""}) + + +def test_tokenize_chat_request_still_validates_dict_bodies(): + with pytest.raises(VLLMValidationError, match="add_generation_prompt"): + TokenizeChatRequest.model_validate( + { + "model": "qwen", + "messages": [{"role": "user", "content": "hello"}], + "continue_final_message": True, + "add_generation_prompt": True, + } + ) diff --git a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py index 68119a5683dd..e5c3616d6e65 100644 --- a/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py +++ b/tests/entrypoints/weight_transfer/test_weight_transfer_llm.py @@ -8,13 +8,13 @@ """ import os +import weakref from dataclasses import dataclass from unittest.mock import patch import pytest import torch -from vllm import LLM from vllm.config import WeightTransferConfig from vllm.distributed.weight_transfer.base import ( WeightTransferEngine, @@ -102,25 +102,24 @@ def mock_create_engine(config, vllm_config, device, model): @create_new_process_for_each_test() -def test_get_world_size_tp1(): +def test_get_world_size_tp1(vllm_runner): """Test world_size is correctly configured for TP=1.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") - llm = LLM( - model=MODEL_NAME, + with vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) - - world_size = llm.llm_engine.vllm_config.parallel_config.world_size - assert world_size == 1 + ) as runner: + world_size = runner.llm.llm_engine.vllm_config.parallel_config.world_size + assert world_size == 1 @create_new_process_for_each_test() -def test_init_weight_transfer_engine_calls_engine(): +def test_init_weight_transfer_engine_calls_engine(vllm_runner): """Test that init_weight_transfer_engine calls the engine's init_transfer_engine method.""" if torch.accelerator.device_count() < 1: @@ -131,17 +130,20 @@ def test_init_weight_transfer_engine_calls_engine(): # Enable insecure serialization to allow pickling functions for collective_rpc os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" - with patch( - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", - mock_create_engine, - ): - llm = LLM( - model=MODEL_NAME, + with ( + patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ), + vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) + ) as runner, + ): + llm = weakref.proxy(runner.llm) # Verify engine was created def check_engine_exists(self): @@ -170,7 +172,7 @@ def check_init_called(self): @create_new_process_for_each_test() -def test_update_weights_calls_engine(): +def test_update_weights_calls_engine(vllm_runner): """Test that update_weights calls the engine's receive_weights method.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") @@ -180,17 +182,20 @@ def test_update_weights_calls_engine(): # Enable insecure serialization to allow pickling functions for collective_rpc os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" - with patch( - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", - mock_create_engine, - ): - llm = LLM( - model=MODEL_NAME, + with ( + patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ), + vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) + ) as runner, + ): + llm = weakref.proxy(runner.llm) # First init the weight transfer llm.init_weight_transfer_engine( @@ -233,7 +238,7 @@ def check_update_called(self): @create_new_process_for_each_test() -def test_full_weight_transfer_flow(): +def test_full_weight_transfer_flow(vllm_runner): """Test the complete weight transfer flow: init -> start -> update -> finish.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") @@ -243,17 +248,20 @@ def test_full_weight_transfer_flow(): # Enable insecure serialization to allow pickling functions for collective_rpc os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" - with patch( - "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", - mock_create_engine, - ): - llm = LLM( - model=MODEL_NAME, + with ( + patch( + "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine", + mock_create_engine, + ), + vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) + ) as runner, + ): + llm = weakref.proxy(runner.llm) assert llm.get_weight_version() == "default" @@ -309,20 +317,19 @@ def check_flow(self): @create_new_process_for_each_test() -def test_weight_transfer_config_backend(): +def test_weight_transfer_config_backend(vllm_runner): """Test that WeightTransferConfig backend is properly configured.""" if torch.accelerator.device_count() < 1: pytest.skip("Need at least 1 GPU for this test") # Test with nccl backend - llm = LLM( - model=MODEL_NAME, + with vllm_runner( + MODEL_NAME, enforce_eager=True, load_format="dummy", tensor_parallel_size=1, weight_transfer_config=WeightTransferConfig(backend="nccl"), - ) - - config = llm.llm_engine.vllm_config.weight_transfer_config - assert config is not None - assert config.backend == "nccl" + ) as runner: + config = runner.llm.llm_engine.vllm_config.weight_transfer_config + assert config is not None + assert config.backend == "nccl" diff --git a/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml b/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml index 61fbaf33e29a..84c7fb30196a 100644 --- a/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml +++ b/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml @@ -15,7 +15,7 @@ server_args: >- --block-size 256 --gpu-memory-utilization 0.5 --kv-cache-dtype fp8 - --attention_config.use_fp4_indexer_cache=True + --attention_config.indexer_kv_dtype=mxfp4 --max-num-batched-tokens 16384 --max-num-seqs 128 --speculative-config '{"method":"dspark", diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt new file mode 100644 index 000000000000..63b428ac0e30 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-0.txt @@ -0,0 +1,4 @@ +Qwen3.5-35B-A3B-experts-int8-humming.yaml +gpt-oss-20b-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming.yaml +Qwen3-30B-A3B-FP8-block-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt new file mode 100644 index 000000000000..e0260c504e75 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-1.txt @@ -0,0 +1,5 @@ +Qwen3.6-35B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml +Qwen3-30B-A3B-Fp8-v1-humming.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-0.6B-MXFP8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt b/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt new file mode 100644 index 000000000000..8453e5d73cab --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-a100-shard-2.txt @@ -0,0 +1,5 @@ +Qwen3.5-35B-A3B-FP8-humming.yaml +NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-MXFP4A16-humming.yaml +Qwen3-30B-A3B-NVFP4-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt new file mode 100644 index 000000000000..60179f1159fe --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-h100-shard-0.txt @@ -0,0 +1,5 @@ +Qwen3.5-35B-A3B-experts-int8-humming.yaml +Qwen3-30B-A3B-Instruct-2507-quantized.w8a8-humming.yaml +Qwen3-30B-A3B-Fp8-v1-humming.yaml +Qwen3-30B-A3B-NVFP4-humming.yaml +Qwen2-1.5B-Instruct-FP8W8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt new file mode 100644 index 000000000000..0676dc7623f5 --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-h100-shard-1.txt @@ -0,0 +1,5 @@ +Qwen3.6-35B-A3B-NVFP4-humming.yaml +Qwen3-30B-A3B-AWQ-humming.yaml +Qwen3-30B-A3B-GPTQ-Int4-humming.yaml +Qwen3-30B-A3B-FP8-block-humming.yaml +Qwen3-0.6B-MXFP8-humming.yaml diff --git a/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt b/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt new file mode 100644 index 000000000000..9d8fe128d0ff --- /dev/null +++ b/tests/evals/gsm8k/configs/humming/config-h100-shard-2.txt @@ -0,0 +1,4 @@ +NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4-humming.yaml +Qwen3.5-35B-A3B-FP8-humming.yaml +Qwen3-30B-A3B-MXFP4A16-humming.yaml +gpt-oss-20b-humming.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml b/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml index 742d9e40b8a0..955f9d3e90ec 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml @@ -2,4 +2,4 @@ model_name: "deepseek-ai/DeepSeek-V4-Flash" accuracy_threshold: 0.95 num_questions: 1319 num_fewshot: 5 -server_args: "--trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 2 --attention_config.use_fp4_indexer_cache=True --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --speculative_config.method=mtp --speculative_config.num_speculative_tokens=2" +server_args: "--trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 2 --attention_config.indexer_kv_dtype=mxfp4 --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --speculative_config.method=mtp --speculative_config.num_speculative_tokens=2" diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-0.txt b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-0.txt new file mode 100644 index 000000000000..fd65e8611241 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-0.txt @@ -0,0 +1,3 @@ +DeepSeek-V4-Flash-deepgemm-mega-moe.yaml +Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml +Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-1.txt b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-1.txt new file mode 100644 index 000000000000..7fd830430d45 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-1.txt @@ -0,0 +1,5 @@ +Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml +Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml +Llama-4-Scout-Fp8-CT-vllm-cutlass.yaml +Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml +Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-2.txt b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-2.txt new file mode 100644 index 000000000000..20fda6edc873 --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-2.txt @@ -0,0 +1,5 @@ +Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml +Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml +Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml +Llama-4-Scout-BF16-triton.yaml +Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml diff --git a/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-3.txt b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-3.txt new file mode 100644 index 000000000000..17c52f8f4d0a --- /dev/null +++ b/tests/evals/gsm8k/configs/moe-refactor/config-b200-shard-3.txt @@ -0,0 +1,6 @@ +Llama-4-Scout-BF16-fi-cutlass.yaml +Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml +Nemotron-Nano-30B-NvFp4-ModelOpt-vllm-cutlass.yaml +Mixtral-8x7B-BF16-fi-cutlass.yaml +Qwen3-30B-A3B-NvFp4-CT-marlin.yaml +Mixtral-8x7B-BF16-triton.yaml diff --git a/tests/evals/gsm8k/test_gsm8k_offloading.py b/tests/evals/gsm8k/test_gsm8k_offloading.py index 7ed97f7efcd9..c140a41e2644 100644 --- a/tests/evals/gsm8k/test_gsm8k_offloading.py +++ b/tests/evals/gsm8k/test_gsm8k_offloading.py @@ -133,6 +133,7 @@ class OffloadingModelConfig: connector="OffloadingConnector", # Baseline ~0.49 on 200 questions (measured on GB200). accuracy_threshold=0.39, + cpu_offload_gib=1, ), OffloadingModelConfig( id="offloading-gemma-4-e4b-it", @@ -140,6 +141,7 @@ class OffloadingModelConfig: connector="OffloadingConnector", # Baseline ~0.64 on 200 questions (measured on GB200). accuracy_threshold=0.55, + cpu_offload_gib=1, ), OffloadingModelConfig( id="offloading-qwen3.5-35b", diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py index 1285b4c89e8d..48cd9fd0f69d 100644 --- a/tests/jit_monitor/test_hooks.py +++ b/tests/jit_monitor/test_hooks.py @@ -12,6 +12,7 @@ import pytest +from vllm.platforms import current_platform from vllm.utils import jit_monitor pytestmark = pytest.mark.cpu_test @@ -322,6 +323,10 @@ def __call__(self, *args, **kwargs): # ------------------------------------------------------------------ +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="TileLang JIT monitoring is disabled on ROCm", +) def test_tilelang_jit_kernel_logs_warning(): with _patch_jit_modules(_make_fake_knobs()): from tilelang.jit.kernel import JITKernel @@ -337,6 +342,10 @@ def test_tilelang_jit_kernel_logs_warning(): assert "tl_kernel" in msg +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="TileLang JIT monitoring is disabled on ROCm", +) def test_tilelang_jit_impl_logs_warning(): with _patch_jit_modules(_make_fake_knobs()): from tilelang.jit import JITImpl @@ -386,6 +395,10 @@ def set_mode(self, mode): assert "tilelang_fn" in msg +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="TileLang JIT monitoring is disabled on ROCm", +) def test_tilelang_jit_impl_does_not_log_on_cache_hit(): with _patch_jit_modules(_make_fake_knobs()): from tilelang.jit import JITImpl @@ -413,6 +426,10 @@ def set_mode(self, mode): warning_once.assert_called_once() +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="TileLang JIT monitoring is disabled on ROCm", +) def test_tilelang_from_database_does_not_log(): with _patch_jit_modules(_make_fake_knobs()): from tilelang.jit.kernel import JITKernel @@ -425,6 +442,10 @@ def test_tilelang_from_database_does_not_log(): warning_once.assert_not_called() +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="TileLang JIT monitoring is disabled on ROCm", +) def test_tilelang_error_mode_raises(): with _patch_jit_modules(_make_fake_knobs()): from tilelang.jit.kernel import JITKernel diff --git a/tests/kernels/attention/test_deepgemm_attention.py b/tests/kernels/attention/test_deepgemm_attention.py index 0cea46d6284f..5d11063c7a2b 100644 --- a/tests/kernels/attention/test_deepgemm_attention.py +++ b/tests/kernels/attention/test_deepgemm_attention.py @@ -13,6 +13,7 @@ fp8_fp4_paged_mqa_logits, get_num_sms, get_paged_mqa_logits_metadata, + native_next_n_supported, ) from vllm.utils.import_utils import has_deep_gemm from vllm.utils.math_utils import cdiv @@ -205,7 +206,12 @@ def _ref_fp8_fp4_paged_mqa_logits( @pytest.mark.skipif( not current_platform.has_device_capability(90), reason="SM90 and SM100 only" ) -def test_deepgemm_fp8_fp4_paged_mqa_logits(): +# next_n = 1 + num_speculative_tokens, so next_n=4 is MTP=3 (issue #35878). +@pytest.mark.parametrize("batch_size,next_n", [(4, 1), (2, 2), (2, 4)]) +def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int): + if not native_next_n_supported(next_n): + pytest.skip(f"next_n={next_n} has no native kernel on this architecture") + # NOTE: clean_logits=True is incompatible with the 2D context_lens # required by csrc/apis/attention.hpp; only the False path is exercised. clean_logits = False @@ -213,98 +219,97 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits(): random.seed(0) max_model_len = 4096 - for batch_size, next_n in [(4, 1), (2, 2)]: - for heads, index_dim in [(32, 128)]: - for avg_kv in (2048,): - num_blocks, blocksize = max_model_len * 2, 64 - - q = torch.randn( - (batch_size, next_n, heads, index_dim), - device="cuda", - dtype=torch.bfloat16, - ) - kv_cache = torch.randn( - (num_blocks, blocksize, 1, index_dim), - device="cuda", - dtype=torch.bfloat16, - ) - weights = torch.randn( - (batch_size * next_n, heads), - device="cuda", - dtype=torch.float32, - ) - - context_lens = ( - torch.randint(int(0.8 * avg_kv), int(1.2 * avg_kv), (batch_size,)) - .cuda() - .to(torch.int32) - ) - max_block_len = ( - (context_lens.max().item() + blocksize - 1) // blocksize * blocksize - ) - block_tables = torch.zeros( - (batch_size, max_block_len), - device="cuda", - dtype=torch.int32, - ) + for heads, index_dim in [(32, 128)]: + for avg_kv in (2048,): + num_blocks, blocksize = max_model_len * 2, 64 - counter = 0 - block_idx_pool = list(range(num_blocks)) - random.shuffle(block_idx_pool) - for i in range(batch_size): - ctx_len = int(context_lens[i].item()) - for j in range((ctx_len + blocksize - 1) // blocksize): - block_tables[i][j] = block_idx_pool[counter] - counter += 1 + q = torch.randn( + (batch_size, next_n, heads, index_dim), + device="cuda", + dtype=torch.bfloat16, + ) + kv_cache = torch.randn( + (num_blocks, blocksize, 1, index_dim), + device="cuda", + dtype=torch.bfloat16, + ) + weights = torch.randn( + (batch_size * next_n, heads), + device="cuda", + dtype=torch.float32, + ) - q_fp8 = q.to(torch.float8_e4m3fn) - kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache) - - # deep_gemm paged MQA logits requires 2D context_lens of - # shape (B, next_n) (csrc/apis/attention.hpp:332-335); - # see indexer.py:607-608. For each batch/next_n token, the - # effective context length is context_lens[b] - next_n + j + 1. - next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32) - context_lens_2d = ( - context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange - ).contiguous() - schedule_metadata = get_paged_mqa_logits_metadata( - context_lens_2d, blocksize, get_num_sms() - ) - logits = fp8_fp4_paged_mqa_logits( - (q_fp8, None), - kv_cache_fp8, - weights, - context_lens_2d, - block_tables, - schedule_metadata, - max_model_len, - clean_logits=clean_logits, - ) + context_lens = ( + torch.randint(int(0.8 * avg_kv), int(1.2 * avg_kv), (batch_size,)) + .cuda() + .to(torch.int32) + ) + max_block_len = ( + (context_lens.max().item() + blocksize - 1) // blocksize * blocksize + ) + block_tables = torch.zeros( + (batch_size, max_block_len), + device="cuda", + dtype=torch.int32, + ) - ref_logits = _ref_fp8_fp4_paged_mqa_logits( - q, - kv_cache, - weights, - context_lens, - block_tables, - max_model_len, - ) + counter = 0 + block_idx_pool = list(range(num_blocks)) + random.shuffle(block_idx_pool) + for i in range(batch_size): + ctx_len = int(context_lens[i].item()) + for j in range((ctx_len + blocksize - 1) // blocksize): + block_tables[i][j] = block_idx_pool[counter] + counter += 1 + + q_fp8 = q.to(torch.float8_e4m3fn) + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache) + + # deep_gemm paged MQA logits requires 2D context_lens of + # shape (B, next_n) (csrc/apis/attention.hpp:332-335); + # see indexer.py:607-608. For each batch/next_n token, the + # effective context length is context_lens[b] - next_n + j + 1. + next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32) + context_lens_2d = ( + context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange + ).contiguous() + schedule_metadata = get_paged_mqa_logits_metadata( + context_lens_2d, + blocksize, + get_num_sms(), + ) + logits = fp8_fp4_paged_mqa_logits( + (q_fp8, None), + kv_cache_fp8, + weights, + context_lens_2d, + block_tables, + schedule_metadata, + max_model_len, + clean_logits=clean_logits, + ) - positions = ( - torch.arange(max_model_len, device="cuda") - .unsqueeze(0) - .expand(batch_size * next_n, -1) - ) - row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n - next_n_offset = ( - torch.arange(batch_size * next_n, device="cuda") % next_n - ) - mask = positions <= ( - context_lens[row_indices] - next_n + next_n_offset - ).unsqueeze(1) + ref_logits = _ref_fp8_fp4_paged_mqa_logits( + q, + kv_cache, + weights, + context_lens, + block_tables, + max_model_len, + ) - logits = logits.masked_fill(~mask, 0) - ref_logits = ref_logits.masked_fill(~mask, 0) - diff = calc_diff(logits, ref_logits) - assert diff < 1e-3, f"{diff=}" + positions = ( + torch.arange(max_model_len, device="cuda") + .unsqueeze(0) + .expand(batch_size * next_n, -1) + ) + row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n + next_n_offset = torch.arange(batch_size * next_n, device="cuda") % next_n + mask = positions <= ( + context_lens[row_indices] - next_n + next_n_offset + ).unsqueeze(1) + + logits = logits.masked_fill(~mask, 0) + ref_logits = ref_logits.masked_fill(~mask, 0) + diff = calc_diff(logits, ref_logits) + assert diff < 1e-3, f"{diff=}" diff --git a/tests/kernels/attention/test_flashmla_sparse.py b/tests/kernels/attention/test_flashmla_sparse.py index 040ad50a5bbf..eb953faa4469 100644 --- a/tests/kernels/attention/test_flashmla_sparse.py +++ b/tests/kernels/attention/test_flashmla_sparse.py @@ -4,43 +4,6 @@ import torch -def test_deepseek_v4_c128a_dynamic_topk_packed_buffers(): - from vllm.models.deepseek_v4.sparse_mla import build_c128a_topk_metadata - - device = torch.device("cuda") - capacity_width = 256 - active_width = 128 - global_decode_buffer = torch.empty( - (2, capacity_width), dtype=torch.int32, device=device - ) - decode_lens_buffer = torch.empty(2, dtype=torch.int32, device=device) - prefill_buffer = torch.empty((2, capacity_width), dtype=torch.int32, device=device) - - global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( - positions=torch.tensor([255, 511], dtype=torch.int64, device=device), - compress_ratio=128, - num_decode_tokens=1, - token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32, device=device), - block_table=torch.tensor([[3]], dtype=torch.int32, device=device), - block_size=capacity_width, - slot_mapping=torch.tensor([0, 1], dtype=torch.int64, device=device), - global_decode_buffer=global_decode_buffer, - decode_lens_buffer=decode_lens_buffer, - prefill_buffer=prefill_buffer, - max_compressed_tokens=active_width, - ) - - assert global_decode.shape == (1, active_width) - assert prefill_local.shape == (1, active_width) - assert global_decode.stride() == (active_width, 1) - assert prefill_local.stride() == (active_width, 1) - assert global_decode[0, :2].cpu().tolist() == [768, 769] - assert decode_lens.cpu().tolist() == [2] - assert prefill_local[0, :4].cpu().tolist() == list(range(4)) - assert torch.all(global_decode[0, 2:] == -1) - assert torch.all(prefill_local[0, 4:] == -1) - - def test_sparse_flashmla_metadata_smoke(): import vllm.v1.attention.ops.flashmla as fm @@ -225,6 +188,7 @@ def make_swa_metadata(): token_to_req_indices=torch.tensor([0, 1, 1], dtype=torch.int32), decode_swa_indices=torch.tensor([[5, 6, -1, -1]], dtype=torch.int32), decode_swa_lens=torch.tensor([2], dtype=torch.int32), + decode_swa_width=4, is_valid_token=torch.tensor([True], dtype=torch.bool), num_decodes=1, num_prefills=1, @@ -331,3 +295,120 @@ def make_flashmla_metadata(): assert builder_calls == 4 assert sparse_indices_third is not sparse_indices_fourth assert sparse_lens_third is not sparse_lens_fourth + + +def test_flashinfer_sparse_index_preserves_logical_window(monkeypatch): + from vllm.models.deepseek_v4.nvidia import flashinfer_sparse as flashinfer_mod + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + + captured_shapes_and_windows: list[tuple[int, int]] = [] + + def fake_build(*args, **kwargs): + # window_size is the 12th positional arg of + # build_flashinfer_mixed_sparse_indices. + captured_shapes_and_windows.append((args[0].shape[-1], args[11])) + num_tokens = args[0].shape[0] + args[3].shape[0] + return ( + torch.zeros((num_tokens, 1), dtype=torch.int32), + torch.zeros((num_tokens,), dtype=torch.int32), + ) + + monkeypatch.setattr( + flashinfer_mod, "build_flashinfer_mixed_sparse_indices", fake_build + ) + + attn = object.__new__(flashinfer_mod.DeepseekV4FlashInferMLAAttention) + attn.compress_ratio = 1 + attn.window_size = 4 + attn.topk_indices_buffer = torch.zeros((4, 0), dtype=torch.int32) + + wide_width = 8 + wide_indices = torch.full((1, wide_width), -1, dtype=torch.int32) + wide_indices[0, :2] = torch.tensor([5, 6], dtype=torch.int32) + wide_metadata = DeepseekSparseSWAMetadata( + block_table=torch.tensor([[0, 1]], dtype=torch.int32), + slot_mapping=torch.tensor([0], dtype=torch.int64), + block_size=64, + seq_lens=torch.tensor([8], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32), + token_to_req_indices=torch.tensor([0], dtype=torch.int32), + decode_swa_indices=wide_indices, + decode_swa_lens=torch.tensor([2], dtype=torch.int32), + decode_swa_width=wide_width, + is_valid_token=torch.tensor([True], dtype=torch.bool), + num_decodes=1, + num_prefills=0, + num_decode_tokens=1, + num_prefill_tokens=0, + ) + attn._build_sparse_index_metadata( + kv_cache=None, + swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16), + swa_metadata=wide_metadata, + attn_metadata=None, + swa_only=True, + ) + assert captured_shapes_and_windows == [(wide_width, attn.window_size)] + + empty_width = 8 + empty_metadata = DeepseekSparseSWAMetadata( + block_table=torch.tensor([[0, 1]], dtype=torch.int32), + slot_mapping=torch.tensor([0, 1], dtype=torch.int64), + block_size=64, + seq_lens=torch.tensor([8], dtype=torch.int32), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 2], dtype=torch.int32), + token_to_req_indices=torch.tensor([0, 0], dtype=torch.int32), + decode_swa_indices=torch.empty((0, 1, empty_width), dtype=torch.int32), + decode_swa_lens=torch.empty((0,), dtype=torch.int32), + decode_swa_width=empty_width, + is_valid_token=torch.tensor([True, True], dtype=torch.bool), + num_decodes=0, + num_prefills=1, + num_decode_tokens=0, + num_prefill_tokens=2, + ) + attn._build_sparse_index_metadata( + kv_cache=None, + swa_k_cache=torch.empty((1, 64, 512), dtype=torch.bfloat16), + swa_metadata=empty_metadata, + attn_metadata=None, + swa_only=True, + ) + assert captured_shapes_and_windows == [ + (wide_width, attn.window_size), + (empty_width, attn.window_size), + ] + + +def test_flashinfer_mixed_sparse_indices_separates_window_and_padded_width(): + from vllm.models.deepseek_v4.common.ops.cache_utils import ( + build_flashinfer_mixed_sparse_indices, + ) + + device = torch.device("cuda") + padded_width = 8 + logical_window = 4 + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices( + decode_swa_indices=torch.empty( + (0, padded_width), dtype=torch.int32, device=device + ), + decode_compressed_indices=None, + decode_compressed_topk_lens=None, + prefill_topk_indices=torch.empty((1, 0), dtype=torch.int32, device=device), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32, device=device), + seq_lens=torch.tensor([logical_window], dtype=torch.int32, device=device), + token_to_req_indices=torch.tensor([0], dtype=torch.int32, device=device), + swa_block_table=torch.tensor([[0]], dtype=torch.int32, device=device), + swa_block_size=64, + compressed_block_table=None, + compressed_block_size=64, + window_size=logical_window, + compress_ratio=1, + topk=0, + ) + + assert sparse_indices.shape == (1, padded_width) + assert sparse_indices[0].cpu().tolist() == [0, 1, 2, 3, -1, -1, -1, -1] + assert sparse_lens.cpu().tolist() == [padded_width] diff --git a/tests/kernels/attention/test_mha_attn.py b/tests/kernels/attention/test_mha_attn.py index d73acfc0ee9c..b6578fa2f915 100644 --- a/tests/kernels/attention/test_mha_attn.py +++ b/tests/kernels/attention/test_mha_attn.py @@ -27,9 +27,17 @@ @pytest.fixture(autouse=True) -def clear_cache(): - """Clear lru cache to ensure each test case runs without caching.""" +def reset_test_state(): + """Clear cached selectors and restore process-wide torch defaults.""" + default_device = torch.get_default_device() + default_dtype = torch.get_default_dtype() _cached_get_attn_backend.cache_clear() + try: + yield + finally: + torch.set_default_device(default_device) + torch.set_default_dtype(default_dtype) + _cached_get_attn_backend.cache_clear() devices = ["cpu"] diff --git a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py index a384af6f1317..0c15cb99f1a7 100644 --- a/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py +++ b/tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py @@ -82,6 +82,7 @@ def _build_decode_metadata(): create_vllm_config, ) from vllm.config.vllm import set_current_vllm_config + from vllm.model_executor.layers.attention.mla_attention import get_mla_dims from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.kv_cache_interface import MLAAttentionSpec from vllm.v1.worker.workspace import init_workspace_manager @@ -110,11 +111,20 @@ def _build_decode_metadata(): builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls() - # The builder reads layer.prefill_backend from static_forward_context; a - # stub with the attribute is enough for metadata construction. + # The builder reads prefill_backend and the MLA latent dims off the layer, + # so the stub carries both; the dims come from the model config to stay in + # step with the checkpoint. + mla_dims = get_mla_dims(vllm_config.model_config) layer_name = "placeholder" vllm_config.compilation_config.static_forward_context[layer_name] = ( - types.SimpleNamespace(prefill_backend=torch.empty((1,))) + types.SimpleNamespace( + prefill_backend=torch.empty((1,)), + q_lora_rank=mla_dims.q_lora_rank, + kv_lora_rank=mla_dims.kv_lora_rank, + qk_nope_head_dim=mla_dims.qk_nope_head_dim, + qk_rope_head_dim=mla_dims.qk_rope_head_dim, + v_head_dim=mla_dims.v_head_dim, + ) ) init_workspace_manager(device) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 6fe2a3e77587..2f1e8d279356 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -24,12 +24,27 @@ def _on_split_decode_arch() -> bool: return False +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 + except ImportError: + return False + + # The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other # architectures take the fallback decode kernel, so its tests are skipped there. requires_split_decode_arch = pytest.mark.skipif( not _on_split_decode_arch(), reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", ) +requires_gfx950 = pytest.mark.skipif( + not _on_gfx950(), + reason="optimized sparse decode partial is gfx950-only", +) NOPE_HEAD_DIM = 448 ROPE_HEAD_DIM = 64 @@ -118,28 +133,48 @@ def _pack_fp8_ds_mla_cache( return cache -def _read_fp8_ds_mla_cache( - cache: torch.Tensor, slot: int, block_size: int, use_fnuz: bool -) -> torch.Tensor: - cache_flat = cache.view(torch.uint8).flatten() +def _poison_fp8_ds_mla_cache_row( + cache: torch.Tensor, block_size: int, slot: int = 0 +) -> None: + flat = cache.flatten() block_idx = slot // block_size pos = slot % block_size block_base = block_idx * cache.stride(0) token_base = block_base + pos * 576 scale_base = block_base + block_size * 576 + pos * 8 + flat[token_base] = 0x7F + flat[scale_base : scale_base + 7] = 255 + flat[token_base + NOPE_HEAD_DIM : token_base + 576].view(torch.bfloat16)[0] = float( + "nan" + ) + + +def _read_fp8_ds_mla_cache_rows( + cache: torch.Tensor, + slots: torch.Tensor, + block_size: int, + use_fnuz: bool, +) -> torch.Tensor: + cache_flat = cache.view(torch.uint8).flatten() + block_idx = slots // block_size + pos = slots % block_size + block_base = block_idx * cache.stride(0) + token_base = block_base + pos * 576 + scale_base = block_base + block_size * 576 + pos * 8 fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn - nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM] + nope_offsets = torch.arange(NOPE_HEAD_DIM, device=cache.device) + nope_u8 = cache_flat[token_base[:, None] + nope_offsets] nope = nope_u8.view(fp8_dtype).to(torch.float32) + scale_offsets = torch.arange(7, device=cache.device) scales = torch.exp2( - cache_flat[scale_base : scale_base + 7].to(torch.float32) - 127.0 + cache_flat[scale_base[:, None] + scale_offsets].to(torch.float32) - 127.0 ) - nope = nope * scales.repeat_interleave(64) - rope_u8 = cache_flat[ - token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2 - ] - rope = rope_u8.view(torch.bfloat16).to(torch.float32) - return torch.cat([nope, rope]) + nope = nope * scales.repeat_interleave(64, dim=1) + rope_offsets = torch.arange(ROPE_HEAD_DIM * 2, device=cache.device) + rope_u8 = cache_flat[token_base[:, None] + NOPE_HEAD_DIM + rope_offsets] + rope = rope_u8.contiguous().view(torch.bfloat16).to(torch.float32) + return torch.cat([nope, rope], dim=1) def _ref_sparse_decode_ragged( @@ -158,19 +193,30 @@ def _ref_sparse_decode_ragged( out = torch.empty_like(q_f32) for query_idx in range(q.shape[0]): - row_kv = [ - _read_fp8_ds_mla_cache(main_cache, int(slot), block_size, main_use_fnuz) - for slot in main_rows[query_idx] - ] - if extra_cache is not None and extra_rows is not None: - row_kv.extend( - _read_fp8_ds_mla_cache( - extra_cache, int(slot), block_size, extra_use_fnuz + row_kv = [] + if main_rows[query_idx]: + main_slots = torch.tensor( + main_rows[query_idx], dtype=torch.int64, device=q.device + ) + row_kv.append( + _read_fp8_ds_mla_cache_rows( + main_cache, main_slots, block_size, main_use_fnuz + ) + ) + if extra_cache is not None and extra_rows is not None and extra_rows[query_idx]: + extra_slots = torch.tensor( + extra_rows[query_idx], dtype=torch.int64, device=q.device + ) + row_kv.append( + _read_fp8_ds_mla_cache_rows( + extra_cache, extra_slots, block_size, extra_use_fnuz ) - for slot in extra_rows[query_idx] ) - kv = torch.stack(row_kv).to(q.device) + if not row_kv: + out[query_idx] = 0 + continue + kv = torch.cat(row_kv) for head_idx in range(q.shape[1]): scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale if attn_sink is not None: @@ -198,6 +244,46 @@ def _ragged_from_rows( ) +def _launch_sparse_decode_reduce( + part_m: torch.Tensor, + part_l: torch.Tensor, + part_acc: torch.Tensor, + adaptive_splits: bool, +) -> torch.Tensor: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + num_queries, num_splits, num_heads = part_m.shape + out = torch.empty( + (num_queries, num_heads, HEAD_DIM), + dtype=torch.bfloat16, + device=part_m.device, + ) + attn_sink = torch.empty(1, dtype=torch.float32, device=part_m.device) + mod._sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + num_heads, + HAS_ATTN_SINK=False, + ADAPTIVE_SPLITS=adaptive_splits, + COMB_DIM=HEAD_DIM, + BLOCK_H=1, + NUM_SPLITS=num_splits, + SPLITS_PAD=1 << (num_splits - 1).bit_length(), + num_warps=4, + ) + return out + + @torch.inference_mode() def test_paged_mqa_logits_do_not_contain_nan(monkeypatch) -> None: from vllm._aiter_ops import rocm_aiter_ops @@ -312,6 +398,19 @@ def test_compute_global_topk_ragged_indices_and_indptr() -> None: torch.testing.assert_close(actual_lens, expected_lens) +def test_extra_cache_nan_free_provenance_gate(monkeypatch) -> None: + from vllm.models.deepseek_v4.amd import rocm as mod + + monkeypatch.setattr(mod, "_ON_GFX950", True) + assert mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", True, True) + assert not mod._trust_dsv4_extra_cache_nan_free("bfloat16", False, True) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, False) + + monkeypatch.setattr(mod, "_ON_GFX950", False) + assert not mod._trust_dsv4_extra_cache_nan_free("fp8_ds_mla", False, True) + + @torch.inference_mode() def test_sparse_attn_prefill_ragged_kernel() -> None: from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( @@ -366,6 +465,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device) scale = HEAD_DIM**-0.5 + out = torch.empty_like(q) actual = _rocm_sparse_attn_decode_ragged_triton( q=q, main_cache=main_cache, @@ -378,6 +478,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None: extra_cache=extra_cache, extra_indices=extra_indices, extra_indptr=extra_indptr, + out=out, ) expected = _ref_sparse_decode_ragged( q=q, @@ -391,9 +492,111 @@ def test_sparse_attn_decode_ragged_kernel() -> None: main_use_fnuz=main_use_fnuz, ) + assert actual.data_ptr() == out.data_ptr() torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_decode_ragged_triton, + ) + + device = torch.device("cuda") + block_size = 4 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + extra_cache = torch.zeros_like(main_cache) + _poison_fp8_ds_mla_cache_row(main_cache, block_size) + _poison_fp8_ds_mla_cache_row(extra_cache, block_size) + indices = torch.zeros(1, dtype=torch.int32, device=device) + indptr = torch.tensor([0, 1], dtype=torch.int32, device=device) + + actual = _rocm_sparse_attn_decode_ragged_triton( + q=torch.ones(1, 1, HEAD_DIM, dtype=torch.bfloat16, device=device), + main_cache=main_cache, + main_indices=indices, + main_indptr=indptr, + scale=HEAD_DIM**-0.5, + attn_sink=None, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=indices, + extra_indptr=indptr, + ) + + assert not torch.isnan(actual).any() + assert torch.equal(actual, torch.zeros_like(actual)) + + +@pytest.mark.parametrize("on_gfx950", [False, True]) +@torch.inference_mode() +def test_rocm_ragged_graph_buffer_view_tracks_source_width( + monkeypatch, on_gfx950: bool +) -> None: + from vllm.models.deepseek_v4.amd import rocm as rocm_mod + + monkeypatch.setattr(rocm_mod, "_ON_GFX950", on_gfx950) + + indices_buffer = torch.full((16,), -1, dtype=torch.int32) + indptr_buffer = torch.full((3,), -1, dtype=torch.int32) + first_indices = torch.tensor([3, 5, 7], dtype=torch.int32) + first_indptr = torch.tensor([0, 1, 3], dtype=torch.int32) + first_view, first_indptr_view = rocm_mod._copy_ragged_to_graph_buffers( + first_indices, + first_indptr, + indices_buffer, + indptr_buffer, + num_rows=2, + max_entries_per_row=8, + ) + + second_indices = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32) + second_indptr = torch.tensor([0, 2, 6], dtype=torch.int32) + second_view, second_indptr_view = rocm_mod._copy_ragged_to_graph_buffers( + second_indices, + second_indptr, + indices_buffer, + indptr_buffer, + num_rows=2, + max_entries_per_row=8, + ) + + expected_first_entries = ( + first_indices.numel() if on_gfx950 else indices_buffer.numel() + ) + expected_second_entries = ( + second_indices.numel() if on_gfx950 else indices_buffer.numel() + ) + assert first_view.numel() == expected_first_entries + assert second_view.numel() == expected_second_entries + assert first_view.data_ptr() == second_view.data_ptr() == indices_buffer.data_ptr() + assert first_indptr_view.data_ptr() == second_indptr_view.data_ptr() + assert torch.equal(second_view[: second_indices.numel()], second_indices) + assert torch.equal(second_indptr_view, second_indptr) + + +def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None: + from vllm.models.deepseek_v4.amd import rocm as rocm_mod + from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4SparseMLAMetadataBuilder, + ) + + metadata = SimpleNamespace(for_cudagraph_capture=False) + monkeypatch.setattr( + DeepseekV4SparseMLAMetadataBuilder, + "build_for_cudagraph_capture", + lambda *_: metadata, + ) + builder = object.__new__(rocm_mod.DeepseekV4ROCMAiterMLASparseMetadataBuilder) + + actual = builder.build_for_cudagraph_capture(SimpleNamespace()) + + assert actual is metadata + assert actual.for_cudagraph_capture is _on_gfx950() + + @requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: @@ -408,6 +611,9 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: # A tiny batch on a large device should split to add parallelism. assert mod._decode_num_splits(2, 1, avg_main_len=256.0, avg_extra_len=0.0) > 1 + # The shared gfx942 selector retains its original 16-split ceiling. + assert mod._decode_num_splits(1, 1, 128.0, 8192.0) == 16 + # The chosen count always stays within the searched [1, 16] range, and a # zero-length workload never splits (no work to parallelize). for num_queries in (1, 4, 24, 224, 1024): @@ -418,6 +624,16 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 +@torch.inference_mode() +def test_decode_num_splits_gfx950(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + monkeypatch.setattr(mod, "_decode_cu_count", lambda: 256) + assert mod._decode_gfx950_num_splits(1, 1, 128, 8192) == 32 + assert mod._decode_gfx950_num_splits(17, 1, 128, 32) == 4 + assert mod._decode_gfx950_num_splits(512, 1, 128, 7812) == 1 + + @requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) @pytest.mark.parametrize("with_extra", [True, False]) @@ -473,7 +689,14 @@ def test_sparse_attn_decode_split_k_kernel( scale = HEAD_DIM**-0.5 # Pin the split count so each parametrized value is exercised deterministically. - monkeypatch.setattr(mod, "_decode_num_splits", lambda *args, **kwargs: num_splits) + split_fn = "_decode_gfx950_num_splits" if _on_gfx950() else "_decode_num_splits" + other_split_fn = ( + "_decode_num_splits" if _on_gfx950() else "_decode_gfx950_num_splits" + ) + monkeypatch.setattr(mod, split_fn, lambda *args, **kwargs: num_splits) + monkeypatch.setattr( + mod, other_split_fn, lambda *args, **kwargs: pytest.fail("wrong selector") + ) actual = mod._rocm_sparse_attn_decode_ragged_triton( q=q, @@ -503,6 +726,228 @@ def test_sparse_attn_decode_split_k_kernel( torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_adaptive_reduce_ignores_stale_scratch() -> None: + device = torch.device("cuda") + part_m = torch.full( + (1, 8, 1), + torch.finfo(torch.float32).min, + dtype=torch.float32, + device=device, + ) + part_l = torch.zeros_like(part_m) + part_acc = torch.full( + (1, 8, 1, HEAD_DIM), + float("nan"), + dtype=torch.float32, + device=device, + ) + part_m[:, :2] = 0 + part_l[:, :2] = 1 + part_acc[:, 0] = 1 + part_acc[:, 1] = 3 + + actual = _launch_sparse_decode_reduce(part_m, part_l, part_acc, True) + + assert torch.isfinite(actual).all() + assert torch.equal(actual, torch.full_like(actual, 2)) + + +@requires_gfx950 +@pytest.mark.parametrize("extra_len", [0, 1, 31, 32, 33, 63, 64, 65]) +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_outer64_boundaries( + monkeypatch, extra_len: int +) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(13) + block_size = 4 + num_heads = 16 + num_extra_rows = 80 + q = torch.randn(2, num_heads, HEAD_DIM, dtype=torch.bfloat16, device=device) + q *= 0.125 + main_cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + main_indices = torch.empty(0, dtype=torch.int32, device=device) + main_indptr = torch.zeros(3, dtype=torch.int32, device=device) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_extra_rows, HEAD_DIM, dtype=torch.bfloat16, device=device) + * 0.125, + block_size, + use_fnuz=False, + ) + _poison_fp8_ds_mla_cache_row(extra_cache, block_size) + + raw_row = list(range(1, extra_len + 1)) + if extra_len > 3: + raw_row[3] = -1 + if extra_len > 40: + raw_row[40] = num_extra_rows + if extra_len > 64: + raw_row[64] = num_extra_rows + 1024 + extra_indices, extra_indptr = _ragged_from_rows([raw_row, []], device) + valid_row = [slot for slot in raw_row if 0 <= slot < num_extra_rows] + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1) + actual = mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + ) + expected = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=[[], []], + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=[valid_row, []], + ) + + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + assert torch.equal(actual[1], torch.zeros_like(actual[1])) + + +@requires_gfx950 +@torch.inference_mode() +def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None: + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + device = torch.device("cuda") + torch.manual_seed(17) + block_size = 64 + num_queries = 16 + num_heads = 16 + num_splits = 8 + extra_per_query = 65 * num_splits + max_extra_per_query = 8192 + q = ( + torch.randn( + num_queries, + num_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + main_cache = _pack_fp8_ds_mla_cache( + torch.randn(num_queries, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125, + block_size, + use_fnuz=False, + ) + extra_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_queries * extra_per_query, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + block_size, + use_fnuz=False, + ) + main_rows = [[query_idx] for query_idx in range(num_queries)] + extra_rows = [ + list(range(query_idx * extra_per_query, (query_idx + 1) * extra_per_query)) + for query_idx in range(num_queries) + ] + main_indices, main_indptr = _ragged_from_rows(main_rows, device) + short_extra_rows = [row[:64] for row in extra_rows] + long_indices, long_indptr = _ragged_from_rows(extra_rows, device) + short_indices, short_indptr = _ragged_from_rows(short_extra_rows, device) + extra_indices = torch.full( + (num_queries * max_extra_per_query,), + -1, + dtype=torch.int32, + device=device, + ) + extra_indices[: long_indices.numel()].copy_(long_indices) + extra_indptr = long_indptr.clone() + extra_indices_ptr = extra_indices.data_ptr() + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + out = torch.empty_like(q) + + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: num_splits) + + def run_decode() -> torch.Tensor: + return mod._rocm_sparse_attn_decode_ragged_triton( + q=q, + main_cache=main_cache, + main_indices=main_indices, + main_indptr=main_indptr, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + extra_cache=extra_cache, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + out=out, + extra_cache_nan_free=True, + adaptive_splits=True, + ) + + run_decode() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_out = run_decode() + torch.accelerator.synchronize() + captured_long = out.clone() + expected_long = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=extra_rows, + ) + torch.testing.assert_close(captured_long, expected_long, atol=2e-2, rtol=2e-2) + + extra_indices[: short_indices.numel()].copy_(short_indices) + extra_indptr.copy_(short_indptr) + graph.replay() + torch.accelerator.synchronize() + short_out = out.clone() + expected_short = _ref_sparse_decode_ragged( + q=q, + main_cache=main_cache, + main_rows=main_rows, + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=extra_cache, + extra_rows=short_extra_rows, + ) + + assert captured_out.data_ptr() == out.data_ptr() + assert extra_indices.data_ptr() == extra_indices_ptr + assert extra_indices.numel() == num_queries * max_extra_per_query + assert not torch.equal(short_out, captured_long) + torch.testing.assert_close(short_out, expected_short, atol=2e-2, rtol=2e-2) + + extra_indices[: long_indices.numel()].copy_(long_indices) + extra_indptr.copy_(long_indptr) + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/tests/kernels/core/test_mrope.py b/tests/kernels/core/test_mrope.py index 6f64fabfe9b9..81e0a916a304 100644 --- a/tests/kernels/core/test_mrope.py +++ b/tests/kernels/core/test_mrope.py @@ -6,6 +6,7 @@ import torch from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.mrope import apply_interleaved_rope from vllm.platforms import current_platform from vllm.transformers_utils.config import get_config from vllm.utils.torch_utils import set_random_seed @@ -60,6 +61,50 @@ class MRoPETestInfo(NamedTuple): num_tokens_list = [11, 8192] +def test_apply_interleaved_rope(): + mrope_section = [3, 1, 1] + x = torch.tensor( + [ + [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], + [[10, 11, 12, 13, 14], [15, 16, 17, 18, 19]], + [[20, 21, 22, 23, 24], [25, 26, 27, 28, 29]], + ] + ) + + result = apply_interleaved_rope(x, mrope_section) + + expected = torch.tensor([[0, 11, 22, 3, 4], [5, 16, 27, 8, 9]]) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only test." +) +def test_apply_interleaved_rope_torch_compile(): + mrope_section = [24, 20, 20] + num_tokens = 8192 + rotary_dim = sum(mrope_section) * 2 + cache = torch.randn( + 3, + num_tokens, + rotary_dim, + device=device, + dtype=torch.bfloat16, + ) + x = cache[..., : rotary_dim // 2] + + expected = apply_interleaved_rope(x, mrope_section) + compiled_fn = torch.compile( + apply_interleaved_rope, + backend="inductor", + fullgraph=True, + ) + + result = compiled_fn(x, mrope_section) + + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests." ) diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py new file mode 100644 index 000000000000..97ed4c48e787 --- /dev/null +++ b/tests/kernels/mamba/test_gdn_fused_mtp.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-path tests for fused Qwen3.5 GDN MTP decode.""" + +from __future__ import annotations + +import types +from unittest.mock import patch + +import pytest +import torch + +from vllm.platforms import current_platform + +if not (current_platform.is_cuda() and current_platform.has_device_capability(80)): + pytest.skip( + reason="Fused GDN MTP decode requires CUDA compute capability 8.0+.", + allow_module_level=True, + ) + +from tests.v1.attention.utils import ( # noqa: E402 + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig, set_current_vllm_config # noqa: E402 +from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402 + ChunkGatedDeltaRule, + QwenGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402 + MambaStateShapeCalculator, +) +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( # noqa: E402 + rmsnorm_fn, +) +from vllm.utils.torch_utils import _encode_layer_name # noqa: E402 +from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 + GDNAttentionMetadataBuilder, +) +from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 + +NUM_SPEC = 3 +SPEC_TOKENS = NUM_SPEC + 1 +H = 1 +HV = 8 +K = 128 +V = 128 +CONV_KERNEL = 4 +CONV_DIM = 2 * H * K + HV * V +BLOCK_SIZE = 16 +PREFIX = "model.layers.0.linear_attn" +EPS = 1e-6 + + +class _TestGatedNorm: + def __init__(self, weight: torch.Tensor) -> None: + self.weight = weight + self.bias = None + self.eps = EPS + self.group_size = None + self.norm_before_gate = True + self.activation = "silu" + + def __call__(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return rmsnorm_fn( + x, + self.weight, + None, + z=z, + eps=self.eps, + norm_before_gate=True, + activation=self.activation, + ) + + +def _make_vllm_config(): + config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + hf_config_override={"linear_key_head_dim": K}, + ) + config.additional_config = {"gdn_prefill_backend": "cutedsl"} + config.cache_config.mamba_cache_mode = "none" + config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=NUM_SPEC + ) + return config + + +def _build_layer( + vllm_config, + conv_state: torch.Tensor, + ssm_state: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + conv_weight: torch.Tensor, + norm_weight: torch.Tensor, +): + layer = types.SimpleNamespace( + prefix=PREFIX, + enable_packed_recurrent_decode=False, + disable_tp_for_ba_proj=False, + tp_size=1, + num_k_heads=H, + num_v_heads=HV, + head_k_dim=K, + head_v_dim=V, + key_dim=K, + value_dim=HV * V, + activation="silu", + A_log=a_log, + dt_bias=dt_bias, + conv1d=types.SimpleNamespace(weight=conv_weight, bias=None), + kv_cache=(conv_state, ssm_state), + norm=_TestGatedNorm(norm_weight), + layer_norm_epsilon=EPS, + gdn_decode_kernel="cuda", + ) + with set_current_vllm_config(vllm_config): + layer.chunk_gated_delta_rule = ChunkGatedDeltaRule() + for name in ( + "rearrange_mixed_qkv", + "_forward_core", + "_forward_core_decode_spec_post_conv_fused_norm", + "_forward_core_decode_spec_fused_norm", + "_can_use_fused_gdn_mtp_decode", + "_rms_norm_gated_cuda", + "_forward_core_fused_norm", + "_forward_core_fused_norm_packed", + "split_ba", + ): + setattr( + layer, + name, + types.MethodType(getattr(QwenGatedDeltaNetAttention, name), layer), + ) + return layer + + +@torch.inference_mode() +def test_fused_forward_uses_packed_entrypoint() -> None: + """Fused mode keeps projected QKVZ and BA packed through the model op.""" + device = torch.device("cuda") + num_tokens = 3 + hidden_states = torch.empty(num_tokens, 1, dtype=torch.bfloat16, device=device) + mixed_qkvz = torch.randn( + num_tokens, CONV_DIM + HV * V, dtype=torch.bfloat16, device=device + ) + ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device) + layer = types.SimpleNamespace( + prefix=PREFIX, + enable_fused_gdn_decode=True, + norm=types.SimpleNamespace( + weight=torch.empty(V, dtype=torch.bfloat16, device=device) + ), + num_v_heads=HV, + tp_size=1, + head_v_dim=V, + in_proj_qkvz=lambda _: (mixed_qkvz, None), + in_proj_ba=lambda _: (ba, None), + out_proj=lambda x: (x, None), + ) + layer.forward_cuda = types.MethodType( + QwenGatedDeltaNetAttention.forward_cuda, layer + ) + + def packed_op( + actual_qkvz: torch.Tensor, + actual_ba: torch.Tensor, + output: torch.Tensor, + *, + layer_name: str, + ) -> None: + assert actual_qkvz is mixed_qkvz + assert actual_ba is ba + assert layer_name == _encode_layer_name(PREFIX) + output.fill_(1) + + with ( + patch.object( + torch.ops.vllm, + "qwen_gdn_attention_core_fused_norm_packed", + side_effect=packed_op, + ) as packed_mock, + patch.object(torch.ops.vllm, "qwen_gdn_attention_core") as triton_mock, + ): + output = layer.forward_cuda(hidden_states) + + packed_mock.assert_called_once() + triton_mock.assert_not_called() + torch.testing.assert_close(output, torch.ones_like(output)) + + +@pytest.mark.parametrize( + "seq_lens,query_lens,draft_tokens,expected_fused_calls", + [ + pytest.param([128], [SPEC_TOKENS], [NUM_SPEC], 1, id="pure-mtp"), + pytest.param( + [128, 96], + [SPEC_TOKENS, 64], + [NUM_SPEC, -1], + 0, + id="mixed-mtp-falls-back", + ), + pytest.param([96], [64], [-1], 0, id="pure-prefill"), + pytest.param([128], [1], [-1], 0, id="pure-decode"), + ], +) +@torch.inference_mode() +def test_fused_model_path_matches_reference( + seq_lens: list[int], + query_lens: list[int], + draft_tokens: list[int], + expected_fused_calls: int, +) -> None: + """Fused MTP and its mixed/prefill/decode fallbacks match the reference.""" + torch.manual_seed(1) + device = torch.device("cuda") + vllm_config = _make_vllm_config() + builder = GDNAttentionMetadataBuilder( + kv_cache_spec=MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=NUM_SPEC, + ), + layer_names=[PREFIX], + vllm_config=vllm_config, + device=device, + ) + batch = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, device, arange_block_indices=True + ) + common.block_table_tensor.add_(1) + with set_current_vllm_config(vllm_config): + metadata = builder.build( + common_prefix_len=0, + common_attn_metadata=common, + num_accepted_tokens=torch.ones( + batch.batch_size, dtype=torch.int32, device=device + ), + num_decode_draft_tokens_cpu=torch.tensor(draft_tokens, dtype=torch.int32), + ) + + state_indices = [ + indices + for indices in ( + metadata.spec_state_indices_tensor, + metadata.non_spec_state_indices_tensor, + ) + if indices is not None and indices.numel() > 0 + ] + pool_size = max(int(indices.max().item()) for indices in state_indices) + 1 + conv_state_shape, temporal_state_shape = ( + MambaStateShapeCalculator.gated_delta_net_state_shape( + 1, H, HV, K, V, CONV_KERNEL, NUM_SPEC + ) + ) + conv_state_seed = 0.05 * torch.randn( + pool_size, *conv_state_shape, dtype=torch.bfloat16, device=device + ) + ssm_state_seed = 0.01 * torch.randn( + pool_size, *temporal_state_shape, dtype=torch.float32, device=device + ) + a_log = 0.1 * torch.randn(HV, dtype=torch.float32, device=device) + dt_bias = 0.1 * torch.randn(HV, dtype=torch.float32, device=device) + conv_weight = 0.1 * torch.randn( + CONV_DIM, 1, CONV_KERNEL, dtype=torch.bfloat16, device=device + ) + norm_weight = torch.randn(V, dtype=torch.float32, device=device) + num_tokens = batch.compute_num_tokens() + mixed_qkv = 0.1 * torch.randn( + num_tokens, CONV_DIM, dtype=torch.bfloat16, device=device + ) + b = 0.1 * torch.randn(num_tokens, HV, dtype=torch.bfloat16, device=device) + a = 0.1 * torch.randn_like(b) + output_gate = 0.1 * torch.randn( + num_tokens, HV, V, dtype=torch.bfloat16, device=device + ) + mixed_qkvz = torch.cat((mixed_qkv, output_gate.flatten(1)), dim=-1) + ba = torch.cat((b, a), dim=-1) + context = types.SimpleNamespace(attn_metadata={PREFIX: metadata}) + + reference_layer = _build_layer( + vllm_config, + conv_state_seed.clone(), + ssm_state_seed.clone(), + a_log, + dt_bias, + conv_weight, + norm_weight, + ) + reference_out = torch.zeros_like(output_gate) + with patch.object( + qwen_gdn_linear_attn, "get_forward_context", return_value=context + ): + reference_layer._forward_core( + mixed_qkv=mixed_qkv.clone(), + b=b, + a=a, + core_attn_out=reference_out, + ) + reference_out = reference_layer.norm(reference_out, output_gate) + + fused_layer = _build_layer( + vllm_config, + conv_state_seed.clone(), + ssm_state_seed.clone(), + a_log, + dt_bias, + conv_weight, + norm_weight, + ) + context.no_compile_layers = {PREFIX: fused_layer} + fused_out = torch.zeros_like(output_gate) + fused_op = qwen_gdn_linear_attn.ops.fused_gdn_decode_post_conv_mtp + with ( + patch.object(qwen_gdn_linear_attn, "get_forward_context", return_value=context), + patch.object( + qwen_gdn_linear_attn.ops, + "fused_gdn_decode_post_conv_mtp", + wraps=fused_op, + ) as fused_mock, + ): + torch.ops.vllm.qwen_gdn_attention_core_fused_norm_packed( + mixed_qkvz.clone(), + ba, + fused_out, + layer_name=_encode_layer_name(PREFIX), + ) + + assert fused_mock.call_count == expected_fused_calls + torch.testing.assert_close( + fused_layer.kv_cache[0], reference_layer.kv_cache[0], atol=0, rtol=0 + ) + torch.testing.assert_close(fused_out, reference_out, atol=3e-2, rtol=3e-2) + torch.testing.assert_close( + fused_layer.kv_cache[1], + reference_layer.kv_cache[1], + atol=3e-2, + rtol=3e-2, + ) diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index 61e8aeb7a94e..e039310b93e1 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -529,6 +529,125 @@ def test_mxfp4_cpu_fused_moe_small_expert_blocks(): torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) +# Both E2M1 zero codes: 0b0000 (+0.0) and 0b1000 (-0.0), in both nibbles. +MXFP4_ZERO_BYTES = [0x00, 0x88, 0x08, 0x80] +# 0 and 255 are the ends of the E8M0 range; 127 is the identity scale. +MXFP4_E8M0_VALUES = [0, 1, 127, 200, 254, 255] + + +@pytest.mark.parametrize("zero_byte", MXFP4_ZERO_BYTES) +@pytest.mark.parametrize("e8m0", MXFP4_E8M0_VALUES) +def test_mxfp4_cpu_zero_codes_stay_zero(zero_byte, e8m0): + """A zero E2M1 code stays zero for every E8M0 exponent. + + The unpack applies the block scale as an integer add on the bf16 exponent + field, which is exact for every value in the E2M1 codebook except the two + zeros: 0x0000 and 0x8000 have no exponent to shift, so adding to them + produces a small finite number instead of zero. Both are special-cased, and + this is the invariant that special case exists for. + + Worth pinning separately from ``test_mxfp4_cpu_fused_moe``: there the zero + codes are a small fraction of random weights and a broken special case + would stay inside the 1e-2 tolerance for the low exponents. Here every + weight is a zero, so the output is exactly zero or it is not. + """ + N, K, E, M = 64, 64, 2, 4 + dtype = torch.bfloat16 + set_random_seed(0) + + a = torch.randn(M, K, dtype=dtype) + w1q = torch.full((E, 2 * N, K // 2), zero_byte, dtype=torch.uint8) + w1s = torch.full((E, 2 * N, K // 32), e8m0, dtype=torch.uint8) + # w2 is ordinary: the zeros have to survive the first GEMM and the + # activation, and a nonzero w2 is what would expose it if they did not. + w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10 + w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16) + w2s = w2s.reshape(E, K, N // 32) + + topk_weight = torch.ones((M, 1), dtype=torch.float32) + topk_ids = torch.zeros((M, 1), dtype=torch.int32) + + pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) + pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) + out = ops.fused_experts_cpu( + a, + pw1, + pw2, + topk_weight, + topk_ids, + False, + ops.CPUQuantMethod.MXFP4, + pw1s, + pw2s, + None, + None, + None, + ) + + # silu(0) * 0 = 0, so the whole layer collapses to exactly zero. Not + # assert_close: any nonzero output here is a wrong unpack, not rounding. + assert torch.equal(out, torch.zeros_like(out)), ( + f"zero code 0x{zero_byte:02x} with e8m0={e8m0} produced " + f"max |out| = {out.abs().max().item()}" + ) + + +# Narrower than MXFP4_E8M0_VALUES on purpose: this test keeps nonzero weights, +# and 6.0 * 2**(255-127) is not representable in bf16, so the ends of the E8M0 +# range would compare inf against inf and prove nothing. The all-zero test above +# is the one that can reach them. +MXFP4_E8M0_FINITE = [107, 127, 137] + + +@pytest.mark.parametrize("e8m0", MXFP4_E8M0_FINITE) +def test_mxfp4_cpu_zero_codes_mixed_with_nonzero(e8m0): + """Zeros and nonzeros in the same 32-element scale block. + + The zero check is per lane, not per block: this fails if it is ever + rewritten as a whole-block branch. Uses a single shared exponent so the + reference is a plain power of two. + """ + N, K, E, M = 64, 64, 2, 4 + dtype = torch.bfloat16 + set_random_seed(0) + + a = torch.randn(M, K, dtype=dtype) + # Alternate a zero byte and a nonzero one along K, so every scale block + # holds both kinds. + pattern = torch.tensor([0x88, 0x21], dtype=torch.uint8).repeat(K // 4) + w1q = pattern.view(1, 1, -1).expand(E, 2 * N, K // 2).contiguous() + w1s = torch.full((E, 2 * N, K // 32), e8m0, dtype=torch.uint8) + w1dq = MXFP4QuantizeUtil.dequantize(w1q, dtype, w1s) + + w2_bf16 = torch.randn(E, K, N, dtype=dtype) / 10 + w2q, w2s = MXFP4QuantizeUtil.quantize(w2_bf16) + w2s = w2s.reshape(E, K, N // 32) + w2dq = MXFP4QuantizeUtil.dequantize(w2q, dtype, w2s) + + topk_weight = torch.ones((M, 1), dtype=torch.float32) + topk_ids = torch.zeros((M, 1), dtype=torch.int32) + ref_out = ref_mxfp4_fused_moe(a, w1dq, w2dq, topk_weight, topk_ids, 1) + + pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) + pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) + out = ops.fused_experts_cpu( + a, + pw1, + pw2, + topk_weight, + topk_ids, + False, + ops.CPUQuantMethod.MXFP4, + pw1s, + pw2s, + None, + None, + None, + ) + + torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("M", [1, 32]) @pytest.mark.parametrize("N,K,E,topk", [(128, 128, 4, 2), (64, 64, 4, 2)]) @pytest.mark.parametrize("seed", [0]) diff --git a/tests/kernels/moe/test_deepep_v2_moe.py b/tests/kernels/moe/test_deepep_v2_moe.py index a2ff56be5126..ab609d7eabbb 100644 --- a/tests/kernels/moe/test_deepep_v2_moe.py +++ b/tests/kernels/moe/test_deepep_v2_moe.py @@ -235,6 +235,10 @@ def _deep_ep_v2_moe( w2_scale = w2_scale.to(device=device_idx) pg = torch.distributed.new_group(list(range(pgi.world_size))) + # The caller's set_random_seed() only seeds the parent; spawn() gives each + # worker a fresh unseeded RNG, so seed here too or the inputs below differ + # every run. Offset by rank to keep the ranks' data distinct. + set_random_seed(7 + pgi.rank) test_tensors = TestTensors.make(config) with set_current_vllm_config(VllmConfig()): @@ -373,6 +377,7 @@ def _deep_ep_v2_moe_cudagraph( init_workspace_manager(device) pg = torch.distributed.new_group(list(range(pgi.world_size))) + set_random_seed(7 + pgi.rank) test_tensors = TestTensors.make(config) num_local_experts = config.num_experts // pgi.world_size hidden_size = config.k diff --git a/tests/kernels/moe/test_fused_moe_kernel_gptq_awq.py b/tests/kernels/moe/test_fused_moe_kernel_gptq_awq.py new file mode 100644 index 000000000000..ef5da9ff0909 --- /dev/null +++ b/tests/kernels/moe/test_fused_moe_kernel_gptq_awq.py @@ -0,0 +1,577 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the tensor-descriptor (TD) load path in +fused_moe_kernel_gptq_awq. + +Covers use_int4_w4a16 (packed-nibble unpack via tl.interleave), the only weight +layout with a TD path in this kernel, gated on VLLM_TRITON_USE_TD. Every test +forces TD on and skips where TD cannot run; standalone pointer-path coverage +lives in tests/kernels/moe/test_moe.py::test_fused_moe_wn16, over a wider shape +sweep. + +The _matches_pointer tests compare against the pointer path rather than the fp32 +reference, which is tight enough to catch a subtly wrong nibble interleave. The +K- and N-tail cases target the two places the paths legitimately read different +B values, and are XPU-only because the launcher keeps TD off unaligned K +elsewhere. +""" + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.fused_moe import fused_topk, override_config +from vllm.model_executor.layers.fused_moe.config import int4_w4a16_moe_quant_config +from vllm.model_executor.layers.fused_moe.fused_moe import ( + fused_experts, + should_moe_wna16_use_cuda, +) +from vllm.model_executor.layers.fused_moe.utils import ( + TD_MIN_GATHER_ROWS, + moe_use_td_hw_supported, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types +from vllm.triton_utils import tl + +DEVICE = "xpu" if current_platform.is_xpu() else "cuda" +_HAS_TL_MAKE_DESC = hasattr(tl, "make_tensor_descriptor") +_TD_SKIP_REASON = ( + "TD path needs tl.make_tensor_descriptor (Triton >= 3.6) and hardware " + "whose A-gather compiles: XPU, or NVIDIA Blackwell (sm100+), since " + "tile::gather4 (tcgen05/TMEM) is rejected by ptxas on Hopper and earlier" +) + + +def _td_unsupported() -> bool: + return not _HAS_TL_MAKE_DESC or not moe_use_td_hw_supported() + + +def _skip_if_diverted_to_cuda_kernel( + m: int, e: int, topk: int, group_size: int +) -> None: + """Skip shapes that never reach this Triton kernel on CUDA. + + ``fused_experts_impl`` consults ``should_moe_wna16_use_cuda()`` per GEMM and + hands the launch to ``invoke_fused_moe_wna16_cuda_kernel`` when it holds, so + neither the TD nor the pointer path of ``fused_moe_kernel_gptq_awq`` runs and + the comparison would pass with the whole TD branch deleted. Both GEMMs see the + same ``num_valid_tokens`` (GEMM1: ``m * topk`` with ``top_k=topk``; GEMM2: + ``A.size(0) == m * topk`` with ``top_k=1``), so one check covers both. The + predicate requires ``is_cuda()``, so this never fires on XPU. + """ + if should_moe_wna16_use_cuda( + num_valid_tokens=m * topk, group_size=group_size, num_experts=e, bit=4 + ): + pytest.skip( + "should_moe_wna16_use_cuda() diverts both GEMMs of this shape to the " + "CUDA WNA16 kernel, so fused_moe_kernel_gptq_awq is never launched" + ) + + +# One bf16 ULP at the magnitudes this kernel produces (2^-13 for values around +# 0.03), doubled to leave headroom over a single-element rounding disagreement. +_ONE_ULP_ATOL = 2.5e-4 + + +@pytest.fixture(scope="module") +def vllm_config(): + return VllmConfig() + + +# m=1 is kept on purpose, but it is XPU-only coverage, and the reason differs +# per platform: +# * On XPU the launcher gates TD off for a single-row A (the M == 1 check in +# invoke_fused_moe_wna16_triton_kernel) while the second GEMM still runs with +# M = m * topk and does take TD -- so the case exercises the TD path and +# guards the gate against regressing into wrong output rather than merely +# slower output. +# * On CUDA it exercises neither path of this kernel: with m == 1 the ratio +# m * topk / e is at most 1, so should_moe_wna16_use_cuda() always holds for +# the group sizes GPTQ/AWQ actually ship, and both GEMMs are diverted to +# invoke_fused_moe_wna16_cuda_kernel. _skip_if_diverted_to_cuda_kernel() +# skips it there instead of letting it pass for the wrong reason. +WN16_MNK = [ + (1, 128, 128), + (32, 2048, 128), + (222, 2048, 1024), +] +NUM_EXPERTS = [8] +TOP_KS = [2] +GROUP_SIZES = [128] +HAS_ZP = [True, False] + +# A skip keyed off production dispatch could silently swallow the whole sweep if +# that predicate is ever retuned, so pin the premise: everything above m=1 must +# still reach this kernel. Trivially true off CUDA, load-bearing on CUDA. +assert not any( + should_moe_wna16_use_cuda( + num_valid_tokens=m * TOP_KS[0], + group_size=GROUP_SIZES[0], + num_experts=NUM_EXPERTS[0], + bit=4, + ) + for m, _, _ in WN16_MNK + if m > 1 +), "only the m=1 shape may be diverted to the CUDA WNA16 kernel" + + +def fused_moe( + hidden_states, + w1, + w2, + score, + topk, + renormalize=False, + quant_config=None, + global_num_experts=-1, + expert_map=None, +): + topk_weights, topk_ids, _ = fused_topk( + hidden_states, score.float(), topk, renormalize + ) + return fused_experts( + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=quant_config, + ) + + +def torch_moe(a, w1, w2, score, topk): + """Pure-PyTorch MoE reference for correctness validation. + + Implements fused MoE with SiLU+Mul activation and expert routing. + Used as reference to validate Triton kernel outputs. + """ + score = torch.softmax(score, dim=-1, dtype=torch.float32) + topk_weight, topk_ids = torch.topk(score, topk) + + m, k = a.shape + a_rep = a.view(m, -1, k).repeat(1, topk, 1).reshape(-1, k) + out = torch.zeros(m * topk, w2.shape[1], dtype=a.dtype, device=a.device) + + topk_flat = topk_ids.view(-1) + act = SiluAndMul() + for i in range(w1.shape[0]): + mask = topk_flat == i + if mask.sum(): + tmp = a_rep[mask] @ w1[i].transpose(0, 1) + tmp = act(tmp) + out[mask] = tmp @ w2[i].transpose(0, 1) + + return ( + (out.view(m, -1, w2.shape[1]).to(torch.float32) * topk_weight.view(m, -1, 1)) + .sum(dim=1) + .to(out.dtype) + ) + + +def _prepare_quantized_weights(e, n, k, group_size, has_zp, device, dtype): + """Prepare int4 quantized MoE weights with scales and zero-points. + + int4 only: it is the sole layout of this kernel with a TD path, so there is + nothing for an int8 branch here to compare against. + + Returns: (w1_ref, w2_ref, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp) + """ + w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 10 + w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 10 + + pack_factor = 2 + quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8 + + w1_ref = w1.clone() + w2_ref = w2.clone() + w1_qw = torch.empty((e, 2 * n, k // pack_factor), device=device, dtype=torch.uint8) + w2_qw = torch.empty((e, k, n // pack_factor), device=device, dtype=torch.uint8) + w1_sc = torch.empty((e, 2 * n, k // group_size), device=device, dtype=dtype) + w2_sc = torch.empty((e, k, n // group_size), device=device, dtype=dtype) + + w1_zp = torch.empty( + (e, 2 * n // pack_factor, k // group_size), device=device, dtype=torch.uint8 + ) + w2_zp = torch.empty( + (e, k // pack_factor, n // group_size), device=device, dtype=torch.uint8 + ) + + for i in range(e * 2): + expert_id = i % e + if i // e == 0: + w, w_ref_arr, w_qw_arr, w_sc_arr, w_zp_arr = w1, w1_ref, w1_qw, w1_sc, w1_zp + else: + w, w_ref_arr, w_qw_arr, w_sc_arr, w_zp_arr = w2, w2_ref, w2_qw, w2_sc, w2_zp + + weight, qweight, scales, qzeros = quantize_weights( + w[expert_id].T, quant_type, group_size, has_zp, False + ) + weight = weight.T + qweight = qweight.T.contiguous().to(torch.uint8) + scales = scales.T + + if has_zp: + qzeros = qzeros.T.contiguous().to(torch.uint8) + + qweight = qweight[:, 1::2] * 16 + qweight[:, ::2] + if has_zp: + qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :] + + w_ref_arr[expert_id] = weight + w_qw_arr[expert_id] = qweight + w_sc_arr[expert_id] = scales + if has_zp: + w_zp_arr[expert_id] = qzeros + + return w1_ref, w2_ref, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp + + +def _assert_td_matches_pointer(td_output, pointer_output): + """Compare the TD and pointer paths at one bf16 ULP. + + Both accumulate in fp32 and round to bf16 at slightly different points, so a + single-element 1-ULP disagreement (2^-13 at the magnitudes here) is expected + rather than a defect; observed on Blackwell with TD the value closer to the + fp32 reference. Anything larger is a real divergence -- fault injection at + 4 ULP fails, as does 1% of elements off by 5%. + """ + torch.testing.assert_close(td_output, pointer_output, atol=_ONE_ULP_ATOL, rtol=1e-4) + + +def _build_quant_config(w1_sc, w2_sc, w1_zp, w2_zp, has_zp, group_size): + return int4_w4a16_moe_quant_config( + w1_scale=w1_sc, + w2_scale=w2_sc, + w1_zp=w1_zp if has_zp else None, + w2_zp=w2_zp if has_zp else None, + block_shape=[0, group_size], + ) + + +@pytest.mark.skipif(_td_unsupported(), reason=_TD_SKIP_REASON) +@pytest.mark.parametrize("m,n,k", WN16_MNK) +@pytest.mark.parametrize("e", NUM_EXPERTS) +@pytest.mark.parametrize("topk", TOP_KS) +@pytest.mark.parametrize("group_size", GROUP_SIZES) +@pytest.mark.parametrize("has_zp", HAS_ZP) +def test_fused_moe_wn16_use_td( + m, n, k, e, topk, group_size, has_zp, monkeypatch, vllm_config +): + """TD-path correctness vs the PyTorch reference. + + TD-on only: the TD-off leg would duplicate + tests/kernels/moe/test_moe.py::test_fused_moe_wn16, which already covers + the pointer path against the same reference over a wider shape sweep. + """ + _skip_if_diverted_to_cuda_kernel(m, e, topk, group_size) + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1") + dtype = torch.bfloat16 + torch.manual_seed(7) + + a = torch.randn((m, k), device=DEVICE, dtype=dtype) / 10 + score = torch.randn((m, e), device=DEVICE, dtype=dtype) + + w1_ref, w2_ref, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp = ( + _prepare_quantized_weights(e, n, k, group_size, has_zp, DEVICE, dtype) + ) + quant_config = _build_quant_config(w1_sc, w2_sc, w1_zp, w2_zp, has_zp, group_size) + + with set_current_vllm_config(vllm_config): + triton_output = fused_moe( + a, + w1_qw, + w2_qw, + score, + topk, + renormalize=False, + global_num_experts=e, + quant_config=quant_config, + ) + torch_output = torch_moe(a, w1_ref, w2_ref, score, topk) + + torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0) + + +@pytest.mark.skipif(_td_unsupported(), reason=_TD_SKIP_REASON) +@pytest.mark.parametrize("m,n,k", WN16_MNK) +@pytest.mark.parametrize("e", NUM_EXPERTS) +@pytest.mark.parametrize("topk", TOP_KS) +@pytest.mark.parametrize("group_size", GROUP_SIZES) +@pytest.mark.parametrize("has_zp", HAS_ZP) +def test_fused_moe_wn16_td_matches_pointer( + m, n, k, e, topk, group_size, has_zp, monkeypatch, vllm_config +): + """Direct TD-vs-pointer-path comparison on identical inputs. + + Tighter than the fp32-reference tolerance check above (atol=2e-2), which + is loose enough to miss a subtly wrong nibble interleave -- a swapped + low/high nibble would often still land within that tolerance for random + weights. The two Triton paths should agree much more closely than either + agrees with the fp32 reference. + """ + _skip_if_diverted_to_cuda_kernel(m, e, topk, group_size) + dtype = torch.bfloat16 + torch.manual_seed(7) + + a = torch.randn((m, k), device=DEVICE, dtype=dtype) / 10 + score = torch.randn((m, e), device=DEVICE, dtype=dtype) + + _, _, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp = _prepare_quantized_weights( + e, n, k, group_size, has_zp, DEVICE, dtype + ) + quant_config = _build_quant_config(w1_sc, w2_sc, w1_zp, w2_zp, has_zp, group_size) + + def run(use_td: bool) -> torch.Tensor: + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0") + with set_current_vllm_config(vllm_config): + return fused_moe( + a, + w1_qw, + w2_qw, + score, + topk, + renormalize=False, + global_num_experts=e, + quant_config=quant_config, + ) + + pointer_output = run(use_td=False) + td_output = run(use_td=True) + + _assert_td_matches_pointer(td_output, pointer_output) + + +@pytest.mark.skipif(_td_unsupported(), reason=_TD_SKIP_REASON) +@pytest.mark.parametrize("has_zp", HAS_ZP) +def test_fused_moe_wn16_td_k_tail_matches_pointer(has_zp, monkeypatch, vllm_config): + """K-tail (block_k_diviable=False) case: forces the automatic + tensor-descriptor zero-fill, compared bit-exact against the pointer + path's explicit K-mask. + + Bypasses get_moe_wna16_block_config's auto block-size selection (which + always keeps K block-aligned for the group_size/BLOCK_SIZE_K combinations + it picks) via override_config, forcing a BLOCK_SIZE_K that does not + divide K. + + As with the N-tail test, the kernel's ``K`` is ``A.size(1)``, not this + ``k``: GEMM1 runs with ``K = k`` and GEMM2 with ``K = n`` (its A is the + intermediate activation). Here the tail lands in GEMM1 (96 % 64 == 32), + which is enough to exercise the descriptor's zero-fill; asserted below so + the premise cannot silently decay. + + XPU-only, for the same reason the N-tail test is: off XPU the launcher + disables TD for exactly the unaligned K this test constructs. + """ + m, n, k = 33, 512, 96 + e, topk, group_size = 8, 2, 32 + # k % group_size == 0 (96 % 32 == 0) keeps the scale-tensor shape valid. + forced_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + } + assert k % forced_config["BLOCK_SIZE_K"] != 0, "GEMM1 (K=k) must have a K tail" + # The unaligned K that makes this test meaningful is the same condition the + # launcher bails out on off XPU, which would leave GEMM1 running + # pointer-vs-pointer while only GEMM2 (aligned, so no tail) took TD. + if not current_platform.is_xpu(): + pytest.skip( + "TD is disabled off-XPU for unaligned K, so the K-tail leg would " + "compare the pointer path against itself" + ) + dtype = torch.bfloat16 + torch.manual_seed(7) + + a = torch.randn((m, k), device=DEVICE, dtype=dtype) / 10 + score = torch.randn((m, e), device=DEVICE, dtype=dtype) + + _, _, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp = _prepare_quantized_weights( + e, n, k, group_size, has_zp, DEVICE, dtype + ) + quant_config = _build_quant_config(w1_sc, w2_sc, w1_zp, w2_zp, has_zp, group_size) + + def run(use_td: bool) -> torch.Tensor: + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0") + with set_current_vllm_config(vllm_config), override_config(forced_config): + return fused_moe( + a, + w1_qw, + w2_qw, + score, + topk, + renormalize=False, + global_num_experts=e, + quant_config=quant_config, + ) + + pointer_output = run(use_td=False) + td_output = run(use_td=True) + + _assert_td_matches_pointer(td_output, pointer_output) + + +@pytest.mark.skipif(_td_unsupported(), reason=_TD_SKIP_REASON) +@pytest.mark.parametrize("has_zp", HAS_ZP) +def test_fused_moe_wn16_td_n_tail_matches_pointer(has_zp, monkeypatch, vllm_config): + """N-tail: the pointer path wraps tail lanes with ``% N`` while TD gets + zero-fill, reconciled only by the ``offs_cn < N`` store mask. + + The kernel's N is ``B.size(1)``, so GEMM1 runs at ``N = 2n`` and GEMM2 at + ``N = k``; both are asserted below to keep a tail. + """ + m, n, k = 33, 48, 48 + e, topk, group_size = 8, 2, 16 + forced_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + } + block_n = forced_config["BLOCK_SIZE_N"] + assert (2 * n) % block_n != 0, "GEMM1 (N=2n) must have an N tail" + assert k % block_n != 0, "GEMM2 (N=k) must have an N tail" + # These shapes also leave a K tail, which off XPU trips the K-alignment + # bail-out and disables TD for both GEMMs -- the comparison would then be + # pointer-vs-pointer and pass for the wrong reason. Skip instead. + if not current_platform.is_xpu() and k % forced_config["BLOCK_SIZE_K"] != 0: + pytest.skip( + "TD is disabled off-XPU for unaligned K, so this would compare the " + "pointer path against itself" + ) + dtype = torch.bfloat16 + torch.manual_seed(7) + + a = torch.randn((m, k), device=DEVICE, dtype=dtype) / 10 + score = torch.randn((m, e), device=DEVICE, dtype=dtype) + + _, _, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp = _prepare_quantized_weights( + e, n, k, group_size, has_zp, DEVICE, dtype + ) + quant_config = _build_quant_config(w1_sc, w2_sc, w1_zp, w2_zp, has_zp, group_size) + + def run(use_td: bool) -> torch.Tensor: + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1" if use_td else "0") + with set_current_vllm_config(vllm_config), override_config(forced_config): + return fused_moe( + a, + w1_qw, + w2_qw, + score, + topk, + renormalize=False, + global_num_experts=e, + quant_config=quant_config, + ) + + _assert_td_matches_pointer(run(use_td=True), run(use_td=False)) + + +# Shape shared by the two fallback tests. m=33 with topk=2 over 8 experts keeps +# the launch clear of the should_moe_wna16_use_cuda diversion on CUDA; the +# kernel's K is 128 for GEMM1 and 512 (= n) for GEMM2. +_FALLBACK_MNK = (33, 512, 128) +_FALLBACK_E_TOPK_GROUP = (8, 2, 32) +assert not should_moe_wna16_use_cuda( + num_valid_tokens=_FALLBACK_MNK[0] * _FALLBACK_E_TOPK_GROUP[1], + group_size=_FALLBACK_E_TOPK_GROUP[2], + num_experts=_FALLBACK_E_TOPK_GROUP[0], + bit=4, +), "fallback shapes must reach this kernel, not the CUDA WNA16 one" + + +def _assert_td_falls_back(forced_config, monkeypatch, vllm_config): + """Force TD on with a block config it cannot serve, and require the launch to + survive on the pointer path instead of aborting. + + Reaching the comparison at all is the regression check; the tolerance mirrors + the fp32-reference comparison used elsewhere in this file. + """ + m, n, k = _FALLBACK_MNK + e, topk, group_size = _FALLBACK_E_TOPK_GROUP + dtype = torch.bfloat16 + torch.manual_seed(7) + + a = torch.randn((m, k), device=DEVICE, dtype=dtype) / 10 + score = torch.randn((m, e), device=DEVICE, dtype=dtype) + w1_ref, w2_ref, w1_qw, w2_qw, w1_sc, w2_sc, w1_zp, w2_zp = ( + _prepare_quantized_weights(e, n, k, group_size, False, DEVICE, dtype) + ) + quant_config = _build_quant_config(w1_sc, w2_sc, w1_zp, w2_zp, False, group_size) + + monkeypatch.setenv("VLLM_TRITON_USE_TD", "1") + with set_current_vllm_config(vllm_config), override_config(forced_config): + out = fused_moe( + a, + w1_qw, + w2_qw, + score, + topk, + renormalize=False, + global_num_experts=e, + quant_config=quant_config, + ) + ref = torch_moe(a, w1_ref, w2_ref, score, topk) + + torch.testing.assert_close(out, ref, atol=2e-2, rtol=0) + + +@pytest.mark.skipif(_td_unsupported(), reason=_TD_SKIP_REASON) +@pytest.mark.parametrize("block_size_m", [2, 4]) +def test_td_skipped_below_min_gather_rows(block_size_m, monkeypatch, vllm_config): + """A BLOCK_SIZE_M below the gather minimum must fall back, not abort. + + tensor_descriptor.gather() asserts at least TD_MIN_GATHER_ROWS rows, so a + smaller tile used to kill the launch outright ("descriptor gather must have + at least 8 rows"); reproduced on B200. Two config sources reach it: the + override_config used here (honoured verbatim by try_get_optimal_moe_config, + so this arrives through fused_experts_impl) and get_default_config's + use_moe_wna16_cuda branch, min(16, next_power_of_2(M)), which is the + production path via TritonWNA16Experts.apply. Every other test here forces + a block config >= 8, which is why a green suite still shipped the crash. + """ + assert block_size_m < TD_MIN_GATHER_ROWS, "premise: must be below the minimum" + _assert_td_falls_back( + { + "BLOCK_SIZE_M": block_size_m, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + }, + monkeypatch, + vllm_config, + ) + + +@pytest.mark.skipif(_td_unsupported(), reason=_TD_SKIP_REASON) +def test_td_skipped_below_min_block_k(monkeypatch, vllm_config): + """A BLOCK_SIZE_K below 32 must fall back, not abort. + + The int4 B descriptor is built at byte granularity, so its innermost block + dim is BLOCK_SIZE_K // 2 -- below 16 bytes, which make_tensor_descriptor + rejects. Only override_config can produce such a config + (get_moe_wna16_block_config returns 32 or 64), which is why this needs a + forced config rather than a shape. + """ + forced_config = { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 16, + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + } + # Both GEMMs must stay K-aligned, or the earlier unaligned-K bail-out would + # disable TD off XPU and mask the gate under test. + _, n, k = _FALLBACK_MNK + assert all(K % forced_config["BLOCK_SIZE_K"] == 0 for K in (k, n)), ( + "premise: neither GEMM may hit the unaligned-K bail-out first" + ) + _assert_td_falls_back(forced_config, monkeypatch, vllm_config) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index 9e74ec2fd386..5b8dfab25bc4 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -1095,6 +1095,101 @@ def record_scale_interleave(scale): torch.testing.assert_close(interleaved_scales[1], w2_scale) +def test_gpt_oss_quant_config_supplies_clamped_swiglu_params(): + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4MoEMethod + + fake_method = types.SimpleNamespace( + mxfp4_backend=Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, + w13_precision_config=None, + w2_precision_config=None, + ) + fake_layer = types.SimpleNamespace( + w13_weight_scale=torch.zeros(2, 4, 2, dtype=torch.uint8), + w2_weight_scale=torch.zeros(2, 4, 2, dtype=torch.uint8), + ) + + quant_config = GptOssMxfp4MoEMethod.get_fused_moe_quant_config( + fake_method, fake_layer + ) + + assert quant_config is not None + assert quant_config.gemm1_alpha == 1.702 + assert quant_config.gemm1_beta == 1.0 + assert quant_config.gemm1_clamp_limit == 7.0 + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="CUDA is required for FlashInferExperts parameter storage", +) +@pytest.mark.parametrize("supplied", [False, True]) +def test_flashinfer_experts_swiglu_params_follow_quant_config(supplied: bool): + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + mxfp4_w4a16_moe_quant_config, + ) + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( + FlashInferExperts, + ) + + num_experts, intermediate_size, hidden_size = 2, 128, 256 + w1_scale = torch.zeros( + (num_experts, 2 * intermediate_size, hidden_size // 32), + dtype=torch.uint8, + device="cuda", + ) + w2_scale = torch.zeros( + (num_experts, hidden_size, intermediate_size // 32), + dtype=torch.uint8, + device="cuda", + ) + extra = ( + {"gemm1_alpha": 0.5, "gemm1_beta": 0.25, "gemm1_clamp_limit": 3.0} + if supplied + else {} + ) + quant_config = mxfp4_w4a16_moe_quant_config( + w1_scale=w1_scale, w2_scale=w2_scale, **extra + ) + moe_config = FusedMoEConfig( + num_experts=num_experts, + experts_per_token=1, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=num_experts, + num_logical_experts=num_experts, + activation=MoEActivation.SILU, + device="cuda", + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + ) + experts = FlashInferExperts(moe_config=moe_config, quant_config=quant_config) + + if supplied: + for name, value in ( + ("gemm1_alpha", 0.5), + ("gemm1_beta", 0.25), + ("gemm1_clamp_limit", 3.0), + ): + parameter = getattr(experts, name) + assert isinstance(parameter, torch.Tensor) + assert parameter.dtype == torch.float32 + assert parameter.shape == (num_experts,) + torch.testing.assert_close( + parameter.cpu(), + torch.full((num_experts,), value, dtype=torch.float32), + ) + else: + assert experts.gemm1_alpha is None + assert experts.gemm1_beta is None + assert experts.gemm1_clamp_limit is None + + @pytest.mark.parametrize("topk", [1, 4]) @pytest.mark.parametrize("num_experts", [32]) @pytest.mark.parametrize("num_tokens", [1, 128]) diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 9b41765de6bf..52f87c403e75 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -14,6 +14,10 @@ ) from tests.kernels.utils import fp8_ulp_distance from vllm.config import VllmConfig +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( + B12xFp8BlockScaledMMKernel, + _run_b12x_fp8_block_scaled_mm, +) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, @@ -353,3 +357,54 @@ def test_w8a8_block_fp8_flashinfer_matmul(M, N, K, block_size, out_dtype, seed): torch.abs(out.to(torch.bfloat16) - ref_out.to(torch.bfloat16)) ) / torch.mean(torch.abs(ref_out.to(torch.bfloat16))) assert rel_diff < 0.001 + + +@pytest.mark.parametrize( + "M,N,K", + [(1, 128, 256), (8, 256, 512), (129, 256, 256), (2, 4096, 4096)], +) +@torch.inference_mode() +def test_w8a8_block_fp8_b12x_matmul(M, N, K): + supported, reason = B12xFp8BlockScaledMMKernel.is_supported() + if not supported: + pytest.skip(reason) + + torch.manual_seed(M) + fp8_max = torch.finfo(torch.float8_e4m3fn).max + A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max + B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max + A_fp8, As = per_token_group_quant_fp8(A_bf16, 128, use_ue8m0=False) + B_fp8, Bs = per_block_cast_to_fp8( + B_bf16, + block_size=[128, 128], + use_ue8m0=False, + ) + As = As.float() + Bs = Bs.float() + + ref_out = native_w8a8_block_matmul( + A_fp8, + B_fp8, + As, + Bs, + [128, 128], + torch.bfloat16, + ) + out = _run_b12x_fp8_block_scaled_mm( + A_fp8, + B_fp8, + As, + Bs, + torch.bfloat16, + ) + + rel_diff = torch.mean(torch.abs(out.float() - ref_out.float())) / torch.mean( + torch.abs(ref_out.float()) + ) + cosine = torch.nn.functional.cosine_similarity( + out.float().flatten(), + ref_out.float().flatten(), + dim=0, + ) + assert rel_diff < 0.002 + assert cosine >= 0.9999 diff --git a/tests/kernels/quantization/test_block_int8.py b/tests/kernels/quantization/test_block_int8.py deleted file mode 100644 index 4a12ecd50dad..000000000000 --- a/tests/kernels/quantization/test_block_int8.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_block_int8.py -import itertools - -import pytest -import torch - -from tests.kernels.quant_utils import native_w8a8_block_matmul -from vllm.config import VllmConfig -from vllm.model_executor.layers.quantization.utils.int8_utils import ( - w8a8_block_int8_matmul, -) -from vllm.platforms import current_platform - -if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): - pytest.skip( - "INT8 Triton kernels require a CUDA-alike or XPU device", - allow_module_level=True, - ) - -if current_platform.is_cuda_alike() and not current_platform.has_device_capability( - (7, 0) -): - pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True) - -vllm_config = VllmConfig() - -DTYPES = [torch.half, torch.bfloat16] -M = [1, 33, 64, 222] -N = [128, 1024] -K = [256, 4096] -# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]] -BLOCK_SIZE = [[128, 128]] -SEEDS = [0] - - -@pytest.mark.parametrize( - "M,N,K,block_size,out_dtype,seed", - itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS), -) -@torch.inference_mode() -def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed): - torch.manual_seed(seed) - device = current_platform.device_type - factor_for_scale = 1e-2 - int8_info = torch.iinfo(torch.int8) - int8_max, int8_min = int8_info.max, int8_info.min - - A_fp32 = torch.rand(M, K, dtype=torch.float32, device=device) - A_fp32 = (A_fp32 - 0.5) * 2 * int8_max - A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) - - B_fp32 = torch.rand(N, K, dtype=torch.float32, device=device) - B_fp32 = (B_fp32 - 0.5) * 2 * int8_max - B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn) - - block_n, block_k = block_size[0], block_size[1] - n_tiles = (N + block_n - 1) // block_n - k_tiles = (K + block_k - 1) // block_k - - As = torch.rand(M, k_tiles, dtype=torch.float32, device=device) * factor_for_scale - Bs = ( - torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device) - * factor_for_scale - ) - - ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) - out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) - - rel_diff = torch.mean( - torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) - ) / torch.mean(torch.abs(ref_out.to(torch.float32))) - assert rel_diff < 0.001 diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py index 42f163dea44a..3c4682911656 100644 --- a/tests/kernels/quantization/test_triton_w4a16.py +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -306,6 +306,7 @@ class DummyLayer(torch.nn.Module): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_triton_w4a16_process_weights_after_loading_keeps_gptq_qzeros_layout(): + """AutoGPTQ qzeros are already [K//G, N//8] (output_dim=1): no transpose.""" if not torch.cuda.is_available(): pytest.skip("CUDA/HIP device not available") @@ -413,6 +414,7 @@ class DummyLayer(torch.nn.Module): @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") def test_triton_w4a16_symmetric_apply_ignores_qzeros(monkeypatch): + """For symmetric (uint4b8) layers, apply_weights must pass qzeros=None.""" from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( MPLinearLayerConfig, ) diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index d8878a0aea90..cfdec96e61d0 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -28,13 +28,91 @@ _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, _launch_two_stage_sparse_attn_compressor, + compress_norm_rope_store_triton, ) from vllm.models.deepseek_v4.compressor import _get_c128_boundary from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, +) +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + cp_gather_indexer_k_quant_cache_triton, + indexer_k_quant_and_cache_triton, +) from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 +def _on_gfx950() -> bool: + if not current_platform.is_rocm(): + return False + try: + from vllm.platforms.rocm import _ON_GFX950 + + return _ON_GFX950 + except Exception: + return False + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +def test_cp_gather_despecialized_kernel_is_gfx950_only(monkeypatch): + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + + class FakeKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(*args): + self.calls.append((grid, args)) + + return launch + + legacy_kernel = FakeKernel() + gfx950_kernel = FakeKernel() + monkeypatch.setattr(mod, "_cp_gather_indexer_quant_cache_kernel", legacy_kernel) + monkeypatch.setattr( + mod, + "_cp_gather_indexer_quant_cache_gfx950_kernel", + gfx950_kernel, + ) + + k_cache = torch.zeros((4, 1, 132), dtype=torch.uint8) + k_fp8 = torch.empty((5, 128), dtype=current_platform.fp8_dtype()) + k_scale = torch.empty((5, 4), dtype=torch.uint8) + block_table = torch.zeros((2, 7), dtype=torch.int32) + cu_seqlen = torch.tensor([0, 2, 5], dtype=torch.int32) + token_to_seq = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32) + args = (k_cache, k_fp8, k_scale, block_table, cu_seqlen, token_to_seq) + + monkeypatch.setattr(mod, "_ON_GFX950", True) + mod.cp_gather_indexer_k_quant_cache_triton(*args) + assert len(gfx950_kernel.calls) == 1 + assert not legacy_kernel.calls + gfx950_grid, gfx950_args = gfx950_kernel.calls[0] + assert gfx950_grid == (5,) + assert len(gfx950_args) == 18 + assert gfx950_args[-3:] == (2, 7, 4) + + monkeypatch.setattr(mod, "_ON_GFX950", False) + mod.cp_gather_indexer_k_quant_cache_triton(*args) + assert len(legacy_kernel.calls) == 1 + legacy_grid, legacy_args = legacy_kernel.calls[0] + assert legacy_grid == (5,) + assert len(legacy_args) == 19 + assert legacy_args[-4:] == (5, 2, 7, 4) + + +@pytest.mark.parametrize( + ("window_size", "num_speculative_tokens", "expected"), + [(128, 5, 192), (512, 5, 576), (1024, 0, 1024)], +) +def test_get_dspark_swa_index_width( + window_size: int, num_speculative_tokens: int, expected: int +): + assert get_dspark_swa_index_width(window_size, num_speculative_tokens) == expected + + def test_compute_global_topk_reuses_output_buffers(): device = "cuda" topk_indices = torch.tensor( @@ -78,6 +156,129 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float): return x_fp8, scales +def _decode_dsv4_cache_row( + cache: torch.Tensor, block_size: int, scrub_nan: bool +) -> torch.Tensor: + flat = cache.flatten() + nope = flat[:448].view(torch.float8_e4m3fn).to(torch.bfloat16) + encoded = flat[block_size * 576 : block_size * 576 + 7] + scales = torch.exp2(encoded.to(torch.float32) - 127.0).to(torch.bfloat16) + nope = nope * scales.repeat_interleave(64) + rope = flat[448:576].view(torch.bfloat16) + decoded = torch.cat((nope, rope)) + if scrub_nan: + decoded = torch.where(decoded == decoded, decoded, 0.0) + return decoded + + +def _assert_nan_free_cache_matches_legacy_scrub( + cache: torch.Tensor, block_size: int +) -> None: + flat = cache.flatten() + scale_base = block_size * 576 + scale_codes = flat[scale_base : scale_base + 8] + assert scale_codes[0].item() == 254 + assert scale_codes[1].item() == 247 + assert scale_codes[:7].max().item() <= 254 + nope_bytes = flat[:448] + assert not ((nope_bytes == 0x7F) | (nope_bytes == 0xFF)).any() + + rope = flat[448:576].view(torch.bfloat16) + assert not torch.isnan(rope).any() + assert torch.isposinf(rope[0]) + assert torch.equal(rope[1:4], torch.zeros_like(rope[1:4])) + + legacy_cache = cache.clone() + legacy_flat = legacy_cache.flatten() + legacy_flat[scale_base] = 255 + legacy_rope = legacy_flat[448:576].view(torch.bfloat16) + legacy_rope[1:4] = float("nan") + canonical = _decode_dsv4_cache_row(cache, block_size, scrub_nan=False) + legacy = _decode_dsv4_cache_row(legacy_cache, block_size, scrub_nan=True) + torch.testing.assert_close(canonical, legacy, rtol=0, atol=0) + assert torch.isinf(canonical[0]) + assert torch.isposinf(canonical[64]) + + +@pytest.mark.skipif( + not _on_gfx950(), + reason="NaN-free fp8_ds_mla compressed-cache contract is gfx950-only", +) +@pytest.mark.parametrize("writer", ["single_pass", "two_stage_finalizer"]) +def test_gfx950_compressed_cache_canonicalizes_nonfinite(writer: str) -> None: + head_dim = 512 + rope_dim = 64 + block_size = 4 + device = "cuda" + + positions = torch.zeros(1, dtype=torch.int64, device=device) + slot_mapping = torch.zeros(1, dtype=torch.int64, device=device) + rms_weight = torch.ones(head_dim, dtype=torch.bfloat16, device=device) + rms_weight[0] = float("inf") + rms_weight[64] = torch.finfo(torch.bfloat16).max + rms_weight[448] = float("inf") + rms_weight[450] = float("nan") + cos_sin_cache = torch.zeros(1, rope_dim, dtype=torch.float32, device=device) + cos_sin_cache[:, : rope_dim // 2] = 1.0 + cache = torch.zeros(1, block_size, 584, dtype=torch.uint8, device=device) + + state_cache = torch.zeros(1, 1, 2 * head_dim, dtype=torch.float32, device=device) + state_cache[..., :head_dim] = 1.0 + token_to_req = torch.zeros(1, dtype=torch.int32, device=device) + block_table = torch.zeros(1, 1, dtype=torch.int32, device=device) + + if writer == "single_pass": + compress_norm_rope_store_triton( + state_cache=state_cache, + num_actual=1, + token_to_req_indices=token_to_req, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=1, + state_width=head_dim, + cos_sin_cache=cos_sin_cache, + kv_cache=cache, + k_cache_metadata=SimpleNamespace(slot_mapping=slot_mapping), + pdl_kwargs={}, + head_dim=head_dim, + rope_head_dim=rope_dim, + compress_ratio=1, + overlap=False, + use_fp4_cache=False, + rms_norm_weight=rms_weight, + rms_norm_eps=1e-6, + quant_block=64, + token_stride=576, + scale_dim=8, + ) + else: + _launch_two_stage_sparse_attn_compressor( + state_cache, + token_to_req, + positions, + slot_mapping, + block_table, + 1, + head_dim, + 1, + cos_sin_cache, + cache, + slot_mapping, + rms_weight, + 1e-6, + 64, + 576, + 8, + head_dim, + rope_dim, + 1, + torch.empty(1, head_dim, dtype=torch.float32, device=device), + ) + + _assert_nan_free_cache_matches_legacy_scrub(cache, block_size) + + @pytest.mark.parametrize( ("starts", "query_start_loc", "expected"), [ @@ -389,7 +590,8 @@ def test_indexer_gather_accepts_upper_bound_output(): valid_tokens = 9 upper_bound_tokens = 13 block_size = 16 - num_blocks = 2 + num_seqs = 3 + num_blocks = num_seqs sentinel = 123 device = "cuda" @@ -397,13 +599,15 @@ def test_indexer_gather_accepts_upper_bound_output(): kv_cache = torch.zeros( num_blocks, block_size, cache_stride, dtype=torch.uint8, device=device ) - slot_mapping = torch.arange(valid_tokens, dtype=torch.int64, device=device) + slot_mapping = torch.tensor( + [0, 1, 2, 16, 17, 18, 32, 33, 34], dtype=torch.int64, device=device + ) ops.indexer_k_quant_and_cache(k, kv_cache, slot_mapping, quant_block_size, "ue8m0") block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze( - 0 + 1 ) - cu_seq_lens = torch.tensor([0, valid_tokens], dtype=torch.int32, device=device) + cu_seq_lens = torch.tensor([0, 3, 6, 9], dtype=torch.int32, device=device) dst_k = torch.full( (upper_bound_tokens, head_dim), sentinel, dtype=torch.uint8, device=device ) @@ -418,8 +622,52 @@ def test_indexer_gather_accepts_upper_bound_output(): ops.cp_gather_indexer_k_quant_cache( kv_cache, dst_k, dst_scale, block_table, cu_seq_lens ) + + if current_platform.is_rocm(): + triton_kv_cache = torch.zeros_like(kv_cache) + indexer_k_quant_and_cache_triton( + k, + triton_kv_cache, + slot_mapping, + quant_block_size, + "ue8m0", + ) + triton_dst_k = torch.full_like(dst_k, sentinel) + triton_dst_scale = torch.full_like(dst_scale, sentinel) + token_to_seq = torch.cat( + ( + torch.repeat_interleave( + torch.arange(num_seqs, dtype=torch.int32, device=device), 3 + ), + torch.full( + (upper_bound_tokens - valid_tokens,), + -1, + dtype=torch.int32, + device=device, + ), + ) + ) + cp_gather_indexer_k_quant_cache_triton( + triton_kv_cache, + triton_dst_k.view(current_platform.fp8_dtype()), + triton_dst_scale, + block_table, + cu_seq_lens, + token_to_seq, + ) torch.accelerator.synchronize() + if current_platform.is_rocm(): + triton_recovered = triton_dst_k[:valid_tokens].view( + current_platform.fp8_dtype() + ).float() * triton_dst_scale[:valid_tokens].view(torch.float32) + triton_error = (triton_recovered - k.float()).abs().amax(dim=1) + max_triton_error = ( + 16.0 * triton_dst_scale[:valid_tokens].view(torch.float32).flatten() + ) + assert torch.all(triton_error <= max_triton_error) + assert torch.all(triton_dst_k[valid_tokens:] == sentinel) + assert torch.all(triton_dst_scale[valid_tokens:] == sentinel) k_recovered = dst_k[:valid_tokens].view(torch.float8_e4m3fn).float() * dst_scale[ :valid_tokens ].view(torch.float32) diff --git a/tests/kernels/test_fused_gdn_post_conv.py b/tests/kernels/test_fused_gdn_post_conv.py index fad77891a4dd..b6dacf2e73f2 100644 --- a/tests/kernels/test_fused_gdn_post_conv.py +++ b/tests/kernels/test_fused_gdn_post_conv.py @@ -1,18 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for fused_gdn_prefill_post_conv kernel. +"""Tests for GDN post-convolution kernels. -Verifies that the fused kernel matches the reference: - split → rearrange → contiguous → l2norm → gating +The prefill tests cover preparation and the MTP tests cover recurrent state +updates plus output normalization and gating. """ import pytest import torch import torch.nn.functional as F +from vllm import _custom_ops as ops +from vllm.third_party.flash_linear_attention.ops import ( + fused_sigmoid_gating_delta_rule_update, +) from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import ( fused_post_conv_prep, ) +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( + rmsnorm_fn, +) def reference_post_conv( @@ -207,3 +214,132 @@ def test_fused_post_conv_l0(): ) assert q.shape == (0, H, K) assert g.shape == (0, HV) + + +@pytest.mark.parametrize( + "tp_size,query_lengths,state_dtype,norm_dtype", + [ + pytest.param(16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), + pytest.param(4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), + pytest.param(16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), + pytest.param(4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), + pytest.param(16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), + pytest.param(4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + ], +) +@torch.inference_mode() +def test_fused_gdn_decode_post_conv_mtp_ratio8( + tp_size: int, + query_lengths: tuple[int, ...], + state_dtype: torch.dtype, + norm_dtype: torch.dtype, +) -> None: + if torch.cuda.get_device_capability() < (8, 0): + pytest.skip("fused GDN decode MTP requires compute capability 8.0+") + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): + pytest.skip("fused GDN decode MTP op is not built") + + torch.manual_seed(0) + device = "cuda" + H = 16 // tp_size + HV = 128 // tp_size + K = V = 128 + num_reqs = len(query_lengths) + state_width = max(query_lengths) + num_tokens = sum(query_lengths) + num_slots = num_reqs * state_width + 1 + scale = K**-0.5 + eps = 1e-6 + + mixed_qkv = torch.randn( + num_tokens, + 2 * H * K + HV * V, + dtype=torch.bfloat16, + device=device, + ) + query, key, value = torch.split( + mixed_qkv, + [H * K, H * K, HV * V], + dim=-1, + ) + query = query.view(1, num_tokens, H, K) + key = key.view(1, num_tokens, H, K) + value = value.view(1, num_tokens, HV, V) + ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device) + b, a = ba.chunk(2, dim=-1) + assert not a.is_contiguous() + assert not b.is_contiguous() + A_log = 0.5 * torch.randn(HV, dtype=torch.float32, device=device) + dt_bias = 0.1 * torch.randn(HV, dtype=torch.float32, device=device) + output_gate = torch.randn(num_tokens, HV, V, dtype=torch.bfloat16, device=device) + norm_weight = torch.randn(V, dtype=norm_dtype, device=device) + state_ref = ( + 0.01 * torch.randn(num_slots, HV, V, K, dtype=torch.float32, device=device) + ).to(state_dtype) + state_actual = state_ref.clone() + state_indices = torch.arange(1, num_slots, dtype=torch.int32, device=device).view( + num_reqs, state_width + ) + cu_seqlens = torch.tensor( + [0, *torch.tensor(query_lengths).cumsum(0).tolist()], + dtype=torch.int32, + device=device, + ) + num_accepted_tokens = torch.ones(num_reqs, dtype=torch.int32, device=device) + if query_lengths[-1] == 0: + state_indices[-1].zero_() + + for step, accepted_tokens in enumerate((1, min(2, state_width), state_width)): + num_accepted_tokens.fill_(accepted_tokens) + if query_lengths[-1] == 0: + num_accepted_tokens[-1] = 1 + raw_ref, _ = fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + b=b, + dt_bias=dt_bias, + q=query, + k=key, + v=value, + initial_state=state_ref, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + expected = rmsnorm_fn( + raw_ref.squeeze(0), + norm_weight, + None, + z=output_gate, + eps=eps, + norm_before_gate=True, + activation="silu", + ) + actual = ops.fused_gdn_decode_post_conv_mtp( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=num_accepted_tokens, + state=state_actual, + output_gate=output_gate, + norm_weight=norm_weight, + out=torch.empty_like(output_gate), + scale=scale, + norm_eps=eps, + ) + + output_error = (actual.float() - expected.float()).norm() + output_relative_l2 = output_error / expected.float().norm().clamp_min(1e-20) + assert output_relative_l2 < 5e-4, ( + f"MTP output relative L2 mismatch at step {step}: " + f"{output_relative_l2.item():.6g}" + ) + + torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2) diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index 928c3e5bd748..33a66034b422 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -100,3 +100,26 @@ def test_fused_recurrent_packed_decode_matches_reference( valid = ssm_state_indices > 0 torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol) torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") +def test_packed_decode_supports_large_batch_head_grid(): + B, H, HV, K, V = 1024, 8, 64, 1, 1 + device = torch.device("cuda") + gates = torch.empty((B, HV), device=device) + params = torch.empty((HV,), device=device) + out = torch.empty((B, 1, HV, V), device=device) + + fused_recurrent_gated_delta_rule_packed_decode( + mixed_qkv=torch.empty((B, 2 * H * K + HV * V), device=device), + a=gates, + b=gates, + A_log=params, + dt_bias=params, + scale=1.0, + initial_state=torch.empty((1, HV, V, K), device=device), + out=out, + ssm_state_indices=torch.zeros((B,), device=device, dtype=torch.int32), + ) + + assert torch.count_nonzero(out).item() == 0 diff --git a/tests/kernels/test_kimi_k3_gemm_rs.py b/tests/kernels/test_kimi_k3_gemm_rs.py new file mode 100644 index 000000000000..5ca54d2ca4de --- /dev/null +++ b/tests/kernels/test_kimi_k3_gemm_rs.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from tests.utils import ensure_current_vllm_config +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, +) +from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import GemmRS +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port +from vllm.utils.system_utils import update_environment_variables + +_SHAPES = ( + (129, 768), + (257, 768), + (1023, 4224), + (1024, 4224), + (8191, 4224), +) +_N = 512 + + +def _reference( + x: torch.Tensor, + weight: torch.Tensor, + world_size: int, + group: dist.ProcessGroup, +) -> torch.Tensor: + M = x.shape[0] + padded_M = (M + world_size - 1) // world_size * world_size + partial = torch.empty( + (padded_M, weight.shape[0]), + dtype=x.dtype, + device=x.device, + ) + torch.mm(x, weight.T, out=partial[:M]) + if padded_M > M: + partial[M:].zero_() + + output = torch.empty( + (padded_M // world_size, weight.shape[0]), + dtype=x.dtype, + device=x.device, + ) + dist.reduce_scatter_single(output, partial, group=group) + return output + + +def _assert_valid_rows_close( + actual: torch.Tensor, + expected: torch.Tensor, + M: int, + rank: int, +) -> None: + rank_start = rank * actual.shape[0] + valid_rows = min(max(M - rank_start, 0), actual.shape[0]) + torch.testing.assert_close( + actual[:valid_rows], + expected[:valid_rows], + rtol=5e-2, + atol=4.0, + ) + + +def _worker(local_rank: int, world_size: int, master_port: int) -> None: + device = torch.device("cuda", local_rank) + torch.accelerator.set_device_index(device) + update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(master_port), + } + ) + + init_distributed_environment() + with ensure_current_vllm_config(): + initialize_model_parallel(tensor_model_parallel_size=world_size) + + group = dist.group.WORLD + rank = dist.get_rank(group) + gemm_rs = GemmRS(max_M=max(M for M, _ in _SHAPES), N=_N) + + weight_generator = torch.Generator(device=device) + input_generator = torch.Generator(device=device) + weights = {} + for K in {K for _, K in _SHAPES}: + weight_generator.manual_seed(1000 + rank * 10 + K) + weights[K] = torch.randn( + _N, + K, + dtype=torch.bfloat16, + device=device, + generator=weight_generator, + ) + + # Alternating shapes exercise producer-flag reuse across different grids, + # CTA-group choices, and both BN=128 and BN=256 dispatches. + for M, K in (*_SHAPES, *_SHAPES[::-1]): + input_generator.manual_seed(2000 + M + K) + x = torch.randn( + M, + K, + dtype=torch.bfloat16, + device=device, + generator=input_generator, + ) + expected = _reference(x, weights[K], world_size, group) + actual = gemm_rs(x, weights[K]) + torch.accelerator.synchronize(device) + _assert_valid_rows_close(actual, expected, M, rank) + + # Exercise GEMM-RS under CUDA graph capture and replay. + graph_M, graph_K = 1025, 4224 + input_generator.manual_seed(3000) + graph_x = torch.randn( + graph_M, + graph_K, + dtype=torch.bfloat16, + device=device, + generator=input_generator, + ) + graph_expected = _reference(graph_x, weights[graph_K], world_size, group) + graph_output = torch.empty_like(graph_expected) + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + for _ in range(3): + graph_output.copy_(gemm_rs(graph_x, weights[graph_K])) + capture_stream.synchronize() + dist.barrier(group=group) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + graph_output.copy_(gemm_rs(graph_x, weights[graph_K])) + torch.cuda.current_stream().wait_stream(capture_stream) + dist.barrier(group=group) + + for _ in range(3): + graph.replay() + torch.accelerator.synchronize(device) + _assert_valid_rows_close(graph_output, graph_expected, graph_M, rank) + + dist.barrier(group=group) + del gemm_rs + cleanup_dist_env_and_memory() + + +@pytest.mark.distributed(num_gpus=2) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Kimi-K3 GEMM-RS requires SM100", +) +def test_kimi_k3_gemm_rs(monkeypatch: pytest.MonkeyPatch) -> None: + world_size = 2 + if torch.accelerator.device_count() < world_size: + pytest.skip("GEMM-RS requires two GPUs") + + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "1") + try: + mp.spawn( + _worker, + args=(world_size, get_open_port()), + nprocs=world_size, + ) + finally: + cleanup_dist_env_and_memory() diff --git a/tests/kernels/test_mhc_tilelang_jit.py b/tests/kernels/test_mhc_tilelang_jit.py new file mode 100644 index 000000000000..15560165ce2d --- /dev/null +++ b/tests/kernels/test_mhc_tilelang_jit.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import ctypes +import importlib +import importlib.util +import sys +from types import ModuleType +from typing import Any + +import pytest + +from vllm.platforms import current_platform +from vllm.utils import import_utils + + +class _PassConfigKey: + TL_DISABLE_WARP_SPECIALIZED = "disable_warp_specialized" + TL_DISABLE_TMA_LOWER = "disable_tma_lower" + TL_PTXAS_REGISTER_USAGE_LEVEL = "ptxas_register_usage_level" + + +def _install_tilelang_stub( + monkeypatch: pytest.MonkeyPatch, +) -> dict[str, int]: + calls = {"jit_decorate": 0, "compiled_call": 0} + + tilelang: Any = ModuleType("tilelang") + + def jit(**kwargs: Any) -> Any: + def decorate(func: Any) -> Any: + calls["jit_decorate"] += 1 + + def compiled(*args: Any, **kw: Any) -> Any: + calls["compiled_call"] += 1 + return func.__name__ + + return compiled + + return decorate + + tilelang.PassConfigKey = _PassConfigKey + tilelang.jit = jit + + monkeypatch.setattr(import_utils, "has_tilelang", lambda: True) + monkeypatch.setitem(sys.modules, "tilelang", tilelang) + monkeypatch.setitem( + sys.modules, "tilelang.language", ModuleType("tilelang.language") + ) + monkeypatch.delitem(sys.modules, "vllm.tilelang_utils", raising=False) + + return calls + + +def test_tilelang_jit_decorator_is_lazy_only_on_rocm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + if not (current_platform.is_cuda() or current_platform.is_rocm()): + pytest.skip("Test requires CUDA or ROCm") + + calls = _install_tilelang_stub(monkeypatch) + module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels" + monkeypatch.delitem(sys.modules, module_name, raising=False) + module = importlib.import_module(module_name) + + if current_platform.is_rocm(): + assert calls["jit_decorate"] == 0 + else: + assert calls["jit_decorate"] > 0 + + decorated_calls = calls["jit_decorate"] + assert module.mhc_post_tilelang() == "mhc_post_tilelang" + if current_platform.is_rocm(): + assert calls["jit_decorate"] == 1 + else: + assert calls["jit_decorate"] == decorated_calls + assert calls["compiled_call"] == 1 + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm") +def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols( + monkeypatch: pytest.MonkeyPatch, +) -> None: + if importlib.util.find_spec("tilelang") is None: + pytest.skip("Test requires TileLang to be installed") + + class DlInfo(ctypes.Structure): + _fields_ = [ + ("dli_fname", ctypes.c_char_p), + ("dli_fbase", ctypes.c_void_p), + ("dli_sname", ctypes.c_char_p), + ("dli_saddr", ctypes.c_void_p), + ] + + libdl = ctypes.CDLL("libdl.so.2") + dlsym = libdl.dlsym + dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p] + dlsym.restype = ctypes.c_void_p + dladdr = libdl.dladdr + dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(DlInfo)] + dladdr.restype = ctypes.c_int + + from vllm.model_executor.layers import mhc # noqa: F401 + from vllm.models import deepseek_v4 # noqa: F401 + from vllm.utils import jit_monitor + + monkeypatch.setattr(jit_monitor, "_active", False) + monkeypatch.setattr(jit_monitor, "_setup_triton_autotuning_print", lambda: None) + monkeypatch.setattr(jit_monitor, "_setup_triton_jit_hook", lambda: None) + monkeypatch.setattr(jit_monitor, "_setup_cutedsl_jit_hook", lambda: None) + monkeypatch.setattr( + jit_monitor, + "_setup_tilelang_jit_hook", + lambda: pytest.fail("TileLang JIT monitor must not run on ROCm"), + ) + jit_monitor.activate() + + assert not any( + name == "tilelang" or name.startswith("tilelang.") for name in sys.modules + ) + + address = dlsym(None, b"hipFree") + assert address is not None, "hipFree is not available in the global symbol table" + info = DlInfo() + assert dladdr(address, ctypes.byref(info)) + source = info.dli_fname.decode() if info.dli_fname else "" + assert "libhip_stub.so" not in source, source + assert "libamdhip64.so" in source, source diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index 899b3129d62f..f14b488a1729 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -253,6 +253,11 @@ def qwen3vl_vision_lora_files(): ) +@pytest.fixture(scope="session") +def gemma4_vision_lora_files(): + return hf_api().snapshot_download(repo_id="EpochEcho/gemma4-e2b-it-lora-pokemon") + + @pytest.fixture(scope="session") def qwen3_meowing_lora_files(): """Download Qwen3 Meow LoRA files once per test session.""" diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py index a70c5434736f..444b609492ba 100644 --- a/tests/lora/test_fused_moe_lora_kernel.py +++ b/tests/lora/test_fused_moe_lora_kernel.py @@ -22,6 +22,10 @@ from vllm.utils.network_utils import get_open_port from vllm.utils.torch_utils import set_random_seed +# The tensor-parallel cases build their distributed environment inside processes +# spawned by torch.multiprocessing.spawn, so this process never has one to tear down. +pytestmark = pytest.mark.skip_global_cleanup + @pytest.fixture(autouse=True) def reset_device(reset_default_device): diff --git a/tests/lora/test_gemma4_tp.py b/tests/lora/test_gemma4_tp.py new file mode 100644 index 000000000000..efbee00d6249 --- /dev/null +++ b/tests/lora/test_gemma4_tp.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE To avoid overloading the CI pipeline, this test script will not +# be triggered on CI and is primarily intended for local testing and verification. + +import vllm +from vllm.assets.image import ImageAsset +from vllm.lora.request import LoRARequest + +from ..utils import multi_gpu_test + +MODEL_PATH = "google/gemma-4-E2B-it" + +PROMPT_TEMPLATE = """<|turn>user +<|image|>What is in the image? +<|turn>model +""" + +TEST_IMAGES = [ + ImageAsset("stop_sign"), + ImageAsset("cherry_blossom"), +] + +EXPECTED_OUTPUTS_VISION = [ + "A red stop sign stands prominently in the foreground.", + "A majestic skyscraper stands tall, partially obscured by a vibrant " + "canopy of cherry blossoms, against a clear blue sky.", +] + + +def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: + prompts = [ + { + "prompt": PROMPT_TEMPLATE, + "multi_modal_data": {"image": asset.pil_image}, + } + for asset in TEST_IMAGES + ] + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=128) + outputs = llm.generate( + prompts, + sampling_params, + lora_request=LoRARequest(str(lora_id), lora_id, lora_path), + ) + + generated_texts = [output.outputs[0].text.strip() for output in outputs] + for generated, expected in zip(generated_texts, EXPECTED_OUTPUTS_VISION): + assert generated.startswith(expected), ( + f"Generated text {generated!r} does not match expected {expected!r}" + ) + + +def test_gemma4_lora(gemma4_vision_lora_files): + llm = vllm.LLM( + MODEL_PATH, + max_model_len=1024, + enable_lora=True, + enforce_eager=True, + max_loras=4, + trust_remote_code=True, + limit_mm_per_prompt={"image": 1}, + mm_processor_cache_gb=0, + enable_tower_connector_lora=True, + ) + + generate_and_test(llm, gemma4_vision_lora_files, lora_id=1) + generate_and_test(llm, gemma4_vision_lora_files, lora_id=2) + + +@multi_gpu_test(num_gpus=2) +def test_gemma4_lora_tp2(gemma4_vision_lora_files): + llm = vllm.LLM( + MODEL_PATH, + max_model_len=1024, + enforce_eager=True, + enable_lora=True, + max_loras=4, + trust_remote_code=True, + tensor_parallel_size=2, + limit_mm_per_prompt={"image": 1}, + mm_processor_cache_gb=0, + enable_tower_connector_lora=True, + ) + + generate_and_test(llm, gemma4_vision_lora_files, lora_id=1) + generate_and_test(llm, gemma4_vision_lora_files, lora_id=2) + + +@multi_gpu_test(num_gpus=4) +def test_gemma4_lora_tp4(gemma4_vision_lora_files): + llm = vllm.LLM( + MODEL_PATH, + max_model_len=1024, + enforce_eager=True, + enable_lora=True, + max_loras=4, + trust_remote_code=True, + tensor_parallel_size=4, + limit_mm_per_prompt={"image": 1}, + mm_processor_cache_gb=0, + enable_tower_connector_lora=True, + ) + + generate_and_test(llm, gemma4_vision_lora_files, lora_id=1) + generate_and_test(llm, gemma4_vision_lora_files, lora_id=2) diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index ab09514274d8..bce91510d428 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -56,14 +56,19 @@ torch.bfloat16: (3e-2, 2e-2), } -pytestmark = pytest.mark.skipif( - not ( - current_platform.is_cuda_alike() - or current_platform.is_cpu() - or current_platform.is_xpu() +pytestmark = [ + pytest.mark.skipif( + not ( + current_platform.is_cuda_alike() + or current_platform.is_cpu() + or current_platform.is_xpu() + ), + reason="Backend not supported", ), - reason="Backend not supported", -) + # Tests here either take dist_init, which tears the distributed environment + # down itself, or never build one, so the global cleanup only repeats it. + pytest.mark.skip_global_cleanup, +] DEVICE_TYPE = current_platform.device_type DEVICES = ( diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index f94b54d9fb17..9e1f99b3f6a3 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -15,6 +15,14 @@ DEVICE_TYPE = current_platform.device_type +# On XPU, oneDNN/oneMKL can return wrong results for these reference matmuls +# after many Triton kernel launches, so the reference stays on CPU there. +_REF_ON_CPU = current_platform.is_xpu() + + +def _to_ref_device(tensor: torch.Tensor) -> torch.Tensor: + return tensor.cpu() if _REF_ON_CPU else tensor + @pytest.fixture(autouse=True) def reset_device(reset_default_device): @@ -33,10 +41,10 @@ def dynamo_reset(): yield -def _cpu_bgmv_shrink( +def _bgmv_shrink( inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0 ): - """Memory-efficient shrink reference: per-LoRA matmul loop on CPU. + """Memory-efficient shrink reference: per-LoRA matmul loop. output[mask] = scaling * inputs[mask] @ weight.T""" exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): @@ -48,7 +56,7 @@ def _cpu_bgmv_shrink( output[mask] = scaling * (inp @ w.T) -def _cpu_bgmv_expand( +def _bgmv_expand( inputs, lora_weight, output, @@ -57,7 +65,7 @@ def _cpu_bgmv_expand( offset=0, add_inputs=False, ): - """Memory-efficient expand reference: per-LoRA matmul loop on CPU. + """Memory-efficient expand reference: per-LoRA matmul loop. output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T""" exploded = torch.repeat_interleave(lora_indices, seq_len_tensor) for lid in exploded.unique(): @@ -88,21 +96,22 @@ def sgmv_shrink_for_nslices( num_tokens: int, scaling: float, ): - """CPU reference for sgmv_shrink using per-LoRA matmul loop.""" - inp_cpu = inputs_tensor.cpu() - seq_cpu = seq_len_tensor.cpu() - idx_cpu = prompt_lora_mapping.cpu() - out_cpu = out_tensor.cpu() + """Reference for sgmv_shrink using per-LoRA matmul loop.""" + inputs = _to_ref_device(inputs_tensor) + seq_len = _to_ref_device(seq_len_tensor) + mapping = _to_ref_device(prompt_lora_mapping) + out = _to_ref_device(out_tensor) for index in range(nslices): - _cpu_bgmv_shrink( - inp_cpu, - lora_weights_lst[index].cpu(), - out_cpu[index], - seq_cpu, - idx_cpu, + _bgmv_shrink( + inputs, + _to_ref_device(lora_weights_lst[index]), + out[index], + seq_len, + mapping, scaling=scaling, ) - out_tensor.copy_(out_cpu) + if _REF_ON_CPU: + out_tensor.copy_(out) def sgmv_expand_for_nslices( @@ -119,21 +128,22 @@ def sgmv_expand_for_nslices( num_tokens: int, add_inputs: bool, ) -> None: - """CPU reference for sgmv_expand using per-LoRA matmul loop.""" - seq_cpu = seq_len_tensor.cpu() - idx_cpu = prompt_lora_mapping.cpu() - out_cpu = out_tensor.cpu() + """Reference for sgmv_expand using per-LoRA matmul loop.""" + seq_len = _to_ref_device(seq_len_tensor) + mapping = _to_ref_device(prompt_lora_mapping) + out = _to_ref_device(out_tensor) for index in range(nslices): - _cpu_bgmv_expand( - inputs_tensor[index].cpu(), - lora_weights_lst[index].cpu(), - out_cpu, - seq_cpu, - idx_cpu, + _bgmv_expand( + _to_ref_device(inputs_tensor[index]), + _to_ref_device(lora_weights_lst[index]), + out, + seq_len, + mapping, offset=hidden_size * index, add_inputs=add_inputs, ) - out_tensor.copy_(out_cpu) + if _REF_ON_CPU: + out_tensor.copy_(out) _dict_lock = Lock() diff --git a/tests/lora/test_punica_ops_fp8.py b/tests/lora/test_punica_ops_fp8.py index 3e7fe7b27582..e447cef9d8a8 100644 --- a/tests/lora/test_punica_ops_fp8.py +++ b/tests/lora/test_punica_ops_fp8.py @@ -17,6 +17,7 @@ import pytest import torch +import torch.nn.functional as F import vllm.lora.ops.torch_ops as torch_ops import vllm.lora.ops.triton_ops as triton_ops @@ -31,6 +32,8 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +pytestmark = pytest.mark.skip_global_cleanup + DEVICE_TYPE = current_platform.device_type DEVICES = [f"{DEVICE_TYPE}:{0}"] SEED = [0] @@ -206,43 +209,32 @@ def quantize_to_fp8_blockwise( if tensor.ndim == 2: M, K = tensor.shape n_blocks_k = math.ceil(K / group_k) - scale = torch.zeros(M, n_blocks_k, dtype=torch.float32, device=tensor.device) - fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE) - for m in range(M): - for bk in range(n_blocks_k): - k_start = bk * group_k - k_end = min(k_start + group_k, K) - block = tensor[m, k_start:k_end].float() - amax = block.abs().max().clamp(min=1e-12) - s = (amax / FP8_MAX).to(torch.float32) - scale[m, bk] = s - fp8_tensor[m, k_start:k_end] = ( - (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) - ) - return fp8_tensor, scale + # Zero padding leaves the per-block amax unchanged, so the padded + # columns are dropped again after quantizing. + padded = F.pad(tensor.float(), (0, n_blocks_k * group_k - K)) + blocks = padded.view(M, n_blocks_k, group_k) + scale = blocks.abs().amax(dim=-1).clamp(min=1e-12) / FP8_MAX + fp8_tensor = ( + (blocks / scale.unsqueeze(-1)).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) + ) + return fp8_tensor.view(M, -1)[:, :K].contiguous(), scale elif tensor.ndim == 3: L, N, K = tensor.shape n_blocks_n = math.ceil(N / group_n) n_blocks_k = math.ceil(K / group_k) - scale = torch.zeros( - L, n_blocks_n, n_blocks_k, dtype=torch.float32, device=tensor.device + padded = F.pad( + tensor.float(), + (0, n_blocks_k * group_k - K, 0, n_blocks_n * group_n - N), + ) + blocks = padded.view(L, n_blocks_n, group_n, n_blocks_k, group_k) + scale = blocks.abs().amax(dim=(2, 4)).clamp(min=1e-12) / FP8_MAX + fp8_tensor = ( + (blocks / scale[:, :, None, :, None]) + .clamp(FP8_MIN, FP8_MAX) + .to(FP8_DTYPE) + .view(L, n_blocks_n * group_n, n_blocks_k * group_k) ) - fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE) - for li in range(L): - for bn in range(n_blocks_n): - for bk in range(n_blocks_k): - n_start = bn * group_n - n_end = min(n_start + group_n, N) - k_start = bk * group_k - k_end = min(k_start + group_k, K) - block = tensor[li, n_start:n_end, k_start:k_end].float() - amax = block.abs().max().clamp(min=1e-12) - s = (amax / FP8_MAX).to(torch.float32) - scale[li, bn, bk] = s - fp8_tensor[li, n_start:n_end, k_start:k_end] = ( - (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) - ) - return fp8_tensor, scale + return fp8_tensor[:, :N, :K].contiguous(), scale else: raise ValueError(f"Unsupported tensor ndim: {tensor.ndim}") @@ -301,33 +293,23 @@ def dequantize_fp8_blockwise( """Dequantize FP8 tensor with block-wise scale back to output_dtype.""" if fp8_tensor.ndim == 2: M, K = fp8_tensor.shape - out = torch.zeros(M, K, dtype=output_dtype, device=fp8_tensor.device) n_blocks_k = math.ceil(K / group_k) - for m in range(M): - for bk in range(n_blocks_k): - k_start = bk * group_k - k_end = min(k_start + group_k, K) - out[m, k_start:k_end] = ( - fp8_tensor[m, k_start:k_end].float() * scale[m, bk].float() - ).to(output_dtype) - return out + padded = F.pad(fp8_tensor.float(), (0, n_blocks_k * group_k - K)) + blocks = padded.view(M, n_blocks_k, group_k) + out = (blocks * scale.float().unsqueeze(-1)).to(output_dtype) + return out.view(M, -1)[:, :K].contiguous() elif fp8_tensor.ndim == 3: L, N, K = fp8_tensor.shape - out = torch.zeros(L, N, K, dtype=output_dtype, device=fp8_tensor.device) n_blocks_n = math.ceil(N / group_n) n_blocks_k = math.ceil(K / group_k) - for l_idx in range(L): - for bn in range(n_blocks_n): - for bk in range(n_blocks_k): - n_start = bn * group_n - n_end = min(n_start + group_n, N) - k_start = bk * group_k - k_end = min(k_start + group_k, K) - out[l_idx, n_start:n_end, k_start:k_end] = ( - fp8_tensor[l_idx, n_start:n_end, k_start:k_end].float() - * scale[l_idx, bn, bk].float() - ).to(output_dtype) - return out + padded = F.pad( + fp8_tensor.float(), + (0, n_blocks_k * group_k - K, 0, n_blocks_n * group_n - N), + ) + blocks = padded.view(L, n_blocks_n, group_n, n_blocks_k, group_k) + out = (blocks * scale.float()[:, :, None, :, None]).to(output_dtype) + out = out.view(L, n_blocks_n * group_n, n_blocks_k * group_k) + return out[:, :N, :K].contiguous() else: raise ValueError(f"Unsupported tensor ndim: {fp8_tensor.ndim}") @@ -519,45 +501,22 @@ def generate_fp8_expand_data( # shared across slices. Compute shared scale across slices, then quantize. # First compute per-token-per-block scale across all slices n_blocks_k = math.ceil(rank / group_k) - a_scale = torch.zeros( - total_tokens, n_blocks_k, dtype=torch.float32, device=device + padded = F.pad(inputs_bf16.float(), (0, n_blocks_k * group_k - rank)) + blocks = padded.view(nslices, total_tokens, n_blocks_k, group_k) + # Take the block amax across all slices so every slice shares the scale. + a_scale = blocks.abs().amax(dim=-1).clamp(min=1e-12).amax(dim=0) / FP8_MAX + fp8_blocks = ( + (blocks / a_scale[None, :, :, None]).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) + ) + inputs_fp8 = fp8_blocks.view(nslices, total_tokens, -1)[ + :, :, :rank + ].contiguous() + inputs_dequant = ( + (fp8_blocks.float() * a_scale[None, :, :, None]) + .to(dtype) + .view(nslices, total_tokens, -1)[:, :, :rank] + .contiguous() ) - for m in range(total_tokens): - for bk in range(n_blocks_k): - k_start = bk * group_k - k_end = min(k_start + group_k, rank) - # Max across all slices for this token and block - block_amax = torch.tensor(0.0, device=device) - for s in range(nslices): - block = inputs_bf16[s, m, k_start:k_end].float() - block_amax = torch.max( - block_amax, block.abs().max().clamp(min=1e-12) - ) - a_scale[m, bk] = (block_amax / FP8_MAX).to(torch.float32) - - # Quantize all slices with the shared scale - inputs_fp8_list = [] - inputs_dequant_list = [] - for s in range(nslices): - slice_2d = inputs_bf16[s] # (total_tokens, rank) - fp8_slice = torch.zeros_like(slice_2d, dtype=FP8_DTYPE) - dequant_slice = torch.zeros_like(slice_2d) - for m in range(total_tokens): - for bk in range(n_blocks_k): - k_start = bk * group_k - k_end = min(k_start + group_k, rank) - block = slice_2d[m, k_start:k_end].float() - s_val = a_scale[m, bk] - fp8_slice[m, k_start:k_end] = ( - (block / s_val).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE) - ) - dequant_slice[m, k_start:k_end] = ( - fp8_slice[m, k_start:k_end].float() * s_val.float() - ).to(dtype) - inputs_fp8_list.append(fp8_slice) - inputs_dequant_list.append(dequant_slice) - inputs_fp8 = torch.stack(inputs_fp8_list, dim=0) - inputs_dequant = torch.stack(inputs_dequant_list, dim=0) elif quant_mode == "per_tensor": # Per-tensor: kernel loads a single scalar from a_scale_ptr inputs_fp8_2d, a_scale = quantize_to_fp8_per_tensor(inputs_2d_all) diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py new file mode 100644 index 000000000000..1efc521bd8b4 --- /dev/null +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -0,0 +1,831 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import importlib +import types +from dataclasses import dataclass + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + _LINEAR_BACKEND_KERNEL_MAP, + _POSSIBLE_FP8_BLOCK_KERNELS, + _POSSIBLE_FP8_KERNELS, + _POSSIBLE_MXFP4_KERNELS, + _POSSIBLE_MXFP8_KERNELS, + _POSSIBLE_NVFP4_KERNELS, + B12xFp8BlockScaledMMKernel, + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, + B12xTensorFP8ScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, + Mxfp8LinearLayerConfig, + init_fp8_linear_kernel, + init_mxfp4_linear_kernel, + init_mxfp8_linear_kernel, + init_nvfp4_linear_kernel, +) +from vllm.model_executor.kernels.linear.nvfp4.marlin import ( + MarlinNvFp4LinearKernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, + kMxfp4Dynamic, +) +from vllm.platforms import PlatformEnum + + +@pytest.mark.parametrize( + ("kernel_cls", "kernels", "before", "after", "initializer", "kwargs"), + [ + ( + B12xMxFp4LinearKernel, + _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA], + "HummingMxFp4LinearKernel", + "EmulationMxfp4LinearKernel", + init_mxfp4_linear_kernel, + {"activation_quant_key": kMxfp4Dynamic}, + ), + ( + B12xNvFp4LinearKernel, + _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA], + "FbgemmNvFp4LinearKernel", + "EmulationNvFp4LinearKernel", + init_nvfp4_linear_kernel, + {}, + ), + ( + B12xMxfp8LinearKernel, + _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA], + "MarlinMxfp8LinearKernel", + "EmulationMxfp8LinearKernel", + init_mxfp8_linear_kernel, + {}, + ), + ( + B12xTensorFP8ScaledMMLinearKernel, + _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA], + "CutlassFP8ScaledMMLinearKernel", + "PerTensorTorchFP8ScaledMMLinearKernel", + init_fp8_linear_kernel, + { + "activation_quant_key": kFp8StaticTensorSym, + "weight_quant_key": kFp8StaticTensorSym, + "input_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "weight_shape": (2048, 2048), + }, + ), + ( + B12xFp8BlockScaledMMKernel, + _POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA], + "CutlassFp8BlockScaledMMKernel", + "MarlinFP8ScaledMMLinearKernel", + init_fp8_linear_kernel, + { + "activation_quant_key": kFp8Dynamic128Sym, + "weight_quant_key": kFp8Static128BlockSym, + "input_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "weight_shape": (2048, 2048), + }, + ), + ], +) +def test_b12x_backend_registration_priority_and_selection( + monkeypatch, + default_vllm_config, + kernel_cls, + kernels, + before: str, + after: str, + initializer, + kwargs: dict, +) -> None: + import vllm.model_executor.kernels.linear as linear_mod + + assert kernel_cls in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + names = [kernel.__name__ for kernel in kernels] + assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after) + + monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) + monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + kernel_cls, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + monkeypatch.setattr( + kernel_cls, + "can_implement", + classmethod(lambda cls, config: (True, None)), + ) + + assert isinstance(initializer(**kwargs), kernel_cls) + + +def test_b12x_tensor_fp8_can_implement_supported_config() -> None: + config = FP8ScaledMMLinearLayerConfig( + activation_quant_key=kFp8StaticTensorSym, + weight_quant_key=kFp8StaticTensorSym, + weight_shape=(64, 128), + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + ) + + can_implement, reason = B12xTensorFP8ScaledMMLinearKernel.can_implement(config) + + assert can_implement + assert reason is None + + +def test_b12x_block_fp8_checks_runtime_support(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + platform = types.SimpleNamespace( + is_cuda=lambda: True, + is_device_capability_family=lambda family: family == 120, + ) + monkeypatch.setattr(b12x_mod, "current_platform", platform) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(is_supported=lambda: False), + ) + + supported, reason = B12xFp8BlockScaledMMKernel.is_supported() + + assert not supported + assert reason == "B12X regular block-FP8 GEMM is not supported" + + +def test_b12x_block_fp8_requires_matching_supported_dtypes() -> None: + def config(input_dtype: torch.dtype, out_dtype: torch.dtype): + return FP8ScaledMMLinearLayerConfig( + activation_quant_key=kFp8Dynamic128Sym, + weight_quant_key=kFp8Static128BlockSym, + weight_shape=(256, 128), + input_dtype=input_dtype, + out_dtype=out_dtype, + ) + + can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement( + config(torch.float32, torch.float32) + ) + assert not can_implement + assert reason == "Supports only bf16/fp16 input dtype" + + can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement( + config(torch.bfloat16, torch.float16) + ) + assert not can_implement + assert reason == "Input and output dtype must match" + + can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement( + config(torch.float16, torch.float16) + ) + assert can_implement + assert reason is None + + +def test_b12x_block_fp8_requires_aligned_features() -> None: + def can_implement(weight_shape: tuple[int, int]): + config = FP8ScaledMMLinearLayerConfig( + activation_quant_key=kFp8Dynamic128Sym, + weight_quant_key=kFp8Static128BlockSym, + weight_shape=weight_shape, + input_dtype=torch.bfloat16, + out_dtype=torch.bfloat16, + ) + return B12xFp8BlockScaledMMKernel.can_implement(config) + + assert can_implement((256, 192)) == ( + False, + "Input features must be a positive multiple of 128", + ) + assert can_implement((192, 256)) == ( + False, + "Output features must be a positive multiple of 128", + ) + + +def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout( + monkeypatch, +) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + packed = types.SimpleNamespace(out_features=64) + + def pack(weight: torch.Tensor, output_scale: torch.Tensor): + calls.append((weight, output_scale)) + return packed + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_tensor_fp8", + lambda: types.SimpleNamespace(pack_weight=pack), + ) + layer = torch.nn.Module() + layer.prefix = "model.layers.0.self_attn.qkv_proj" + original_weight = ( + torch.randn((128, 64), dtype=torch.float32).clamp(-4, 4).to(torch.float8_e4m3fn) + ) + layer.weight = torch.nn.Parameter(original_weight, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False) + layer.input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False) + weight_loader = object() + scale_loader = object() + layer.weight.weight_loader = weight_loader + layer.weight_scale.weight_loader = scale_loader + kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel) + kernel.config = types.SimpleNamespace(weight_shape=(64, 128)) + kernel.layer_param_names = ( + "weight", + "weight_scale", + "input_scale", + "input_scale_ub", + ) + + kernel.process_weights_after_loading(layer) + + assert layer.b12x_tensor_fp8_packed_weight is packed + assert layer.b12x_warmup_provider is kernel + assert len(calls) == 1 + weight, output_scale = calls[0] + torch.testing.assert_close(weight, original_weight.T.contiguous()) + torch.testing.assert_close(output_scale, torch.tensor([0.125])) + assert layer.weight.numel() == 0 + assert layer.weight_scale.numel() == 0 + assert layer.weight.weight_loader is weight_loader + assert layer.weight_scale.weight_loader is scale_loader + torch.testing.assert_close(layer.input_scale, torch.tensor(0.5)) + + +def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight( + monkeypatch, +) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + + def mm( + source: torch.Tensor, + packed_weight, + *, + bias: torch.Tensor | None = None, + out_dtype: torch.dtype, + expected_m: int, + stream: object = None, + ) -> torch.Tensor: + del stream + calls.append((source, packed_weight, bias, out_dtype, expected_m)) + return torch.full( + (source.shape[0], packed_weight.out_features), + 3.0, + dtype=out_dtype, + ) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_tensor_fp8", + lambda: types.SimpleNamespace(mm=mm), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: types.SimpleNamespace(cuda_stream=object()), + ) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) + + layer = torch.nn.Module() + packed = types.SimpleNamespace(out_features=48) + layer.b12x_tensor_fp8_packed_weight = packed + layer.weight = torch.nn.Parameter( + torch.empty((128, 48), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False) + layer.input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False) + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + x_q = torch.empty((6, 128), dtype=torch.float8_e4m3fn) + bias = torch.empty((48,), dtype=torch.bfloat16) + kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel) + kernel.config = types.SimpleNamespace(out_dtype=torch.bfloat16) + kernel.layer_param_names = ( + "weight", + "weight_scale", + "input_scale", + "input_scale_ub", + ) + kernel.quant_fp8 = lambda source, scale, scale_ub: (x_q, scale) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + assert output.dtype == torch.bfloat16 + assert len(calls) == 1 + source, called_packed, called_bias, out_dtype, expected_m = calls[0] + assert source.data_ptr() == x_q.data_ptr() + assert called_packed is packed + assert called_bias is bias + assert out_dtype == torch.bfloat16 + assert expected_m == 6 + + +def test_b12x_mxfp8_can_implement_supported_config() -> None: + can_implement, reason = B12xMxfp8LinearKernel.can_implement( + Mxfp8LinearLayerConfig() + ) + + assert can_implement + assert reason is None + + +def test_b12x_mxfp8_support_check_reports_missing_import(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + b12x_mod.current_platform, + "is_device_capability_family", + lambda family: family == 120, + ) + monkeypatch.setattr(b12x_mod, "_import_b12x_mxfp8", lambda: None) + + is_supported, reason = B12xMxfp8LinearKernel.is_supported() + + assert not is_supported + assert reason == "Install the B12X backend with `pip install vllm[b12x]`" + + +def test_b12x_mxfp8_support_respects_runtime_probe(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + monkeypatch.setattr(b12x_mod.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + b12x_mod.current_platform, + "is_device_capability_family", + lambda family: family == 120, + ) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: types.SimpleNamespace(is_supported=lambda: False), + ) + + is_supported, reason = B12xMxfp8LinearKernel.is_supported() + + assert not is_supported + assert reason == "b12x.gemm.mxfp8_linear is not supported" + + +def test_b12x_mxfp8_process_weights_packs_modelopt_layout(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + calls = [] + packed = types.SimpleNamespace(out_features=48) + + def pack(weight: torch.Tensor, weight_scale: torch.Tensor): + calls.append((weight, weight_scale)) + return packed + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: types.SimpleNamespace(pack_weight=pack), + ) + + layer = torch.nn.Module() + layer.prefix = "model.layers.0.self_attn.qkv_proj" + layer.weight = torch.nn.Parameter( + torch.empty((48, 128), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.empty((64, 8), dtype=torch.uint8), + requires_grad=False, + ) + weight_loader = object() + scale_loader = object() + layer.weight.weight_loader = weight_loader + layer.weight_scale.weight_loader = scale_loader + kernel = object.__new__(B12xMxfp8LinearKernel) + + kernel.process_weights_after_loading(layer) + + assert layer.b12x_mxfp8_packed_weight is packed + assert layer.b12x_warmup_provider is kernel + assert len(calls) == 1 + weight, weight_scale = calls[0] + assert weight.shape == (48, 128) + assert weight_scale.shape == (48, 4) + assert weight.dtype == torch.float8_e4m3fn + assert weight_scale.dtype == torch.uint8 + assert layer.weight.numel() == 0 + assert layer.weight_scale.numel() == 0 + assert layer.weight.weight_loader is weight_loader + assert layer.weight_scale.weight_loader is scale_loader + + +def test_b12x_mxfp8_reload_reuses_packed_tensor_addresses(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + @dataclass(frozen=True) + class PackedWeight: + values: torch.Tensor + scales: torch.Tensor + out_features: int + + def pack(weight: torch.Tensor, weight_scale: torch.Tensor) -> PackedWeight: + return PackedWeight( + values=weight.clone(), + scales=weight_scale.clone(), + out_features=int(weight.shape[0]), + ) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: types.SimpleNamespace(pack_weight=pack), + ) + layer = torch.nn.Module() + layer.prefix = "model.layers.0.mlp.down_proj" + layer.weight = torch.nn.Parameter( + torch.zeros((48, 128), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.zeros((48, 4), dtype=torch.uint8), + requires_grad=False, + ) + kernel = object.__new__(B12xMxfp8LinearKernel) + + kernel.process_weights_after_loading(layer) + packed = layer.b12x_mxfp8_packed_weight + values_ptr = packed.values.data_ptr() + scales_ptr = packed.scales.data_ptr() + + layer.weight = torch.nn.Parameter( + torch.ones((48, 128), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + torch.full((48, 4), 3, dtype=torch.uint8), + requires_grad=False, + ) + kernel.process_weights_after_loading(layer) + + assert layer.b12x_mxfp8_packed_weight is packed + assert packed.values.data_ptr() == values_ptr + assert packed.scales.data_ptr() == scales_ptr + torch.testing.assert_close( + packed.values, + torch.ones((48, 128), dtype=torch.float8_e4m3fn), + ) + torch.testing.assert_close( + packed.scales, + torch.full((48, 4), 3, dtype=torch.uint8), + ) + assert layer.weight.numel() == 0 + assert layer.weight_scale.numel() == 0 + + +@pytest.fixture +def _mock_b12x_cuda_fp8_platform(monkeypatch: pytest.MonkeyPatch) -> None: + import vllm.model_executor.layers.quantization.utils.fp8_utils as fp8_utils + + monkeypatch.setattr( + fp8_utils, + "current_platform", + types.SimpleNamespace( + is_fp8_fnuz=lambda: False, + is_rocm=lambda: False, + fp8_dtype=lambda: torch.float8_e4m3fn, + is_xpu=lambda: False, + is_cuda_alike=lambda: True, + ), + ) + + +@pytest.mark.usefixtures("_mock_b12x_cuda_fp8_platform") +def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None: + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.empty((128, 128), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + layer.weight_scale_inv = torch.nn.Parameter( + torch.empty((1, 1), dtype=torch.float32), + requires_grad=False, + ) + layer.weight_block_size = [128, 128] + weight_loader = object() + scale_loader = object() + layer.weight.weight_loader = weight_loader + layer.weight_scale_inv.weight_loader = scale_loader + kernel = object.__new__(B12xFp8BlockScaledMMKernel) + + kernel.process_weights_after_loading(layer) + + assert layer.b12x_warmup_provider is kernel + assert layer.weight.shape == (128, 128) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight_scale_inv.shape == (1, 1) + assert layer.weight_scale_inv.dtype == torch.float32 + assert layer.weight.weight_loader is weight_loader + assert layer.weight_scale_inv.weight_loader is scale_loader + + +@pytest.mark.parametrize("scale_dtype", [torch.float8_e8m0fnu, torch.uint8]) +@pytest.mark.usefixtures("_mock_b12x_cuda_fp8_platform") +def test_b12x_block_fp8_upcasts_e8m0_weight_scales(scale_dtype) -> None: + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.empty((128, 128), dtype=torch.float8_e4m3fn), + requires_grad=False, + ) + scale_bytes = torch.tensor([[125]], dtype=torch.uint8) + layer.weight_scale_inv = torch.nn.Parameter( + scale_bytes.view(scale_dtype), + requires_grad=False, + ) + layer.weight_block_size = [128, 128] + kernel = object.__new__(B12xFp8BlockScaledMMKernel) + + kernel.process_weights_after_loading(layer) + + assert layer.weight_scale_inv.dtype == torch.float32 + torch.testing.assert_close( + layer.weight_scale_inv, + torch.tensor([[0.25]], dtype=torch.float32), + ) + + +def test_b12x_mxfp8_apply_uses_packed_weight(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + calls = [] + + def mxfp8_linear( + source: torch.Tensor, + packed_weight, + *, + bias: torch.Tensor | None = None, + expected_m: int | None = None, + stream: object = None, + ) -> torch.Tensor: + del stream + calls.append((source, packed_weight, bias, expected_m)) + return source.new_full((source.shape[0], packed_weight.out_features), 3.0) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: types.SimpleNamespace(mm=mxfp8_linear), + ) + + layer = torch.nn.Module() + packed = types.SimpleNamespace(out_features=48) + layer.b12x_mxfp8_packed_weight = packed + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + bias = torch.empty((48,), dtype=torch.bfloat16) + kernel = object.__new__(B12xMxfp8LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + assert output.dtype == x.dtype + assert len(calls) == 1 + source, called_packed, called_bias, expected_m = calls[0] + assert source.shape == (6, 128) + assert called_packed is packed + assert called_bias is bias + assert expected_m == 6 + + +def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + + def mm_block_fp8(*args, **kwargs): + calls.append((args, kwargs)) + return torch.full( + (args[0].shape[0], args[2].shape[0]), + 13.0, + dtype=kwargs["out_dtype"], + ) + + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8), + ) + + a = torch.empty((6, 128), dtype=torch.float8_e4m3fn) + weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn) + a_scale = torch.empty((6, 1), dtype=torch.float32) + weight_scale = torch.empty((2, 1), dtype=torch.float32) + kernel = object.__new__(B12xFp8BlockScaledMMKernel) + kernel.config = types.SimpleNamespace(out_dtype=torch.bfloat16) + + output = kernel.apply_block_scaled_mm(a, weight, a_scale, weight_scale) + + assert output.shape == (6, 256) + assert output.dtype == torch.bfloat16 + assert len(calls) == 1 + assert calls[0] == ( + (a, a_scale, weight, weight_scale), + {"out_dtype": torch.bfloat16}, + ) + torch.testing.assert_close(output, torch.full_like(output, 13.0)) + + +def test_b12x_mxfp4_requires_dynamic_activations() -> None: + config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic) + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) + + assert can_implement + assert reason is None + + config.activation_quant_key = None + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) + + assert not can_implement + assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations" + + +@pytest.mark.parametrize( + ("kernel_cls", "module_name", "scale_dtype"), + [ + ( + B12xMxFp4LinearKernel, + "vllm.model_executor.kernels.linear.mxfp4.b12x", + torch.uint8, + ), + ( + B12xNvFp4LinearKernel, + "vllm.model_executor.kernels.linear.nvfp4.b12x", + torch.float8_e4m3fn, + ), + ], +) +def test_b12x_fp4_processes_scale_and_preserves_loader( + monkeypatch, + kernel_cls, + module_name: str, + scale_dtype: torch.dtype, +) -> None: + scale = torch.empty((48, 8), dtype=scale_dtype) + swizzled_scale = torch.empty((128, 8), dtype=scale_dtype) + intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale) + monkeypatch.setattr( + importlib.import_module(module_name), + "_import_b12x_intrinsics", + lambda: intrinsics, + ) + layer = torch.nn.Module() + layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" + layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + weight_loader = object() + layer.weight_scale.weight_loader = weight_loader + kernel = object.__new__(kernel_cls) + + kernel.process_weights_after_loading(layer) + + assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() + assert layer.weight_scale.weight_loader is weight_loader + assert layer.b12x_warmup_provider is kernel + + +def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod + import vllm.utils.flashinfer as flashinfer_utils + + calls: list[tuple] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 4), dtype=torch.uint8) + + def mm_mxfp4(*args, **kwargs): + calls.append((args, kwargs)) + return torch.full((6, 48), 3.0, dtype=torch.bfloat16) + + monkeypatch.setattr( + flashinfer_utils, + "flashinfer_mxfp4_quantize", + lambda *args, **kwargs: (x_packed, x_scale_storage), + ) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4), + ) + + layer = torch.nn.Module() + layer.output_size_per_partition = 48 + layer.weight = torch.empty((48, 64), dtype=torch.uint8) + layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8) + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + bias = torch.ones(48, dtype=torch.bfloat16) + kernel = object.__new__(B12xMxFp4LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + torch.testing.assert_close(output, torch.full_like(output, 4.0)) + assert len(calls) == 1 + args, kwargs = calls[0] + assert args == ( + x_packed, + x_scale_storage, + layer.weight, + layer.weight_scale, + ) + assert kwargs == {"out_dtype": torch.bfloat16} + + +def test_b12x_nvfp4_can_implement_supported_config() -> None: + can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) + + assert can_implement + assert reason is None + + +def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: + import vllm.model_executor.kernels.linear as linear_mod + + monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) + monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + MarlinNvFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_nvfp4_linear_kernel(use_a16=True) + + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + calls: list[tuple] = [] + quant_calls: list[tuple] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + + def quant(*args, **kwargs): + quant_calls.append((args, kwargs)) + return x_packed, x_scale_storage + + def mm_nvfp4(*args, **kwargs): + calls.append((args, kwargs)) + return torch.full((6, 48), 3.0, dtype=torch.bfloat16) + + monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", quant) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4), + ) + + layer = torch.nn.Module() + layer.output_size_per_partition = 48 + layer.weight = torch.empty((48, 64), dtype=torch.uint8) + layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + layer.input_global_scale_inv = torch.tensor(2.0) + layer.alpha = torch.tensor(0.25) + x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2] + bias = torch.ones(48, dtype=torch.bfloat16) + kernel = object.__new__(B12xNvFp4LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + torch.testing.assert_close(output, torch.full_like(output, 4.0)) + assert len(quant_calls) == 1 + quant_args, quant_kwargs = quant_calls[0] + assert quant_args[0].shape == (6, 128) + assert quant_args[0].data_ptr() == x.data_ptr() + assert quant_args[1] is layer.input_global_scale_inv + assert not quant_args[0].is_contiguous() + assert quant_kwargs == {"is_sf_swizzled_layout": True} + assert len(calls) == 1 + args, kwargs = calls[0] + assert args == ( + x_packed, + x_scale_storage, + layer.weight, + layer.weight_scale, + layer.alpha, + ) + assert kwargs == {"out_dtype": torch.bfloat16} diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py new file mode 100644 index 000000000000..26069e0a9ad1 --- /dev/null +++ b/tests/model_executor/test_b12x_warmup.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + B12xFp8BlockScaledMMKernel, + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, + B12xTensorFP8ScaledMMLinearKernel, +) +from vllm.model_executor.warmup.b12x_warmup import b12x_warmup +from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts + + +def test_b12x_warmup_token_counts_cover_serving_regimes() -> None: + assert b12x_warmup_token_counts( + max_tokens=2048, + cudagraph_capture_sizes=[1, 2, 8, 128], + ) == (1, 2, 8, 128, 2048) + + +@pytest.mark.parametrize( + ("kernel_cls", "module_name", "call_name", "layer", "name"), + [ + ( + B12xMxFp4LinearKernel, + "vllm.model_executor.kernels.linear.mxfp4.b12x", + "_apply_b12x_mxfp4_linear", + SimpleNamespace( + weight=torch.empty((48, 64), dtype=torch.uint8), + weight_scale=torch.empty((128, 4), dtype=torch.uint8), + ), + "MXFP4", + ), + ( + B12xNvFp4LinearKernel, + "vllm.model_executor.kernels.linear.nvfp4.b12x", + "_apply_b12x_nvfp4_linear", + SimpleNamespace( + weight=torch.empty((48, 64), dtype=torch.uint8), + weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn), + input_global_scale_inv=torch.tensor(2.0), + alpha=torch.tensor(0.25), + ), + "NVFP4", + ), + ( + B12xFp8BlockScaledMMKernel, + "vllm.model_executor.kernels.linear.scaled_mm.b12x", + "_run_b12x_fp8_block_scaled_mm", + SimpleNamespace( + weight=torch.empty((256, 128), dtype=torch.float8_e4m3fn), + weight_scale_inv=torch.empty((2, 1), dtype=torch.float32), + ), + "block-FP8", + ), + ], +) +def test_b12x_warmup_units_cover_token_counts( + monkeypatch, + kernel_cls, + module_name: str, + call_name: str, + layer, + name: str, +) -> None: + calls = [] + monkeypatch.setattr( + importlib.import_module(module_name), + call_name, + lambda *args: calls.append(args), + ) + kernel = object.__new__(kernel_cls) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16) + unit.compile() + + assert unit.name == name + assert [args[0].shape[0] for args in calls] == [1, 8] + assert unit.key[-1] == torch.bfloat16 + + +def test_b12x_mxfp8_warmup_unit(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + calls = [] + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: SimpleNamespace( + mm=lambda *args, **kwargs: calls.append((args, kwargs)) + ), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: SimpleNamespace(cuda_stream=object()), + ) + packed_weight = SimpleNamespace( + in_features=128, + padded_in_features=128, + out_features=256, + weight=SimpleNamespace(values=torch.empty(1)), + ) + layer = SimpleNamespace(b12x_mxfp8_packed_weight=packed_weight) + kernel = object.__new__(B12xMxfp8LinearKernel) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.float16) + unit.compile() + + assert [args[0].shape for args, _ in calls] == [(1, 128), (8, 128)] + assert [kwargs["expected_m"] for _, kwargs in calls] == [1, 8] + + +def test_b12x_tensor_fp8_warmup_unit(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + monkeypatch.setattr( + b12x_mod, + "_import_b12x_tensor_fp8", + lambda: SimpleNamespace( + prewarm=lambda *args, **kwargs: calls.append((args, kwargs)) + ), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: SimpleNamespace(cuda_stream=object()), + ) + packed_weight = SimpleNamespace( + in_features=128, + padded_in_features=128, + out_features=256, + values=torch.empty(1), + ) + layer = SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_weight) + kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16) + unit.compile() + + assert calls[0][0] == (packed_weight, (1, 8)) + assert calls[0][1]["out_dtype"] == torch.bfloat16 + + +def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None: + import vllm.model_executor.warmup.b12x_warmup as warmup_mod + + calls: list[tuple[str, tuple[int, ...], torch.dtype]] = [] + + class Provider: + def get_b12x_warmup_unit(self, layer, token_counts, output_dtype): + return B12xWarmupUnit( + name="fake", + key=(type(self), layer.shape, output_dtype), + compile=lambda: calls.append((layer.name, token_counts, output_dtype)), + ) + + provider = Provider() + layers = [ + SimpleNamespace(name="first", shape=(128, 256), b12x_warmup_provider=provider), + SimpleNamespace( + name="duplicate", shape=(128, 256), b12x_warmup_provider=provider + ), + SimpleNamespace(name="second", shape=(256, 256), b12x_warmup_provider=provider), + SimpleNamespace(), + ] + scans = 0 + + def modules(): + nonlocal scans + scans += 1 + return iter(layers) + + worker = SimpleNamespace( + get_model=lambda: SimpleNamespace(modules=modules), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace(dtype=torch.float32), + ) + platform = SimpleNamespace( + is_cuda=lambda: True, + is_device_capability_family=lambda family: family == 120, + ) + synchronized = [] + monkeypatch.setattr(warmup_mod, "current_platform", platform) + monkeypatch.setattr( + warmup_mod.torch.accelerator, + "synchronize", + lambda: synchronized.append(True), + ) + + b12x_warmup(worker, [1, 2]) + + assert scans == 1 + assert calls == [ + ("first", (1, 2, 8), torch.bfloat16), + ("second", (1, 2, 8), torch.bfloat16), + ] + assert synchronized == [True] diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py index f93de50a3d30..b5c6f90ea491 100644 --- a/tests/model_executor/test_jit_warmup.py +++ b/tests/model_executor/test_jit_warmup.py @@ -9,11 +9,15 @@ import pytest from vllm.model_executor.warmup.jit_warmup import ( + JitWarmupRegistry, VllmJitKernel, WarmupIntRange, get_ast_full_name, zip_inputs, ) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + triton_scalar_specialization_rep, +) def _next_power_of_2(value: int) -> int: @@ -94,6 +98,35 @@ def compile(self, compile_key: ToyKernel.CompileKey) -> None: self.compiled.append(compile_key) +@pytest.mark.parametrize( + ("value", "expected"), + [ + (-(1 << 63), 1 << 31), + (-(1 << 31) - 1, (1 << 31) + 1), + (-(1 << 31), 16), + (0, 16), + (1, 1), + (2, 2), + (16, 16), + ((1 << 31) - 1, 2), + (1 << 31, 1 << 31), + ((1 << 31) + 1, (1 << 31) + 1), + ((1 << 63) - 1, (1 << 31) + 1), + (1 << 63, 1 << 63), + ((1 << 63) + 1, (1 << 63) + 1), + ((1 << 64) - 1, (1 << 63) + 1), + ], +) +def test_triton_scalar_specialization_rep(value: int, expected: int) -> None: + assert triton_scalar_specialization_rep(value) == expected + + +@pytest.mark.parametrize("value", [-(1 << 63) - 1, 1 << 64]) +def test_triton_scalar_specialization_rep_rejects_out_of_range(value: int) -> None: + with pytest.raises(OverflowError, match="outside Triton's scalar range"): + triton_scalar_specialization_rep(value) + + def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None: cfg = _config() @@ -287,6 +320,31 @@ def compile(self, compile_key: CompileKey) -> None: StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4}) +def test_dispatch_helper_calls_resolve_python_builtins() -> None: + class BuiltinKernel(VllmJitKernel["BuiltinKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + limit: int, + ) -> CompileKey: + return self.CompileKey(value=max(1, min(tokens, limit))) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + assert BuiltinKernel().compile_key({"tokens": 8, "limit": 4}) == ( + BuiltinKernel.CompileKey(value=4) + ) + + def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None: class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]): @dataclass(frozen=True) @@ -320,10 +378,85 @@ def compile(self, compile_key: CompileKey) -> None: with pytest.raises(ValueError, match="local assignments"): BranchKernel() - with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"): + with pytest.raises( + ValueError, + match=r"may unpack only its own \*\*kwargs parameter once", + ): KwargsReturnKernel() +def test_dispatch_can_forward_compile_key_fields() -> None: + class ForwardingKernel(VllmJitKernel["ForwardingKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + mode: int + block_size: int + + def dispatch( # type: ignore[override] + self, + *, + tokens: int, + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + block_size=_round_up(tokens, multiple=4), + ) + + def get_warmup_keys(self) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)( + tokens=(1, 5), + mode=(2, 3), + ) + + def compile(self, compile_key: CompileKey) -> None: + pass + + kernel = ForwardingKernel() + expected = kernel.CompileKey( + mode=2, + block_size=8, + ) + assert kernel.dispatch(tokens=5, mode=2) == expected + assert kernel.compile_key({"tokens": 5, "mode": 2}) == expected + assert kernel.get_warmup_keys() == [ + kernel.CompileKey(mode=2, block_size=4), + kernel.CompileKey(mode=3, block_size=4), + kernel.CompileKey(mode=2, block_size=8), + kernel.CompileKey(mode=3, block_size=8), + ] + with pytest.raises(TypeError, match="field 'block_size' is specified twice"): + kernel.compile_key({"tokens": 5, "mode": 2, "block_size": 4}) + with pytest.raises(TypeError, match="unexpected keyword argument 'extra'"): + kernel.compile_key({"tokens": 5, "mode": 2, "extra": 1}) + + +def test_dispatch_supports_tuple_and_mapping_subscriptions() -> None: + class SubscriptKernel(VllmJitKernel["SubscriptKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + first: int + named: int + + def dispatch( # type: ignore[override] + self, + *, + values: tuple[int, ...], + config: dict[str, int], + ) -> CompileKey: + return self.CompileKey(first=values[0], named=config["named"]) + + def get_warmup_keys(self) -> list[CompileKey]: + return [] + + def compile(self, compile_key: CompileKey) -> None: + pass + + assert SubscriptKernel().compile_key( + {"values": (3, 5), "config": {"named": 7}} + ) == SubscriptKernel.CompileKey(first=3, named=7) + + def test_dispatch_reports_unsupported_expression_with_context() -> None: class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]): @dataclass(frozen=True) @@ -361,6 +494,102 @@ def test_warmup_compiles_all_returned_keys_in_order() -> None: ] +def test_runtime_cache_miss_compiles_and_caches_executor() -> None: + class CachedKernel(VllmJitKernel["CachedKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def __init__(self) -> None: + self.compiled: list[CachedKernel.CompileKey] = [] + super().__init__() + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(value=value) + + def get_warmup_keys(self) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)(value=1) + + def compile(self, compile_key: CompileKey) -> None: + self.compiled.append(compile_key) + self._compiled_cache[compile_key] = object() + + def __call__(self, value: int) -> None: + compile_key = self.dispatch(value=value) + self._get_or_compile(compile_key) + + kernel = CachedKernel() + kernel.warmup() + kernel(1) + kernel(2) + + assert kernel.compiled == [ + CachedKernel.CompileKey(value=1), + CachedKernel.CompileKey(value=2), + ] + + +def test_registry_records_only_inside_model_setup_context() -> None: + registry = JitWarmupRegistry(_config()) + kernel = RecordingToyKernel() + + kernel.register_warmup(3, _config()) + with registry.activate(): + kernel.register_warmup(3, _config()) + + assert len(registry) == 1 + assert kernel.compiled == [] + + +def test_registry_expands_requests_and_deduplicates_owner_keys() -> None: + registry = JitWarmupRegistry(_config()) + kernel = RecordingToyKernel() + + with registry.activate(): + kernel.register_warmup(3, _config()) + kernel.register_warmup(5, _config()) + + registry.warmup() + + assert kernel.compiled == [ + ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True), + ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True), + ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True), + ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True), + ] + + +def test_registry_passes_vllm_config_to_default_requests() -> None: + class ConfigKernel(VllmJitKernel["ConfigKernel.CompileKey"]): + @dataclass(frozen=True) + class CompileKey: + value: int + + def __init__(self) -> None: + self.compiled: list[ConfigKernel.CompileKey] = [] + super().__init__() + + def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override] + return self.CompileKey(value=value) + + def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: + return [self.dispatch(value=vllm_config.bias)] + + def compile(self, compile_key: CompileKey) -> None: + self.compiled.append(compile_key) + + registry = JitWarmupRegistry(_config(bias=7)) + kernel = ConfigKernel() + + with registry.activate(): + kernel.register_warmup() + kernel.register_warmup() + assert len(registry) == 1 + registry.warmup() + + assert kernel.compiled == [ConfigKernel.CompileKey(value=7)] + + def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None: dotted_expr = ast.parse("foo.bar.baz").body[0] call_expr = ast.parse("foo()").body[0] diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py new file mode 100644 index 000000000000..7227ed6fd9ea --- /dev/null +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which value the DFlash drafter is fed under AttnRes. + +`_capture_aux_hidden_stream` picks the weights it mixes against from one of +three places depending on where the tapped layer sits, and returns the plain +running prefix when the feature is off. The mixture itself is the kernel's +job and is covered by ``test_attn_res.py``; what is asserted here is the +selection, which is the part that can silently feed the drafter the wrong +tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia import model as k3_model + +END_LAYER = 4 + + +def _weights(tag: float) -> SimpleNamespace: + """A norm/projection pair that is identifiable by value.""" + return SimpleNamespace( + weight=torch.full((2,), tag), + variance_epsilon=tag, + ) + + +def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: + """A stand-in carrying only what the tap reads. + + Constructing the real model needs a distributed init and weights, and none + of it participates in the selection under test. + """ + consumers = [] + for i in range(END_LAYER): + consumers.append( + SimpleNamespace( + self_attention_res_norm=_weights(float(i)), + self_attention_res_proj=SimpleNamespace( + weight=torch.full((1, 2), float(i)) + ), + prev_valid_blocks=i, + ) + ) + return SimpleNamespace( + _aux_attn_res_stream=enabled, + use_attn_res=use_attn_res, + end_layer=END_LAYER, + layers=consumers, + output_attn_res_norm=_weights(99.0), + output_attn_res_proj=SimpleNamespace(weight=torch.full((1, 2), 99.0)), + num_attn_res_blocks=99, + ) + + +@pytest.fixture +def recorder(monkeypatch): + """Replace the kernel so the call it would have made is inspectable.""" + calls = [] + + def _fake_attn_res( + prefix, + delta, + block_residual, + norm_weight, + proj_weight, + output_norm_weight, + **kwargs, + ): + calls.append( + SimpleNamespace( + prefix=prefix, + delta=delta, + block_residual=block_residual, + norm_weight=norm_weight, + proj_weight=proj_weight, + kwargs=kwargs, + ) + ) + return torch.full_like(prefix, -1.0) + + monkeypatch.setattr(k3_model, "attn_res", _fake_attn_res) + return calls + + +def _set_last_rank(monkeypatch, is_last: bool): + monkeypatch.setattr( + k3_model, + "get_pp_group", + lambda: SimpleNamespace(is_last_rank=is_last), + ) + + +def _call(stub, layer_idx, prefix_sum, pending_mlp_out, block_residual): + return k3_model.KimiLinearModel._capture_aux_hidden_stream( + stub, layer_idx, prefix_sum, pending_mlp_out, block_residual + ) + + +@pytest.mark.parametrize( + "enabled,use_attn_res", [(False, True), (True, False), (False, False)] +) +def test_disabled_reproduces_the_plain_residual_sum( + recorder, monkeypatch, enabled, use_attn_res +): + """Off, the tap must be exactly the sum it replaced. + + Both conditions matter. `use_attn_res` is what constructs the norm and + projection weights, so without it the lookups below would raise rather + than fall back. + """ + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + got = _call( + _stub_model(enabled=enabled, use_attn_res=use_attn_res), + 0, + prefix_sum, + pending, + torch.zeros(2), + ) + + torch.testing.assert_close(got, prefix_sum + pending) + assert not recorder, "the kernel must not run when the tap is off" + + +def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): + """The value the next layer reads is the mixture against *its* weights, + so the tap has to reach forward rather than use the current layer's.""" + _set_last_rank(monkeypatch, True) + + _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) + + assert len(recorder) == 1 + call = recorder[0] + # Layer 2's weights, not layer 1's. + torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0)) + assert call.kwargs["num_blocks"] == 2 + + +def test_last_layer_on_the_final_rank_uses_the_output_aggregation( + recorder, monkeypatch +): + """Nothing downstream but the model's own output-side mixture.""" + _set_last_rank(monkeypatch, True) + + _call( + _stub_model(enabled=True), END_LAYER - 1, torch.zeros(2), None, torch.zeros(2) + ) + + assert len(recorder) == 1 + torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0)) + assert recorder[0].kwargs["num_blocks"] == 99 + + +def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch): + """The consumer lives on the next rank and the output aggregation only + exists on the last one, so there is nothing here to mix against. + + This is the case that would otherwise reach for weights this rank never + constructs. The forward guard is `layer_idx + 1 < end_layer`, where + `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a + `PPMissingLayer` is unreachable by construction -- the fallback below is + what makes that true rather than merely likely. + """ + _set_last_rank(monkeypatch, False) + prefix_sum = torch.tensor([3.0, 4.0]) + + got = _call( + _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2) + ) + + torch.testing.assert_close(got, prefix_sum) + assert not recorder, "no weights exist on this rank to mix against" + + +def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( + recorder, monkeypatch +): + """The kernel writes an applied delta back into the prefix in place, which + would double-add it into the live residual stream, so the pending output + has to arrive already summed into the prefix with `delta` left None.""" + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2)) + + assert len(recorder) == 1 + assert recorder[0].delta is None + torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending) + # And the caller's tensor is not mutated on the way. + torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0])) diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 61a24a83e041..652f6652b353 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -10,6 +10,7 @@ from vllm.models.kimi_k3.nvidia import model as kimi_model from vllm.models.kimi_k3.nvidia.model import ( KimiK3ForConditionalGeneration, + KimiLinearForCausalLM, KimiLinearModel, ) @@ -18,6 +19,7 @@ def _make_kimi_linear_model() -> KimiLinearModel: model = object.__new__(KimiLinearModel) object.__setattr__(model, "aux_hidden_state_layers", (2,)) object.__setattr__(model, "use_sequence_parallel", False) + object.__setattr__(model, "use_attn_res", False) return model @@ -25,6 +27,14 @@ def test_kimi_k3_advertises_eagle3_support(): assert supports_eagle3(KimiK3ForConditionalGeneration) +def test_kimi_linear_advertises_eagle3_support(): + # The text-only architecture serves the same inner KimiLinearModel, which + # already carries the EagleModelMixin tap machinery - only the interface + # declaration was missing, so EAGLE3-family speculative decoding (dspark) + # was rejected at startup with "Model does not support EAGLE3 interface". + assert supports_eagle3(KimiLinearForCausalLM) + + def test_kimi_k3_uses_shared_eagle3_layer_configuration(): target = object.__new__(KimiK3ForConditionalGeneration) torch.nn.Module.__init__(target) @@ -128,3 +138,64 @@ def test_kimi_linear_forward_extracts_attn_res_aux_hidden_states(monkeypatch): torch.testing.assert_close(aux_hidden_states[0], initial_hidden_states) torch.testing.assert_close(aux_hidden_states[1], prefix_sum + layer_hidden_states) assert final_attn_res.call_args.args[2] is block_residual + + +def test_attn_res_stream_capture_receives_the_layer_outputs_in_order(monkeypatch): + """Pin the argument mapping at the call site. + + The capture helper's own tests invoke it directly by keyword, so they + cannot catch a swap where `forward` hands it the residual as the pending + MLP output. Both are tensors of the same shape, so a swap is silent: it + feeds the drafter a wrong but well-formed tensor. + """ + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + captured = torch.tensor([[11.0, 12.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (1,)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr(kimi_model, "attn_res", Mock(return_value=torch.zeros(1, 2))) + monkeypatch.setenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "1") + + capture = Mock(return_value=captured) + monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture) + + _, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args + assert layer_idx == 0 + assert got_prefix is prefix_sum + assert got_pending is layer_hidden_states + assert got_residual is block_residual + torch.testing.assert_close(aux_hidden_states[0], captured) diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 9fae1c2bcd3d..7416cca02a32 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -6,6 +6,8 @@ Uses torch.rand for q/k/v to match FLA's test pattern. """ +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F @@ -22,6 +24,12 @@ is_flashkda_supported, is_fused_kda_decode_supported, ) +from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM +from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops +from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + kda_recoverssm_verify, +) from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( chunk_kda, chunk_kda_with_fused_gate, @@ -30,9 +38,15 @@ fused_recurrent_kda_fwd, fused_recurrent_kda_packed_decode, ) +from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd -DEVICE = "cuda" +DEVICE = current_platform.device_type + +pytestmark = pytest.mark.skipif( + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), + reason="The KDA kernels require a CUDA-alike or XPU device.", +) # The AMD and NVIDIA copies of the KDA kernels are vendored separately and are # free to diverge, so the shared-semantics tests below run against both. @@ -42,6 +56,38 @@ } +def test_kda_recoverssm_config_state_layout(): + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + dtype=torch.bfloat16, + hf_config=SimpleNamespace( + linear_attn_config={ + "num_heads": 4, + "head_dim": 32, + "short_conv_kernel_size": 4, + } + ), + ), + cache_config=SimpleNamespace( + mamba_cache_dtype="auto", + use_kda_recoverssm=True, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=1), + speculative_config=SimpleNamespace(num_speculative_tokens=2), + ) + + assert KimiLinearForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == ( + torch.bfloat16, + torch.float32, + torch.float32, + torch.bfloat16, + ) + assert KimiLinearForCausalLM.get_mamba_state_shape_from_config(vllm_config)[2:] == ( + (4, 3, 32), + (4, 3, 64), + ) + + @torch.inference_mode() def test_gather_initial_states_correctness(): row_size = 8 * 128 * 128 @@ -529,6 +575,308 @@ def test_kda_spec_decode_correctness( assert torch.isnan(output_storage[..., H * D :]).all() +@pytest.mark.parametrize( + ( + "conv_state_dim_first", + "use_request_indices", + "lower_bound", + "align_mode", + ), + [ + pytest.param(False, False, None, False, id="baseline"), + pytest.param(True, True, -5.0, True, id="all-features"), + pytest.param(False, True, -5.0, False, id="request-indexed"), + pytest.param(True, False, None, True, id="aligned"), + ], +) +@torch.inference_mode() +def test_kda_recoverssm_verify_and_group_commit( + monkeypatch: pytest.MonkeyPatch, + lower_bound: float | None, + use_request_indices: bool, + conv_state_dim_first: bool, + align_mode: bool, +): + monkeypatch.setattr( + recoverssm_ops, + "is_conv_state_dim_first", + lambda: conv_state_dim_first, + ) + num_layers, num_seqs, query_len = 2, 2, 8 + num_blocks, num_heads, dim = (7 if align_mode else 3), 4, 128 + total_tokens = num_seqs * query_len + torch.manual_seed(20260808) + + q, k, v, raw_g = [ + torch.randn( + 1, + total_tokens, + num_heads, + dim, + dtype=torch.bfloat16, + device=DEVICE, + ) + for _ in range(4) + ] + raw_beta = torch.randn( + 1, + total_tokens, + num_heads, + dtype=torch.bfloat16, + device=DEVICE, + ) + query_start_loc = torch.arange( + 0, + total_tokens + 1, + query_len, + dtype=torch.int32, + device=DEVICE, + ) + state_indices = torch.tensor( + [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE + ) + accepted = [2, 8] + if use_request_indices: + global_num_accepted = torch.tensor( + [0, accepted[0], 0, accepted[1]], + dtype=torch.int32, + device=DEVICE, + ) + request_indices = torch.tensor([1, 3], dtype=torch.int32, device=DEVICE) + else: + global_num_accepted = torch.tensor(accepted, dtype=torch.int32, device=DEVICE) + request_indices = None + + block_table = None + num_computed_tokens = None + mamba_block_size = None + if align_mode: + batch_size = 4 if use_request_indices else num_seqs + block_table = torch.full((batch_size, 2), -1, dtype=torch.int32, device=DEVICE) + rows = ( + request_indices + if request_indices is not None + else torch.arange(num_seqs, device=DEVICE) + ) + block_table[rows] = torch.tensor( + [[1, 5], [2, 6]], + dtype=torch.int32, + device=DEVICE, + ) + num_computed_tokens = torch.zeros(batch_size, dtype=torch.int32, device=DEVICE) + num_computed_tokens[rows] = 4 + mamba_block_size = 8 + + layers = [] + expected_outputs = [] + expected_states = [] + initial_states = [] + initial_conv_states = [] + history_len, conv_dim = 3, 12 + for layer_idx in range(num_layers): + A_log = ( + 0.2 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE) + + layer_idx * 0.03 + ).contiguous() + dt_bias = ( + 0.1 * torch.randn(num_heads, dim, dtype=torch.float32, device=DEVICE) + ).contiguous() + checkpoint = 0.01 * torch.randn( + num_blocks, + num_heads, + dim, + dim, + dtype=torch.float32, + device=DEVICE, + ) + conv_shape = ( + (num_blocks, conv_dim, history_len + query_len - 1) + if conv_state_dim_first + else (num_blocks, history_len + query_len - 1, conv_dim) + ) + conv_state = torch.randn(conv_shape, dtype=torch.bfloat16, device=DEVICE) + correction_cache = torch.empty( + num_blocks, + num_heads, + query_len, + dim, + dtype=torch.float32, + device=DEVICE, + ) + kg_cache = torch.empty( + num_blocks, + num_heads, + query_len, + 2 * dim, + dtype=torch.bfloat16, + device=DEVICE, + ) + layer = SimpleNamespace( + kv_cache=( + conv_state, + checkpoint, + correction_cache, + kg_cache, + ), + A_log=A_log, + dt_bias=dt_bias, + local_num_heads=num_heads, + head_dim=dim, + gate_lower_bound=lower_bound, + ) + layers.append(layer) + initial_states.append(checkpoint.clone()) + initial_conv_states.append(conv_state.clone()) + + actual_output = kda_recoverssm_verify( + q=q, + k=k, + v=v, + raw_g=raw_g, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + checkpoint_state=checkpoint, + correction_cache=correction_cache, + kg_cache=kg_cache, + query_start_loc=query_start_loc, + state_indices=state_indices, + spec_query_len=query_len, + ) + + normalized_q = q.float() * torch.rsqrt( + q.float().square().sum(dim=-1, keepdim=True) + 1e-6 + ) + normalized_k = k.float() * torch.rsqrt( + k.float().square().sum(dim=-1, keepdim=True) + 1e-6 + ) + gate_input = raw_g.float() + dt_bias.view(1, 1, num_heads, dim) + if lower_bound is None: + gate = -A_log.exp().view(1, 1, num_heads, 1) * F.softplus(gate_input) + else: + gate = lower_bound * torch.sigmoid( + A_log.exp().view(1, 1, num_heads, 1) * gate_input + ) + beta = raw_beta.float().sigmoid() + + reference_output = [] + committed_states = checkpoint.clone() + for seq_idx, commit_len in enumerate(accepted): + start = seq_idx * query_len + end = start + query_len + output, _ = naive_recurrent_kda( + normalized_q[:, start:end], + normalized_k[:, start:end], + v[:, start:end], + gate[:, start:end], + beta[:, start:end], + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), + ) + reference_output.append(output) + _, committed_state = naive_recurrent_kda( + normalized_q[:, start : start + commit_len], + normalized_k[:, start : start + commit_len], + v[:, start : start + commit_len], + gate[:, start : start + commit_len], + beta[:, start : start + commit_len], + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), + output_final_state=True, + ) + assert committed_state is not None + final_block = state_indices[seq_idx] + if align_mode: + assert block_table is not None + row = request_indices[seq_idx] if use_request_indices else seq_idx + final_block = block_table[row, (4 + commit_len) // 8] + committed_states[final_block] = committed_state.transpose(-1, -2) + if align_mode and 4 + commit_len >= 8: + _, boundary_state = naive_recurrent_kda( + normalized_q[:, start : start + 4], + normalized_k[:, start : start + 4], + v[:, start : start + 4], + gate[:, start : start + 4], + beta[:, start : start + 4], + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2), + output_final_state=True, + ) + assert boundary_state is not None + assert block_table is not None + row = request_indices[seq_idx] if use_request_indices else seq_idx + committed_states[block_table[row, 0]] = boundary_state.transpose(-1, -2) + expected_outputs.append(torch.cat(reference_output, dim=1)) + expected_states.append(committed_states) + torch.testing.assert_close(checkpoint, initial_states[-1]) + torch.testing.assert_close( + actual_output, + expected_outputs[-1], + atol=3e-2, + rtol=3e-2, + ) + + context = KDARecoverSSMCommitContext.create( + layers, + spec_query_len=query_len, + max_num_reqs=global_num_accepted.shape[0], + ) + context.commit( + global_num_accepted, + state_indices, + query_start_loc, + request_indices=request_indices, + block_table=block_table, + num_computed_tokens=num_computed_tokens, + mamba_block_size=mamba_block_size, + ) + + for layer_idx, layer in enumerate(layers): + torch.testing.assert_close( + layer.kv_cache[1], + expected_states[layer_idx], + atol=3e-3, + rtol=3e-3, + ) + for seq_idx, commit_len in enumerate(accepted): + block = state_indices[seq_idx] + if align_mode: + assert block_table is not None + row = request_indices[seq_idx] if use_request_indices else seq_idx + block = block_table[row, (4 + commit_len) // 8] + source_block = state_indices[seq_idx] if align_mode else block + if conv_state_dim_first: + actual_conv = layer.kv_cache[0][block, :, :history_len] + expected_conv = initial_conv_states[layer_idx][ + source_block, + :, + commit_len - 1 : commit_len - 1 + history_len, + ] + else: + actual_conv = layer.kv_cache[0][block, :history_len] + expected_conv = initial_conv_states[layer_idx][ + source_block, + commit_len - 1 : commit_len - 1 + history_len, + ] + torch.testing.assert_close(actual_conv, expected_conv) + if align_mode and 4 + commit_len >= 8: + assert block_table is not None + boundary_block = block_table[row, 0] + if conv_state_dim_first: + actual_boundary_conv = layer.kv_cache[0][ + boundary_block, :, :history_len + ] + expected_boundary_conv = initial_conv_states[layer_idx][ + state_indices[seq_idx], :, 3 : 3 + history_len + ] + else: + actual_boundary_conv = layer.kv_cache[0][ + boundary_block, :history_len + ] + expected_boundary_conv = initial_conv_states[layer_idx][ + state_indices[seq_idx], 3 : 3 + history_len + ] + torch.testing.assert_close(actual_boundary_conv, expected_boundary_conv) + + @pytest.mark.parametrize( ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"), [ diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 5352ef7a7a64..069c9a4d7a8e 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import fields +from unittest.mock import Mock import pytest import torch @@ -26,6 +27,9 @@ GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMPostprocessMetadata, +) from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, mamba_get_block_table_tensor, @@ -44,8 +48,12 @@ } -def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata): - for field in fields(KimiK3KDAMetadata): +def _assert_matches_shared_gdn( + reference: GDNAttentionMetadata, actual: KimiK3KDAMetadata +): + assert actual.recoverssm_commit is None + assert actual.recoverssm_context is None + for field in fields(GDNAttentionMetadata): actual_value = getattr(actual, field.name) expected_value = getattr(reference, field.name) if field.name in PRUNED_METADATA_FIELDS: @@ -78,6 +86,7 @@ def _make_builder( full_cuda_graph: bool, device: torch.device = DEVICE, mamba_cache_mode: str = "none", + use_recoverssm: bool = False, ) -> AttentionMetadataBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -92,17 +101,24 @@ def _make_builder( CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode - return builder_cls( + vllm_config.cache_config.use_replayssm = use_recoverssm + vllm_config.cache_config.use_kda_recoverssm = use_recoverssm + builder = builder_cls( kv_cache_spec=MambaSpec( block_size=BLOCK_SIZE, shapes=((16, 64),), dtypes=(torch.float16,), - num_speculative_blocks=num_speculative_tokens, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens), ), layer_names=["layer.0"], vllm_config=vllm_config, device=device, ) + if use_recoverssm: + assert isinstance(builder, KimiK3KDAMetadataBuilder) + builder.recoverssm_context = Mock() + return builder @pytest.mark.parametrize( @@ -244,6 +260,99 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding(): torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3])) +@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) +def test_recoverssm_spec_uses_one_state_slot_and_current_window( + mamba_cache_mode: str, +): + if mamba_cache_mode == "align" and not torch.cuda.is_available(): + pytest.skip("align metadata construction requires CUDA") + device = torch.device("cuda") if mamba_cache_mode == "align" else DEVICE + batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device + ).replace(is_prefilling=torch.tensor([True, True, False])) + builder = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + device=device, + mamba_cache_mode=mamba_cache_mode, + use_recoverssm=True, + ) + assert isinstance(builder, KimiK3KDAMetadataBuilder) + context = builder.recoverssm_context + assert context is not None + actual = builder.build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32), + num_accepted_tokens=torch.tensor([3, 2, 2], dtype=torch.int32, device=device), + ) + + assert actual.spec_state_indices_tensor is not None + assert actual.spec_state_indices_tensor.shape == (1, 1) + torch.testing.assert_close( + actual.num_accepted_tokens, + torch.ones(1, dtype=torch.int32, device=device), + ) + commit_metadata = actual.recoverssm_commit + assert commit_metadata is not None + torch.testing.assert_close( + commit_metadata.request_indices, + torch.tensor([2], dtype=torch.int32, device=device), + ) + assert actual.recoverssm_context is context + num_accepted_tokens = torch.tensor([3, 2, 1], dtype=torch.int32, device=device) + + postprocess = actual.commit_recoverssm_state(num_accepted_tokens) + + if mamba_cache_mode == "none": + assert commit_metadata.align is None + assert postprocess is None + else: + assert isinstance(postprocess, RecoverSSMPostprocessMetadata) + assert postprocess.num_spec_decodes == 1 + assert postprocess.request_indices is commit_metadata.request_indices + assert postprocess.block_table is common_attn_metadata.block_table_tensor + assert ( + postprocess.num_computed_tokens + is common_attn_metadata.compute_num_computed_tokens() + ) + assert postprocess.block_size == BLOCK_SIZE + args = context.commit.call_args.args + assert args[0] is num_accepted_tokens + torch.testing.assert_close(args[1], commit_metadata.state_indices[:, 0]) + torch.testing.assert_close(args[2], commit_metadata.query_start_loc) + + +def test_recoverssm_distinguishes_draftless_decode_from_one_token_prefill(): + batch = BatchSpec(seq_lens=[40, 30], query_lens=[1, 1]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE + ).replace(is_prefilling=torch.tensor([False, True])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=False, + use_recoverssm=True, + ).build( + 0, + common_attn_metadata, + num_decode_draft_tokens_cpu=torch.full((2,), -1, dtype=torch.int32), + num_accepted_tokens=torch.ones(2, dtype=torch.int32), + ) + + assert actual.num_spec_decodes == 1 + assert actual.num_decodes == 0 + assert actual.num_prefills == 1 + assert actual.spec_state_indices_tensor is not None + assert actual.spec_state_indices_tensor.shape == (1, 1) + torch.testing.assert_close( + actual.spec_query_start_loc, + torch.tensor([0, 1], dtype=torch.int32), + ) + + @pytest.mark.parametrize( ("seq_len", "expected_has_initial_state"), [ @@ -307,6 +416,38 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn(): _assert_matches_shared_gdn(reference, actual) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_recoverssm_spec_cudagraph_stages_one_checkpoint_per_request(): + device = torch.device("cuda") + batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device + ).replace(is_prefilling=torch.tensor([False, False])) + builder = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=2, + full_cuda_graph=True, + device=device, + use_recoverssm=True, + ) + assert isinstance(builder, KimiK3KDAMetadataBuilder) + assert builder.spec_state_indices_tensor.shape == ( + builder.vllm_config.scheduler_config.max_num_seqs, + 1, + ) + actual = builder.build_for_cudagraph_capture(common_attn_metadata) + + assert actual.spec_state_indices_tensor is not None + assert actual.spec_state_indices_tensor.shape == (batch.batch_size, 1) + assert actual.num_accepted_tokens is not None + torch.testing.assert_close( + actual.num_accepted_tokens, + torch.ones(batch.batch_size, dtype=torch.int32, device=device), + ) + assert actual.recoverssm_commit is not None + assert actual.recoverssm_commit.request_indices is None + + def test_kimi_k3_kda_backend_uses_private_metadata_builder(): assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder assert KimiK3KDAAttentionBackend.is_ssm() diff --git a/tests/models/kimi_k3/test_sequence_parallel.py b/tests/models/kimi_k3/test_sequence_parallel.py index 97d4f7b74abe..9e8faea099cb 100644 --- a/tests/models/kimi_k3/test_sequence_parallel.py +++ b/tests/models/kimi_k3/test_sequence_parallel.py @@ -267,19 +267,21 @@ def test_kimi_mtp_restores_sequence_parallel_output(monkeypatch): @pytest.mark.parametrize( - ("enabled", "use_sequence_parallel", "tp_size", "expected"), + ("enabled", "use_sequence_parallel", "eligible", "tp_size", "expected"), [ - (True, True, 8, True), - (False, True, 8, False), # opt-in only - (True, False, 8, False), # replication only exists under SP - (True, True, 1, False), # nothing to shard - (True, True, 5, False), # 6144 % 5 -- would fail divide() + (True, True, True, 8, True), + (False, True, True, 8, False), # opt-in only + (True, False, True, 8, False), # replication only exists under SP + (True, True, False, 8, False), # FusedMoE path owns the reduction + (True, True, True, 1, False), # nothing to shard + (True, True, True, 5, False), # 6144 % 5 -- would fail divide() ], ) def test_shard_sequence_parallel_mlp_gating( monkeypatch, enabled: bool, use_sequence_parallel: bool, + eligible: bool, tp_size: int, expected: bool, ): @@ -293,6 +295,7 @@ def test_shard_sequence_parallel_mlp_gating( hidden_size=7168, intermediate_size=6144, use_sequence_parallel=use_sequence_parallel, + eligible=eligible, ) is expected ) diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index 3057e14060c6..bb7afb23365b 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -186,40 +186,39 @@ def _assert_embeddings_close(vllm_outputs, hf_embeddings): ) -@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module") +@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="class") def colbert_spec(request): """Return the model spec dict for the current parametrization.""" return COLBERT_MODELS[request.param] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_model_name(colbert_spec): return colbert_spec["model"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_dim(colbert_spec): return colbert_spec["colbert_dim"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_max_model_len(colbert_spec): return colbert_spec["max_model_len"] -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def colbert_extra_kwargs(colbert_spec): return colbert_spec["extra_kwargs"] -def test_colbert_token_embed( +@pytest.fixture(scope="class") +def colbert_model( vllm_runner, colbert_model_name, - colbert_dim, colbert_max_model_len, colbert_extra_kwargs, ): - """Test that ColBERT model produces token embeddings.""" with vllm_runner( colbert_model_name, runner="pooling", @@ -228,7 +227,19 @@ def test_colbert_token_embed( enforce_eager=True, **colbert_extra_kwargs, ) as vllm_model: - outputs = vllm_model.token_embed([TEXTS_1[0]]) + yield vllm_model + + +class TestColbertSharedEngine: + """Tests sharing one engine per model. + + Class-scoped so the engine is released before `test_colbert_hf_comparison`, + which needs a runner of its own and cannot start while this one holds VRAM. + """ + + def test_colbert_token_embed(self, colbert_model, colbert_dim): + """Test that ColBERT model produces token embeddings.""" + outputs = colbert_model.token_embed([TEXTS_1[0]]) assert len(outputs) == 1 emb = torch.as_tensor(outputs[0]) @@ -236,53 +247,25 @@ def test_colbert_token_embed( assert emb.shape[1] == colbert_dim assert emb.shape[0] > 1 - -def test_colbert_late_interaction_1_to_1( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, -): - """Test ColBERT late interaction scoring with 1:1 query-document pair.""" - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXTS_1[0]]) - d_outputs = vllm_model.token_embed([TEXTS_2[0]]) + def test_colbert_late_interaction_1_to_1(self, colbert_model): + """Test ColBERT late interaction scoring with 1:1 query-document pair.""" + q_outputs = colbert_model.token_embed([TEXTS_1[0]]) + d_outputs = colbert_model.token_embed([TEXTS_2[0]]) q_emb = torch.as_tensor(q_outputs[0]) d_emb = torch.as_tensor(d_outputs[0]) manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0]) + vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0]) assert len(vllm_scores) == 1 assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) - -def test_colbert_late_interaction_1_to_N( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, -): - """Test ColBERT late interaction scoring with 1:N query-documents.""" - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXTS_1[0]]) - d_outputs = vllm_model.token_embed(TEXTS_2) + def test_colbert_late_interaction_1_to_N(self, colbert_model): + """Test ColBERT late interaction scoring with 1:N query-documents.""" + q_outputs = colbert_model.token_embed([TEXTS_1[0]]) + d_outputs = colbert_model.token_embed(TEXTS_2) q_emb = torch.as_tensor(q_outputs[0]) @@ -291,30 +274,16 @@ def test_colbert_late_interaction_1_to_N( d_emb = torch.as_tensor(d_out) manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) - vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2) + vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2) assert len(vllm_scores) == 2 for i in range(2): assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) - -def test_colbert_late_interaction_N_to_N( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, -): - """Test ColBERT late interaction scoring with N:N query-documents.""" - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - q_outputs = vllm_model.token_embed(TEXTS_1) - d_outputs = vllm_model.token_embed(TEXTS_2) + def test_colbert_late_interaction_N_to_N(self, colbert_model): + """Test ColBERT late interaction scoring with N:N query-documents.""" + q_outputs = colbert_model.token_embed(TEXTS_1) + d_outputs = colbert_model.token_embed(TEXTS_2) manual_scores = [] for q_out, d_out in zip(q_outputs, d_outputs): @@ -322,61 +291,31 @@ def test_colbert_late_interaction_N_to_N( d_emb = torch.as_tensor(d_out) manual_scores.append(compute_maxsim_score(q_emb, d_emb).item()) - vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2) + vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2) assert len(vllm_scores) == 2 for i in range(2): assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01) + def test_colbert_relevance_ordering(self, colbert_model): + """Test that ColBERT scores relevant documents higher than irrelevant.""" + query = "What is machine learning?" + documents = [ + "Machine learning is a subset of artificial intelligence.", + "Python is a programming language.", + "Deep learning uses neural networks.", + ] -def test_colbert_relevance_ordering( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, -): - """Test that ColBERT scores relevant documents higher than irrelevant.""" - query = "What is machine learning?" - documents = [ - "Machine learning is a subset of artificial intelligence.", - "Python is a programming language.", - "Deep learning uses neural networks.", - ] - - with vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = colbert_model.score(query, documents) assert len(scores) == 3 assert scores[0] > scores[1], "ML doc should score higher than Python doc" assert scores[2] > scores[1], "DL doc should score higher than Python doc" - -def test_colbert_embed_not_supported( - vllm_runner, - colbert_model_name, - colbert_max_model_len, - colbert_extra_kwargs, -): - """Test that ColBERT model does not support 'embed' task.""" - with ( - vllm_runner( - colbert_model_name, - runner="pooling", - dtype=DTYPE, - max_model_len=colbert_max_model_len, - enforce_eager=True, - **colbert_extra_kwargs, - ) as vllm_model, - pytest.raises(ValueError, match="Embedding API is not supported"), - ): - vllm_model.embed([TEXTS_1[0]]) + def test_colbert_embed_not_supported(self, colbert_model): + """Test that ColBERT model does not support the embed task.""" + with pytest.raises(ValueError, match="Embedding API is not supported"): + colbert_model.embed([TEXTS_1[0]]) @pytest.mark.parametrize( diff --git a/tests/models/language/pooling/test_truncation_control.py b/tests/models/language/pooling/test_truncation_control.py index 50e8cdbd064c..c4485195558a 100644 --- a/tests/models/language/pooling/test_truncation_control.py +++ b/tests/models/language/pooling/test_truncation_control.py @@ -22,60 +22,47 @@ field.""" +@pytest.fixture(scope="module") +def vllm_model(vllm_runner): + with vllm_runner( + MODEL_NAME, runner="pooling", max_model_len=max_model_len + ) as model: + yield model + + def test_smaller_truncation_size( - vllm_runner, model_name=MODEL_NAME, input_str=input_str + vllm_model, ): truncate_prompt_tokens = 10 - with vllm_runner( - model_name, runner="pooling", max_model_len=max_model_len - ) as vllm_model: - vllm_output = vllm_model.llm.embed( - input_str, - tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), - ) + vllm_output = vllm_model.llm.embed( + input_str, + tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), + ) prompt_tokens = vllm_output[0].prompt_token_ids assert len(prompt_tokens) == truncate_prompt_tokens -def test_max_truncation_size(vllm_runner, model_name=MODEL_NAME, input_str=input_str): +def test_max_truncation_size(vllm_model): truncate_prompt_tokens = -1 - with vllm_runner( - model_name, runner="pooling", max_model_len=max_model_len - ) as vllm_model: - vllm_output = vllm_model.llm.embed( - input_str, - tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), - ) + vllm_output = vllm_model.llm.embed( + input_str, + tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), + ) prompt_tokens = vllm_output[0].prompt_token_ids assert len(prompt_tokens) == max_model_len -def test_bigger_truncation_size( - vllm_runner, model_name=MODEL_NAME, input_str=input_str -): +def test_bigger_truncation_size(vllm_model): truncate_prompt_tokens = max_model_len + 1 - with ( - pytest.raises(VLLMValidationError), - vllm_runner( - model_name, runner="pooling", max_model_len=max_model_len - ) as vllm_model, - ): - llm_output = vllm_model.llm.embed( + with pytest.raises(VLLMValidationError): + vllm_model.llm.embed( input_str, tokenization_kwargs=dict(truncate_prompt_tokens=truncate_prompt_tokens), ) - - assert ( - llm_output - == f"""truncate_prompt_tokens value - ({truncate_prompt_tokens}) is greater than - max_model_len ({max_model_len}). Please, select - a smaller truncation size.""" - ) diff --git a/tests/models/multimodal/generation/test_mm_prefix_lm.py b/tests/models/multimodal/generation/test_mm_prefix_lm.py index 8d3f5b77b715..10ff63f6c0fe 100644 --- a/tests/models/multimodal/generation/test_mm_prefix_lm.py +++ b/tests/models/multimodal/generation/test_mm_prefix_lm.py @@ -8,6 +8,7 @@ from transformers import AutoModelForImageTextToText from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from ....conftest import HfRunner, ImageTestAssets, VllmRunner from .vlm_utils import model_utils @@ -30,7 +31,9 @@ def _install_prefill_hidden_capture(model): def forward(*args, **kwargs): hidden_states = original_forward(*args, **kwargs) if model._prefill_hidden is None and torch.is_tensor(hidden_states): - model._prefill_hidden = hidden_states.detach().float().cpu() + # Capturing hidden states for comparison is a deliberate D2H. + with gpu_sync_allowed(): + model._prefill_hidden = hidden_states.detach().float().cpu() return hidden_states language_model.forward = forward diff --git a/tests/models/multimodal/generation/test_phi4mm.py b/tests/models/multimodal/generation/test_phi4mm.py index 5ab75e145aee..71dbfcd5f357 100644 --- a/tests/models/multimodal/generation/test_phi4mm.py +++ b/tests/models/multimodal/generation/test_phi4mm.py @@ -66,6 +66,11 @@ def vllm_to_hf_output( target_dtype = "half" +IMAGE_SIZE_FACTOR_GROUPS = ( + (1.0,), + (1.0, 1.0, 1.0), + (0.25, 0.5, 1.0), +) def run_test( @@ -167,17 +172,6 @@ def patch_hf_processor( @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [1.0], - # Single-scale, batched - [1.0, 1.0, 1.0], - # Multi-scale - [0.25, 0.5, 1.0], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_model_len", [12800]) @pytest.mark.parametrize("max_tokens", [128]) @@ -187,7 +181,6 @@ def test_models( vllm_runner, image_assets, model, - size_factors, dtype: str, max_model_len: int, max_tokens: int, @@ -201,6 +194,7 @@ def test_models( [rescale_image_size(image, factor) for factor in size_factors], None, ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS for image, prompt in zip(images, HF_IMAGE_PROMPTS) ] @@ -220,19 +214,6 @@ def test_models( @large_gpu_test(min_gb=48) @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # No image - # [], - # Single-scale - [1.0], - # Single-scale, batched - [1.0, 1.0, 1.0], - # Multi-scale - [0.25, 0.5, 1.0], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_model_len", [25600]) @pytest.mark.parametrize("max_tokens", [128]) @@ -242,7 +223,6 @@ def test_multi_images_models( vllm_runner, image_assets, model, - size_factors, dtype: str, max_model_len: int, max_tokens: int, @@ -258,7 +238,8 @@ def test_multi_images_models( for factor in size_factors ], None, - ), + ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS ] run_test( diff --git a/tests/models/multimodal/generation/test_qwen2_vl.py b/tests/models/multimodal/generation/test_qwen2_vl.py index 6148c0bcda7d..19328cad7a23 100644 --- a/tests/models/multimodal/generation/test_qwen2_vl.py +++ b/tests/models/multimodal/generation/test_qwen2_vl.py @@ -29,6 +29,16 @@ def enable_pickle(monkeypatch): models = ["Qwen/Qwen2-VL-2B-Instruct"] target_dtype = "half" +IMAGE_SIZE_FACTOR_GROUPS = ( + (0.5,), + (0.5, 0.5), + (0.25, 0.5, 0.5), +) +VIDEO_SIZE_FACTOR_GROUPS = ( + (0.5,), + (0.5, 0.5), + (0.25, 0.25, 0.5), +) IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>" @@ -323,17 +333,6 @@ def run_embedding_input_test( @pytest.mark.core_model @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [0.5], - # Single-scale, batched - [0.5, 0.5], - # Multi-scale - [0.25, 0.5, 0.5], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [10]) @@ -341,7 +340,6 @@ def test_qwen2_vl_image_embeddings_input( vllm_runner, image_assets, model, - size_factors, dtype, max_tokens, num_logprobs, @@ -355,6 +353,7 @@ def test_qwen2_vl_image_embeddings_input( [rescale_image_size(image, factor) for factor in size_factors], [], ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS for image, prompt in zip(images, IMAGE_PROMPTS) ] @@ -372,17 +371,6 @@ def test_qwen2_vl_image_embeddings_input( @pytest.mark.core_model @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [0.5], - # Single-scale, batched - [0.5, 0.5], - # Multi-scale - [0.25, 0.5, 0.5], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [10]) @@ -390,7 +378,6 @@ def test_qwen2_vl_multiple_image_embeddings_input( vllm_runner, image_assets, model, - size_factors, dtype: str, max_tokens: int, num_logprobs: int, @@ -406,6 +393,7 @@ def test_qwen2_vl_multiple_image_embeddings_input( ], [], ) + for size_factors in IMAGE_SIZE_FACTOR_GROUPS ] run_embedding_input_test( @@ -422,17 +410,6 @@ def test_qwen2_vl_multiple_image_embeddings_input( @pytest.mark.core_model @pytest.mark.parametrize("model", models) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [0.5], - # Single-scale, batched - [0.5, 0.5], - # Multi-scale - [0.25, 0.25, 0.5], - ], -) @pytest.mark.parametrize("dtype", [target_dtype]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [10]) @@ -440,7 +417,6 @@ def test_qwen2_vl_video_embeddings_input( vllm_runner, video_assets, model, - size_factors, dtype: str, max_tokens: int, num_logprobs: int, @@ -457,6 +433,7 @@ def test_qwen2_vl_video_embeddings_input( [], [rescale_video_size(video, factor) for factor in size_factors], ) + for size_factors in VIDEO_SIZE_FACTOR_GROUPS for video, prompt in zip(sampled_vids, VIDEO_PROMPTS) ] diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index d927c547f49f..8e12790925dd 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -286,6 +286,12 @@ def ernie45_vl_chat_template(content: str) -> str: "user\n<|video|>\nDescribe this video in one sentence." "\nmodel\n" ), + # The 16-frame test video produces 1056 vision tokens. Capture only + # the smallest supported bucket that covers it instead of all default + # buckets through max_model_len, which adds unrelated memory pressure. + compilation_config_overrides={ + "encoder_cudagraph_token_budgets": [1120], + }, needs_video_metadata=True, marks=[pytest.mark.core_model], ), @@ -307,7 +313,9 @@ def get_compilation_config(config: VitCudagraphTestConfig): @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS)) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) def test_vit_cudagraph_image(model_id, vllm_runner, image_assets): config = MODEL_CONFIGS[model_id] @@ -351,7 +359,9 @@ def test_vit_cudagraph_image(model_id, vllm_runner, image_assets): @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS)) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) def test_vit_cudagraph_video(model_id, vllm_runner, video_assets): config = MODEL_CONFIGS[model_id] diff --git a/tests/models/multimodal/generation/test_whisper.py b/tests/models/multimodal/generation/test_whisper.py index 310c7fe0f563..fef6b69911b7 100644 --- a/tests/models/multimodal/generation/test_whisper.py +++ b/tests/models/multimodal/generation/test_whisper.py @@ -20,6 +20,7 @@ HF_PROMPT = "" # Whisper expects 16kHz audio WHISPER_SAMPLE_RATE = 16000 +BEAM_WIDTHS = (1, 2) @pytest.fixture(autouse=True) @@ -127,13 +128,11 @@ def check_model_available(model: str) -> None: @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", [64]) -@pytest.mark.parametrize("beam_width", [1, 2]) def test_beam_search_encoder_decoder( hf_runner, vllm_runner, dtype: str, max_tokens: int, - beam_width: int, resampled_assets, ) -> None: """Test beam search with encoder-decoder models (Whisper).""" @@ -146,12 +145,15 @@ def test_beam_search_encoder_decoder( ] with hf_runner(model, dtype=dtype, auto_cls=AutoModelForSpeechSeq2Seq) as hf_model: - hf_outputs = hf_model.generate_beam_search( - hf_prompts, - beam_width=beam_width, - max_tokens=max_tokens, - audios=resampled_assets, - ) + hf_outputs_by_beam_width = [ + hf_model.generate_beam_search( + hf_prompts, + beam_width=beam_width, + max_tokens=max_tokens, + audios=resampled_assets, + ) + for beam_width in BEAM_WIDTHS + ] # Test both explicit encoder/decoder prompts vllm_prompts = [ @@ -179,38 +181,46 @@ def test_beam_search_encoder_decoder( limit_mm_per_prompt={"audio": 2}, enforce_eager=True, ) as vllm_model: - vllm_outputs = vllm_model.generate_beam_search( - vllm_prompts, - beam_width=beam_width, - max_tokens=max_tokens, - ) - - for i in range(len(vllm_prompts)): - hf_output_ids, hf_output_texts = hf_outputs[i] - vllm_output_ids, vllm_output_texts = vllm_outputs[i] - - for j, (hf_text, vllm_text) in enumerate( - zip(hf_output_texts, vllm_output_texts) - ): - print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:") - print(hf_text) - print(f">>>{j}-th vllm output:") - print(vllm_text) - - # Check that we got the same number of beams - assert len(hf_output_ids) == len(vllm_output_ids) - - # For encoder-decoder models, we primarily want to verify that: - # 1. Beam search completes without errors - # 2. We get the expected number of beams - # 3. Outputs are reasonable (non-empty, diverse beams) - for j in range(len(vllm_output_ids)): - # Check that outputs are not empty - assert len(vllm_output_ids[j]) > 0, f"Prompt {i}, beam {j}: empty output" - # Check that decoded text is not empty - assert len(vllm_output_texts[j].strip()) > 0, ( - f"Prompt {i}, beam {j}: empty text output" + vllm_outputs_by_beam_width = [ + vllm_model.generate_beam_search( + vllm_prompts, + beam_width=beam_width, + max_tokens=max_tokens, ) + for beam_width in BEAM_WIDTHS + ] + + for beam_width, hf_outputs, vllm_outputs in zip( + BEAM_WIDTHS, hf_outputs_by_beam_width, vllm_outputs_by_beam_width + ): + for i in range(len(vllm_prompts)): + hf_output_ids, hf_output_texts = hf_outputs[i] + vllm_output_ids, vllm_output_texts = vllm_outputs[i] + + for j, (hf_text, vllm_text) in enumerate( + zip(hf_output_texts, vllm_output_texts) + ): + print(f">>>{j}-th hf output [NOTE: special tokens are filtered]:") + print(hf_text) + print(f">>>{j}-th vllm output:") + print(vllm_text) + + # Check that we got the same number of beams + assert len(hf_output_ids) == len(vllm_output_ids) == beam_width + + # For encoder-decoder models, we primarily want to verify that: + # 1. Beam search completes without errors + # 2. We get the expected number of beams + # 3. Outputs are reasonable (non-empty, diverse beams) + for j in range(len(vllm_output_ids)): + # Check that outputs are not empty + assert len(vllm_output_ids[j]) > 0, ( + f"Prompt {i}, beam {j}: empty output" + ) + # Check that decoded text is not empty + assert len(vllm_output_texts[j].strip()) > 0, ( + f"Prompt {i}, beam {j}: empty text output" + ) def test_parse_language_detection_output(): diff --git a/tests/models/multimodal/generation/vlm_utils/case_filtering.py b/tests/models/multimodal/generation/vlm_utils/case_filtering.py index 116eead7a70a..cbe660425751 100644 --- a/tests/models/multimodal/generation/vlm_utils/case_filtering.py +++ b/tests/models/multimodal/generation/vlm_utils/case_filtering.py @@ -94,7 +94,8 @@ def get_model_type_cases( test_info.needs_video_metadata ) - # No sizes passed for custom inputs, since inputs are directly provided + # Keep all size batches in one test case so they share the same model + # instances. No sizes are passed for preprocessed audio or custom inputs. if test_type not in ( VLMTestType.CUSTOM_INPUTS, VLMTestType.AUDIO, @@ -102,7 +103,9 @@ def get_model_type_cases( wrapped_sizes = get_wrapped_test_sizes(test_info, test_type) if wrapped_sizes is None: raise ValueError(f"Sizes must be set for test type {test_type}") - iter_kwargs["size_wrapper"] = wrapped_sizes + if not wrapped_sizes: + return [] + iter_kwargs["size_wrappers"] = (wrapped_sizes,) # Otherwise expand the custom test options instead elif test_type == VLMTestType.CUSTOM_INPUTS: @@ -127,9 +130,8 @@ def get_parametrized_options( create_new_process_for_each_test: bool, ): """Converts all of our VLMTestInfo into an expanded list of parameters. - This is similar to nesting pytest parametrize calls, but done directly - through an itertools product so that each test can set things like - size factors etc, while still running in isolated test cases. + Runner configuration values are expanded through an itertools product. + Input size batches stay grouped so they can share model instances. """ matching_tests = get_filtered_test_settings( test_settings, test_type, create_new_process_for_each_test @@ -149,8 +151,7 @@ def get_wrapped_test_sizes( test_info: VLMTestInfo, test_type: VLMTestType ) -> tuple[ImageSizeWrapper, ...]: """Given a test info which may have size factors or fixed sizes, wrap them - and combine them into an iterable, each of which will be used in parameter - expansion. + and combine them into an iterable of request batches. Args: test_info: Test configuration to be expanded. diff --git a/tests/models/multimodal/generation/vlm_utils/runners.py b/tests/models/multimodal/generation/vlm_utils/runners.py index 218339ef1dff..571ea54250b7 100644 --- a/tests/models/multimodal/generation/vlm_utils/runners.py +++ b/tests/models/multimodal/generation/vlm_utils/runners.py @@ -4,6 +4,7 @@ types / modalities. """ +import itertools from pathlib import PosixPath from .....conftest import ( @@ -27,9 +28,14 @@ def run_single_image_test( vllm_runner: type[VllmRunner], image_assets: ImageTestAssets, ): - assert test_case.size_wrapper is not None - inputs = builders.build_single_image_inputs_from_test_info( - model_test_info, image_assets, test_case.size_wrapper, tmp_path + assert test_case.size_wrappers + inputs = list( + itertools.chain.from_iterable( + builders.build_single_image_inputs_from_test_info( + model_test_info, image_assets, size_wrapper, tmp_path + ) + for size_wrapper in test_case.size_wrappers + ) ) core.run_test( @@ -55,9 +61,14 @@ def run_multi_image_test( vllm_runner: type[VllmRunner], image_assets: ImageTestAssets, ): - assert test_case.size_wrapper is not None - inputs = builders.build_multi_image_inputs_from_test_info( - model_test_info, image_assets, test_case.size_wrapper, tmp_path + assert test_case.size_wrappers + inputs = list( + itertools.chain.from_iterable( + builders.build_multi_image_inputs_from_test_info( + model_test_info, image_assets, size_wrapper, tmp_path + ) + for size_wrapper in test_case.size_wrappers + ) ) core.run_test( @@ -82,9 +93,20 @@ def run_embedding_test( vllm_runner: type[VllmRunner], image_assets: ImageTestAssets, ): - assert test_case.size_wrapper is not None - inputs, vllm_embeddings = builders.build_embedding_inputs_from_test_info( - model_test_info, image_assets, test_case.size_wrapper + assert test_case.size_wrappers + inputs_and_embeddings = [ + builders.build_embedding_inputs_from_test_info( + model_test_info, image_assets, size_wrapper + ) + for size_wrapper in test_case.size_wrappers + ] + inputs = list( + itertools.chain.from_iterable(inputs for inputs, _ in inputs_and_embeddings) + ) + vllm_embeddings = list( + itertools.chain.from_iterable( + embeddings for _, embeddings in inputs_and_embeddings + ) ) core.run_test( @@ -110,14 +132,19 @@ def run_video_test( vllm_runner: type[VllmRunner], video_assets: VideoTestAssets, ): - assert test_case.size_wrapper is not None + assert test_case.size_wrappers assert test_case.num_video_frames is not None - inputs = builders.build_video_inputs_from_test_info( - model_test_info, - video_assets, - test_case.size_wrapper, - test_case.num_video_frames, - test_case.needs_video_metadata, + inputs = list( + itertools.chain.from_iterable( + builders.build_video_inputs_from_test_info( + model_test_info, + video_assets, + size_wrapper, + test_case.num_video_frames, + test_case.needs_video_metadata, + ) + for size_wrapper in test_case.size_wrappers + ) ) core.run_test( diff --git a/tests/models/multimodal/generation/vlm_utils/types.py b/tests/models/multimodal/generation/vlm_utils/types.py index af48a1479bad..3722caed2cb3 100644 --- a/tests/models/multimodal/generation/vlm_utils/types.py +++ b/tests/models/multimodal/generation/vlm_utils/types.py @@ -158,11 +158,9 @@ class VLMTestInfo(NamedTuple): num_video_frames: int | tuple[int] = 16 needs_video_metadata: bool = False - # Fixed image sizes / image size factors; most tests use image_size_factors - # The values provided for these two fields will be stacked and expanded - # such that each model will consider each image size factor / image size - # once per tests (much like concatenating and wrapping in one parametrize - # call) + # Fixed image sizes / image size factors; most tests use image_size_factors. + # Each inner iterable defines one request batch. All batches are run against + # the same model instance. image_size_factors: Iterable[Iterable[float]] = IMAGE_SIZE_FACTORS image_sizes: Iterable[Iterable[tuple[int, int]]] | None = None @@ -211,8 +209,8 @@ class ExpandableVLMTestArgs(NamedTuple): num_logprobs: int dtype: str distributed_executor_backend: str | None - # Sizes are used for everything except for custom input tests - size_wrapper: ImageSizeWrapper | None = None + # Sizes are used for everything except audio and custom input tests. + size_wrappers: tuple[ImageSizeWrapper, ...] = () # Video only num_video_frames: int | None = None needs_video_metadata: bool = False diff --git a/tests/models/multimodal/pooling/test_clip.py b/tests/models/multimodal/pooling/test_clip.py index 14ede6c1d328..54dd8556e903 100644 --- a/tests/models/multimodal/pooling/test_clip.py +++ b/tests/models/multimodal/pooling/test_clip.py @@ -26,8 +26,7 @@ def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], - input_texts: list[str], - input_images: PromptImageInput, + input_cases: list[tuple[list[str], PromptImageInput]], model: str, *, dtype: str, @@ -39,105 +38,75 @@ def _run_test( with vllm_runner( model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77 ) as vllm_model: - vllm_outputs = vllm_model.embed(input_texts, images=input_images) - - with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model: - all_inputs = hf_model.get_inputs(input_texts, images=input_images) - - all_outputs = [] - for inputs in all_inputs: - inputs = hf_model.wrap_device(inputs) - - if "pixel_values" in inputs: - pooled_output = hf_model.model.get_image_features( - pixel_values=inputs.pixel_values, - ) - else: - pooled_output = hf_model.model.get_text_features( - input_ids=inputs.input_ids, - attention_mask=inputs.attention_mask, - ) - - if not isinstance(pooled_output, torch.Tensor): - pooled_output = pooled_output.pooler_output - pooled_output = pooled_output.squeeze(0) - all_outputs.append(pooled_output.tolist()) - - hf_outputs = all_outputs - - check_embeddings_close( - embeddings_0_lst=hf_outputs, - embeddings_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) + vllm_outputs_per_case = [ + vllm_model.embed(input_texts, images=input_images) + for input_texts, input_images in input_cases + ] + texts = [HF_TEXT_PROMPTS[0]] + images = [input_cases[1][1][0]] + with pytest.raises(ValueError, match="not both"): + vllm_model.embed(texts, images=images) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_text( - hf_runner, - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] + # Should still be able to run subsequent requests + vllm_model.embed(texts) + vllm_model.embed([""], images=images) - _run_test( - hf_runner, - vllm_runner, - input_texts, - input_images, # type: ignore - model, - dtype=dtype, - ) + with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model: + hf_outputs_per_case = [] + for input_texts, input_images in input_cases: + all_inputs = hf_model.get_inputs(input_texts, images=input_images) + + hf_outputs = [] + for inputs in all_inputs: + inputs = hf_model.wrap_device(inputs) + + if "pixel_values" in inputs: + pooled_output = hf_model.model.get_image_features( + pixel_values=inputs.pixel_values, + ) + else: + pooled_output = hf_model.model.get_text_features( + input_ids=inputs.input_ids, + attention_mask=inputs.attention_mask, + ) + + if not isinstance(pooled_output, torch.Tensor): + pooled_output = pooled_output.pooler_output + pooled_output = pooled_output.squeeze(0) + hf_outputs.append(pooled_output.tolist()) + + hf_outputs_per_case.append(hf_outputs) + + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["float"]) -def test_models_image( +def test_models( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - input_texts_images = [ - (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) + text_images = [None] * len(HF_TEXT_PROMPTS) + images = [asset.pil_image for asset in image_assets] + input_cases = [ + (HF_TEXT_PROMPTS, text_images), + (HF_IMAGE_PROMPTS, images), ] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] _run_test( hf_runner, vllm_runner, - input_texts, - input_images, + input_cases, # type: ignore[arg-type] model, dtype=dtype, ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_text_image_no_crash( - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - texts = [HF_TEXT_PROMPTS[0]] - images = [image_assets[0].pil_image] - - with vllm_runner( - model, runner="pooling", dtype=dtype, enforce_eager=True, max_model_len=77 - ) as vllm_model: - with pytest.raises(ValueError, match="not both"): - vllm_model.embed(texts, images=images) - - # Should still be able to run subsequent requests - vllm_model.embed(texts) - vllm_model.embed([""], images=images) diff --git a/tests/models/multimodal/pooling/test_colmodernvbert.py b/tests/models/multimodal/pooling/test_colmodernvbert.py index efeb3195b15b..3dffc66e7a8d 100644 --- a/tests/models/multimodal/pooling/test_colmodernvbert.py +++ b/tests/models/multimodal/pooling/test_colmodernvbert.py @@ -17,29 +17,34 @@ DTYPE = "half" -# ----------------------------------------------------------------------- -# Text-only tests -# ----------------------------------------------------------------------- - - -def test_colmodernvbert_text_token_embed(vllm_runner): - """Text query produces per-token embeddings with shape (seq_len, 128).""" +@pytest.fixture(scope="module") +def colmodernvbert_model(vllm_runner): with vllm_runner( MODEL_NAME, runner="pooling", dtype=DTYPE, enforce_eager=True, ) as vllm_model: - outputs = vllm_model.token_embed(["What is machine learning?"]) + yield vllm_model + + +# ----------------------------------------------------------------------- +# Text-only tests +# ----------------------------------------------------------------------- + - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - assert emb.dim() == 2 - assert emb.shape[1] == COLBERT_DIM - assert emb.shape[0] > 1 +def test_colmodernvbert_text_token_embed(colmodernvbert_model): + """Text query produces per-token embeddings with shape (seq_len, 128).""" + outputs = colmodernvbert_model.token_embed(["What is machine learning?"]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == COLBERT_DIM + assert emb.shape[0] > 1 -def test_colmodernvbert_text_relevance_ordering(vllm_runner): +def test_colmodernvbert_text_relevance_ordering(colmodernvbert_model): """Relevant documents score higher than irrelevant ones.""" query = "What is machine learning?" documents = [ @@ -47,40 +52,28 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner): "The weather in Paris is mild in spring.", ] - with vllm_runner( - MODEL_NAME, - runner="pooling", - dtype=DTYPE, - enforce_eager=True, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = colmodernvbert_model.score(query, documents) - assert len(scores) == 2 - assert scores[0] > scores[1], "ML doc should score higher than weather doc" + assert len(scores) == 2 + assert scores[0] > scores[1], "ML doc should score higher than weather doc" -def test_colmodernvbert_text_late_interaction(vllm_runner): +def test_colmodernvbert_text_late_interaction(colmodernvbert_model): """MaxSim scoring via vLLM matches manual computation.""" query = "What is the capital of France?" doc = "The capital of France is Paris." - with vllm_runner( - MODEL_NAME, - runner="pooling", - dtype=DTYPE, - enforce_eager=True, - ) as vllm_model: - q_out = vllm_model.token_embed([query]) - d_out = vllm_model.token_embed([doc]) + q_out = colmodernvbert_model.token_embed([query]) + d_out = colmodernvbert_model.token_embed([doc]) - q_emb = torch.tensor(q_out[0]) - d_emb = torch.tensor(d_out[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + q_emb = torch.tensor(q_out[0]) + d_emb = torch.tensor(d_out[0]) + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(query, doc) + vllm_scores = colmodernvbert_model.score(query, doc) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) # ----------------------------------------------------------------------- @@ -88,28 +81,22 @@ def test_colmodernvbert_text_late_interaction(vllm_runner): # ----------------------------------------------------------------------- -def test_colmodernvbert_image_token_embed(vllm_runner, image_assets): +def test_colmodernvbert_image_token_embed(colmodernvbert_model, image_assets): """Image input produces per-token embeddings including vision tokens.""" - with vllm_runner( - MODEL_NAME, - runner="pooling", - dtype=DTYPE, - enforce_eager=True, - ) as vllm_model: - image = image_assets[0].pil_image - inputs = vllm_model.get_inputs( - [""], - images=[image], - ) - req_outputs = vllm_model.llm.encode( - inputs, - pooling_task="token_embed", - ) - outputs = [req_output.outputs.data for req_output in req_outputs] - - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - assert emb.dim() == 2 - assert emb.shape[1] == COLBERT_DIM - # Should have at least the image tokens (64 after pixel shuffle) - assert emb.shape[0] >= 64 + image = image_assets[0].pil_image + inputs = colmodernvbert_model.get_inputs( + [""], + images=[image], + ) + req_outputs = colmodernvbert_model.llm.encode( + inputs, + pooling_task="token_embed", + ) + outputs = [req_output.outputs.data for req_output in req_outputs] + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + assert emb.dim() == 2 + assert emb.shape[1] == COLBERT_DIM + # Should have at least the image tokens (64 after pixel shuffle) + assert emb.shape[0] >= 64 diff --git a/tests/models/multimodal/pooling/test_colpali.py b/tests/models/multimodal/pooling/test_colpali.py index 7c91731065bb..9e040bf9a1e5 100644 --- a/tests/models/multimodal/pooling/test_colpali.py +++ b/tests/models/multimodal/pooling/test_colpali.py @@ -19,6 +19,7 @@ ChatCompletionContentPartTextParam, ) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.platforms import current_platform from ....conftest import VllmRunner @@ -74,75 +75,51 @@ def _make_image_mm_param( def _run_token_embed_test( - vllm_runner: type[VllmRunner], + vllm_model: VllmRunner, model: str, - *, - dtype: str, ) -> None: """Verify per-token embedding shape and L2 normalization.""" - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) - - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - # Token embeddings should be 2D: [num_tokens, embed_dim] - assert emb.dim() == 2 - assert emb.shape[1] == EMBED_DIMS[model] - assert emb.shape[0] > 1 - - # Verify L2 normalization - norms = torch.norm(emb, p=2, dim=-1) - torch.testing.assert_close( - norms, - torch.ones_like(norms), - rtol=1e-2, - atol=1e-2, - ) + outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + # Token embeddings should be 2D: [num_tokens, embed_dim] + assert emb.dim() == 2 + assert emb.shape[1] == EMBED_DIMS[model] + assert emb.shape[0] > 1 + + # Verify L2 normalization + norms = torch.norm(emb, p=2, dim=-1) + torch.testing.assert_close( + norms, + torch.ones_like(norms), + rtol=1e-2, + atol=1e-2, + ) def _run_late_interaction_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify MaxSim scoring matches manual computation.""" from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) - d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) + q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) - q_emb = torch.tensor(q_outputs[0]) - d_emb = torch.tensor(d_outputs[0]) + q_emb = torch.tensor(q_outputs[0]) + d_emb = torch.tensor(d_outputs[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) + vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) def _run_relevance_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify that relevant documents score higher than irrelevant ones.""" query = "What is machine learning?" @@ -152,59 +129,18 @@ def _run_relevance_test( "Deep learning uses neural networks for complex tasks.", ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - scores = vllm_model.score(query, documents) - - assert len(scores) == 3 - assert scores[0] > scores[1], "ML doc should score higher than weather doc" - assert scores[2] > scores[1], "DL doc should score higher than weather doc" - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_token_embed( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_token_embed_test(vllm_runner, model, dtype=dtype) - + scores = vllm_model.score(query, documents) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_late_interaction_scoring( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_late_interaction_test(vllm_runner, model, dtype=dtype) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_relevance_ordering( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_relevance_test(vllm_runner, model, dtype=dtype) + assert len(scores) == 3 + assert scores[0] > scores[1], "ML doc should score higher than weather doc" + assert scores[2] > scores[1], "DL doc should score higher than weather doc" # ── Multimodal scoring tests ──────────────────────────────── def _run_multimodal_text_query_image_docs_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Score a text query against image documents via the multimodal path.""" red_image = _make_base64_image(64, 64, color=(255, 0, 0)) @@ -215,27 +151,15 @@ def _run_multimodal_text_query_image_docs_test( _make_image_mm_param(red_image), _make_image_mm_param(blue_image), ] + scores = vllm_model.llm.score(query, image_docs) - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - scores = vllm_model.llm.score(query, image_docs) - - assert len(scores) == 2 - for s in scores: - assert isinstance(s.outputs.score, float) + assert len(scores) == 2 + for s in scores: + assert isinstance(s.outputs.score, float) def _run_multimodal_mixed_docs_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Score a text query against a mix of text and image documents.""" red_image = _make_base64_image(64, 64, color=(255, 0, 0)) @@ -246,28 +170,17 @@ def _run_multimodal_mixed_docs_test( _make_image_mm_param(red_image), ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - scores = vllm_model.llm.score(query, documents) + scores = vllm_model.llm.score(query, documents) - assert len(scores) == 2 - for s in scores: - assert isinstance(s.outputs.score, float) - # Text document about France should score higher than a random image - assert scores[0].outputs.score > scores[1].outputs.score + assert len(scores) == 2 + for s in scores: + assert isinstance(s.outputs.score, float) + # Text document about France should score higher than a random image + assert scores[0].outputs.score > scores[1].outputs.score def _run_multimodal_image_query_text_docs_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Score an image query against text documents.""" red_image = _make_base64_image(64, 64, color=(255, 0, 0)) @@ -278,46 +191,54 @@ def _run_multimodal_image_query_text_docs_test( "The weather forecast shows rain tomorrow.", ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - gpu_memory_utilization=GPU_MEMORY_UTILIZATION, - ) as vllm_model: - scores = vllm_model.llm.score(image_query, documents) - - assert len(scores) == 2 - for s in scores: - assert isinstance(s.outputs.score, float) + scores = vllm_model.llm.score(image_query, documents) - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_multimodal_text_query_image_docs( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype) + assert len(scores) == 2 + for s in scores: + assert isinstance(s.outputs.score, float) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_multimodal_mixed_docs( +def test_colpali_default_runner( vllm_runner, model: str, dtype: str, ) -> None: - _run_multimodal_mixed_docs_test(vllm_runner, model, dtype=dtype) + with vllm_runner( + model, + runner="pooling", + dtype=dtype, + max_model_len=4096, + enforce_eager=True, + gpu_memory_utilization=GPU_MEMORY_UTILIZATION, + ) as vllm_model: + _run_token_embed_test(vllm_model, model) + _run_late_interaction_test(vllm_model) + _run_relevance_test(vllm_model) + _run_multimodal_mixed_docs_test(vllm_model) + _run_multimodal_image_query_text_docs_test(vllm_model) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", [DTYPE]) -def test_colpali_multimodal_image_query_text_docs( +def test_colpali_v2_multimodal_text_query_image_docs( vllm_runner, + monkeypatch: pytest.MonkeyPatch, model: str, dtype: str, ) -> None: - _run_multimodal_image_query_text_docs_test(vllm_runner, model, dtype=dtype) + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + attention_backend = "FLASH_ATTN" if current_platform.is_cuda() else None + with vllm_runner( + model, + runner="pooling", + dtype=dtype, + max_model_len=4096, + enforce_eager=True, + gpu_memory_utilization=GPU_MEMORY_UTILIZATION, + attention_backend=attention_backend, + kernel_config={"enable_flashinfer_autotune": False}, + ) as vllm_model: + assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner + _run_multimodal_text_query_image_docs_test(vllm_model) diff --git a/tests/models/multimodal/pooling/test_colqwen3_5.py b/tests/models/multimodal/pooling/test_colqwen3_5.py index 43914d819b8a..2aac465d4d0a 100644 --- a/tests/models/multimodal/pooling/test_colqwen3_5.py +++ b/tests/models/multimodal/pooling/test_colqwen3_5.py @@ -35,74 +35,65 @@ DTYPE = "half" -def _run_token_embed_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, -) -> None: - """Verify per-token embedding shape and L2 normalization.""" +@pytest.fixture(scope="module", params=MODELS) +def colqwen3_5_model(request, vllm_runner): + model = request.param with vllm_runner( model, runner="pooling", - dtype=dtype, + dtype=DTYPE, max_model_len=4096, enforce_eager=True, ) as vllm_model: - outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + yield model, vllm_model - assert len(outputs) == 1 - emb = torch.tensor(outputs[0]) - # Token embeddings should be 2D: [num_tokens, embed_dim] - assert emb.dim() == 2 - assert emb.shape[1] == EMBED_DIMS[model] - assert emb.shape[0] > 1 - # Verify L2 normalization - norms = torch.norm(emb, p=2, dim=-1) - torch.testing.assert_close( - norms, - torch.ones_like(norms), - rtol=1e-2, - atol=1e-2, - ) +def _run_token_embed_test( + vllm_model: VllmRunner, + model: str, +) -> None: + """Verify per-token embedding shape and L2 normalization.""" + outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + + assert len(outputs) == 1 + emb = torch.tensor(outputs[0]) + # Token embeddings should be 2D: [num_tokens, embed_dim] + assert emb.dim() == 2 + assert emb.shape[1] == EMBED_DIMS[model] + assert emb.shape[0] > 1 + + # Verify L2 normalization + norms = torch.norm(emb, p=2, dim=-1) + torch.testing.assert_close( + norms, + torch.ones_like(norms), + rtol=1e-2, + atol=1e-2, + ) def _run_late_interaction_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify MaxSim scoring matches manual computation.""" from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - ) as vllm_model: - q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) - d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) + q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]]) + d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]]) - q_emb = torch.tensor(q_outputs[0]) - d_emb = torch.tensor(d_outputs[0]) + q_emb = torch.tensor(q_outputs[0]) + d_emb = torch.tensor(d_outputs[0]) - manual_score = compute_maxsim_score(q_emb, d_emb).item() + manual_score = compute_maxsim_score(q_emb, d_emb).item() - vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) + vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0]) - assert len(vllm_scores) == 1 - assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) + assert len(vllm_scores) == 1 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01) def _run_relevance_test( - vllm_runner: type[VllmRunner], - model: str, - *, - dtype: str, + vllm_model: VllmRunner, ) -> None: """Verify that relevant documents score higher than irrelevant ones.""" query = "What is machine learning?" @@ -112,48 +103,26 @@ def _run_relevance_test( "Deep learning uses neural networks for complex tasks.", ] - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - max_model_len=4096, - enforce_eager=True, - ) as vllm_model: - scores = vllm_model.score(query, documents) + scores = vllm_model.score(query, documents) - assert len(scores) == 3 - assert scores[0] > scores[1], "ML doc should score higher than weather doc" - assert scores[2] > scores[1], "DL doc should score higher than weather doc" + assert len(scores) == 3 + assert scores[0] > scores[1], "ML doc should score higher than weather doc" + assert scores[2] > scores[1], "DL doc should score higher than weather doc" -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colqwen3_5_token_embed( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_token_embed_test(vllm_runner, model, dtype=dtype) +def test_colqwen3_5_token_embed(colqwen3_5_model) -> None: + model, vllm_model = colqwen3_5_model + _run_token_embed_test(vllm_model, model) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colqwen3_5_late_interaction_scoring( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_late_interaction_test(vllm_runner, model, dtype=dtype) +def test_colqwen3_5_late_interaction_scoring(colqwen3_5_model) -> None: + _, vllm_model = colqwen3_5_model + _run_late_interaction_test(vllm_model) -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", [DTYPE]) -def test_colqwen3_5_relevance_ordering( - vllm_runner, - model: str, - dtype: str, -) -> None: - _run_relevance_test(vllm_runner, model, dtype=dtype) +def test_colqwen3_5_relevance_ordering(colqwen3_5_model) -> None: + _, vllm_model = colqwen3_5_model + _run_relevance_test(vllm_model) @pytest.mark.parametrize( diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index a2f1d3424c34..9516ab95d52d 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -9,12 +9,14 @@ Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone. """ +from collections.abc import Sequence from io import BytesIO from pathlib import Path import pybase64 as base64 import pytest import torch +from PIL import Image from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor from vllm.entrypoints.chat_utils import ( @@ -54,17 +56,15 @@ def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], - input_texts: list[str], - input_images: PromptImageInput, + input_cases: list[tuple[list[str], PromptImageInput]], model: str, *, dtype: str, ) -> None: - """Run embedding comparison test between HF and vLLM. + """Compare HF and vLLM embeddings for all input cases. NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing. """ - # Run vLLM inference first with vllm_runner( model, runner="pooling", @@ -74,91 +74,70 @@ def _run_test( trust_remote_code=True, **ROCM_ENGINE_KWARGS, ) as vllm_model: - vllm_outputs = vllm_model.embed(input_texts, images=input_images) + vllm_outputs_per_case = [ + vllm_model.embed(input_texts, images=input_images) + for input_texts, input_images in input_cases + ] - # Run HF inference using the model's encode_queries/encode_documents API with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model: - hf_outputs = [] - for text, image in zip(input_texts, input_images): - with torch.inference_mode(): - if text.startswith(QUERY_PREFIX): - # Strip prefix and use encode_queries for query texts - query_text = text[len(QUERY_PREFIX) :] - embedding = hf_model.model.encode_queries([query_text]) - elif text.startswith(PASSAGE_PREFIX): - # Strip prefix and use encode_documents for passages/images - passage_text = text[len(PASSAGE_PREFIX) :] - if image is not None: - # Image document - pass image to encode_documents - embedding = hf_model.model.encode_documents( - images=[image], - texts=[passage_text], - ) + hf_outputs_per_case = [] + for input_texts, input_images in input_cases: + hf_outputs = [] + for text, image in zip(input_texts, input_images): + with torch.inference_mode(): + if text.startswith(QUERY_PREFIX): + query_text = text[len(QUERY_PREFIX) :] + embedding = hf_model.model.encode_queries([query_text]) + elif text.startswith(PASSAGE_PREFIX): + passage_text = text[len(PASSAGE_PREFIX) :] + if image is not None: + embedding = hf_model.model.encode_documents( + images=[image], + texts=[passage_text], + ) + else: + embedding = hf_model.model.encode_documents( + texts=[passage_text] + ) else: - # Text-only document - embedding = hf_model.model.encode_documents( - texts=[passage_text] + raise ValueError( + f"Text must start with {QUERY_PREFIX!r} " + f"or {PASSAGE_PREFIX!r}" ) - else: - raise ValueError( - f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'" - ) - - hf_outputs.append(embedding[0].tolist()) - - check_embeddings_close( - embeddings_0_lst=hf_outputs, - embeddings_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -def test_models_text( - hf_runner, - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - """Test text-only embedding.""" - input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] + hf_outputs.append(embedding[0].tolist()) + hf_outputs_per_case.append(hf_outputs) - _run_test( - hf_runner, - vllm_runner, - input_texts, - input_images, # type: ignore - model, - dtype=dtype, - ) + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", + ) @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["half"]) -def test_models_image( +def test_models( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - """Test image embedding.""" - input_texts_images = [ - (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) + """Test text and image embedding.""" + text_images = [None] * len(HF_TEXT_PROMPTS) + images = [asset.pil_image for asset in image_assets] + input_cases = [ + (HF_TEXT_PROMPTS, text_images), + (HF_IMAGE_PROMPTS, images), ] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] _run_test( hf_runner, vllm_runner, - input_texts, - input_images, + input_cases, # type: ignore[arg-type] model, dtype=dtype, ) @@ -188,8 +167,11 @@ def test_models_image( RERANKER_IMAGE_QUERY = "photo of a red stop sign on a street" +RerankerDocument = tuple[str | None, Image.Image | None] +RerankerCase = tuple[str, Sequence[RerankerDocument]] + -def _pil_to_data_uri(image) -> str: +def _pil_to_data_uri(image: Image.Image) -> str: buf = BytesIO() image.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() @@ -200,10 +182,9 @@ def _run_hf_reranker( hf_runner: type[HfRunner], model: str, dtype: str, - query: str, - docs: list, -) -> list[float]: - """Run HF reranker inference; docs is a list of (doc_text, doc_image|None).""" + input_cases: Sequence[RerankerCase], +) -> list[list[float]]: + """Run all HF reranker cases in one model lifecycle.""" with hf_runner( model, dtype=dtype, @@ -217,35 +198,37 @@ def _run_hf_reranker( use_thumbnail=True, rerank_max_length=2048, ) - examples = [ - { - "question": query, - "doc_text": doc_text if doc_text is not None else "", - "doc_image": doc_image if doc_image is not None else "", + scores_per_case = [] + for query, docs in input_cases: + examples = [ + { + "question": query, + "doc_text": doc_text if doc_text is not None else "", + "doc_image": doc_image if doc_image is not None else "", + } + for doc_text, doc_image in docs + ] + batch_dict = processor.process_queries_documents_crossencoder(examples) + batch_dict = { + k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v + for k, v in batch_dict.items() } - for doc_text, doc_image in docs - ] - batch_dict = processor.process_queries_documents_crossencoder(examples) - batch_dict = { - k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v - for k, v in batch_dict.items() - } - with torch.inference_mode(): - logits = hf_model.model(**batch_dict, return_dict=True).logits - # vLLM applies sigmoid activation to the raw logits before returning - # scores; apply the same here so both sides are comparable. - scores = torch.sigmoid(logits.squeeze(-1).float()) - return scores.detach().cpu().tolist() + with torch.inference_mode(): + logits = hf_model.model(**batch_dict, return_dict=True).logits + # vLLM applies sigmoid activation to raw logits before returning scores. + scores = torch.sigmoid(logits.squeeze(-1).float()) + scores_per_case.append(scores.detach().cpu().tolist()) + + return scores_per_case def _run_vllm_reranker( vllm_runner: type[VllmRunner], model: str, dtype: str, - query: str, - docs: list, -) -> list[float]: - """Run vLLM reranker inference; docs is a list of (doc_text, doc_image|None).""" + input_cases: Sequence[RerankerCase], +) -> list[list[float]]: + """Run all vLLM reranker cases in one model lifecycle.""" with vllm_runner( model, runner="pooling", @@ -255,57 +238,59 @@ def _run_vllm_reranker( trust_remote_code=True, **ROCM_ENGINE_KWARGS, ) as vllm_model: - has_images = any(img is not None for _, img in docs) - - if not has_images: - # Text-only path: use the simple string score API. - queries = [query] * len(docs) - doc_texts = [doc_text for doc_text, _ in docs] - outputs = vllm_model.score( - queries, - doc_texts, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) - else: - # Multimodal path: build ScoreMultiModalParam for each pair. - query_params = [ - ScoreMultiModalParam( - content=[ - ChatCompletionContentPartTextParam( - type="text", - text=query, - ) - ] + scores_per_case = [] + for query, docs in input_cases: + has_images = any(img is not None for _, img in docs) + + if not has_images: + queries = [query] * len(docs) + doc_texts = [doc_text for doc_text, _ in docs] + outputs = vllm_model.score( + queries, + doc_texts, + chat_template=_RERANKER_SCORE_TEMPLATE, ) - ] * len(docs) - - doc_params = [] - for doc_text, doc_image in docs: - content: list = [] - if doc_image is not None: - content.append( - ChatCompletionContentPartImageParam( - type="image_url", - image_url={"url": _pil_to_data_uri(doc_image)}, - ) + else: + query_params = [ + ScoreMultiModalParam( + content=[ + ChatCompletionContentPartTextParam( + type="text", + text=query, + ) + ] ) - if doc_text: - content.append( - ChatCompletionContentPartTextParam( - type="text", - text=doc_text, + ] * len(docs) + + doc_params = [] + for doc_text, doc_image in docs: + content: list = [] + if doc_image is not None: + content.append( + ChatCompletionContentPartImageParam( + type="image_url", + image_url={"url": _pil_to_data_uri(doc_image)}, + ) ) - ) - doc_params.append(ScoreMultiModalParam(content=content)) + if doc_text: + content.append( + ChatCompletionContentPartTextParam( + type="text", + text=doc_text, + ) + ) + doc_params.append(ScoreMultiModalParam(content=content)) - raw_outputs = vllm_model.llm.score( - query_params, - doc_params, - chat_template=_RERANKER_SCORE_TEMPLATE, - ) - outputs = [o.outputs.score for o in raw_outputs] + raw_outputs = vllm_model.llm.score( + query_params, + doc_params, + chat_template=_RERANKER_SCORE_TEMPLATE, + ) + outputs = [output.outputs.score for output in raw_outputs] + + scores_per_case.append(outputs) - return outputs + return scores_per_case def _run_reranker_test( @@ -313,50 +298,44 @@ def _run_reranker_test( vllm_runner: type[VllmRunner], model: str, dtype: str, - query: str, - docs: list, + input_cases: Sequence[RerankerCase], ) -> None: - """Compare HF and vLLM reranker scores. + """Compare HF and vLLM reranker scores for all input cases. NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing. """ - vllm_scores = _run_vllm_reranker(vllm_runner, model, dtype, query, docs) - hf_scores = _run_hf_reranker(hf_runner, model, dtype, query, docs) + vllm_scores_per_case = _run_vllm_reranker(vllm_runner, model, dtype, input_cases) + hf_scores_per_case = _run_hf_reranker(hf_runner, model, dtype, input_cases) - assert len(hf_scores) == len(vllm_scores), ( - f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}" - ) - # NOTE: ROCm shows slightly higher numerical variance dues to different attention - # backend between vLLM and HF; use a marginally looser tolerance + # ROCm has slightly higher variance because vLLM and HF use different + # attention backends. rel_tol = 0.022 if current_platform.is_rocm() else 0.02 - for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)): - assert hf_score == pytest.approx(vllm_score, rel=rel_tol), ( - f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}" + for hf_scores, vllm_scores in zip(hf_scores_per_case, vllm_scores_per_case): + assert len(hf_scores) == len(vllm_scores), ( + f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}" ) + for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)): + assert hf_score == pytest.approx(vllm_score, rel=rel_tol), ( + f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}" + ) @pytest.mark.parametrize("model", RERANKER_MODELS) @pytest.mark.parametrize("dtype", ["half"]) -def test_reranker_text( - hf_runner, - vllm_runner, - model: str, - dtype: str, -) -> None: - """Test reranking with text-only query and text documents.""" - docs = [(text, None) for text in RERANKER_TEXT_DOCS] - _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_TEXT_QUERY, docs) - - -@pytest.mark.parametrize("model", RERANKER_MODELS) -@pytest.mark.parametrize("dtype", ["half"]) -def test_reranker_image_doc( +def test_reranker( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - """Test reranking with text query against image documents.""" - docs = [(None, asset.pil_image) for asset in image_assets] - _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_IMAGE_QUERY, docs) + """Test reranking with text and image documents.""" + text_docs: list[RerankerDocument] = [(text, None) for text in RERANKER_TEXT_DOCS] + image_docs: list[RerankerDocument] = [ + (None, asset.pil_image) for asset in image_assets + ] + input_cases: list[RerankerCase] = [ + (RERANKER_TEXT_QUERY, text_docs), + (RERANKER_IMAGE_QUERY, image_docs), + ] + _run_reranker_test(hf_runner, vllm_runner, model, dtype, input_cases) diff --git a/tests/models/multimodal/pooling/test_siglip.py b/tests/models/multimodal/pooling/test_siglip.py index bca598b42c64..8eeb594db226 100644 --- a/tests/models/multimodal/pooling/test_siglip.py +++ b/tests/models/multimodal/pooling/test_siglip.py @@ -33,16 +33,11 @@ def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], - input_texts: list[str], - input_images: PromptImageInput, + input_cases: list[tuple[list[str], PromptImageInput, dict[str, Any]]], model: str, *, dtype: str, - tokenization_kwargs: dict[str, Any] | None = None, ) -> None: - if tokenization_kwargs is None: - tokenization_kwargs = {} - with vllm_runner( model, runner="pooling", @@ -51,116 +46,88 @@ def _run_test( max_model_len=64, gpu_memory_utilization=0.7, ) as vllm_model: - vllm_outputs = vllm_model.embed( - input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs - ) + vllm_outputs_per_case = [ + vllm_model.embed( + input_texts, + images=input_images, + tokenization_kwargs=tokenization_kwargs, + ) + for input_texts, input_images, tokenization_kwargs in input_cases + ] + + texts = [HF_TEXT_PROMPTS[0]] + images = [input_cases[1][1][0]] + with pytest.raises(ValueError, match="not both"): + vllm_model.embed(texts, images=images) + + vllm_model.embed(texts) + vllm_model.embed([""], images=images) with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model: - all_inputs = hf_model.get_inputs( - input_texts, images=input_images, tokenization_kwargs=tokenization_kwargs + hf_outputs_per_case = [] + for input_texts, input_images, tokenization_kwargs in input_cases: + all_inputs = hf_model.get_inputs( + input_texts, + images=input_images, + tokenization_kwargs=tokenization_kwargs, + ) + + hf_outputs = [] + for inputs in all_inputs: + inputs = hf_model.wrap_device(inputs) + + if "pixel_values" in inputs: + pooled_output = hf_model.model.get_image_features( + pixel_values=inputs.pixel_values, + ) + else: + pooled_output = hf_model.model.get_text_features( + input_ids=inputs.input_ids, + ) + + if not isinstance(pooled_output, torch.Tensor): + pooled_output = pooled_output.pooler_output + pooled_output = pooled_output.squeeze(0) + hf_outputs.append(pooled_output.tolist()) + + hf_outputs_per_case.append(hf_outputs) + + for hf_outputs, vllm_outputs in zip(hf_outputs_per_case, vllm_outputs_per_case): + check_embeddings_close( + embeddings_0_lst=hf_outputs, + embeddings_1_lst=vllm_outputs, + name_0="hf", + name_1="vllm", ) - all_outputs = [] - for inputs in all_inputs: - inputs = hf_model.wrap_device(inputs) - - if "pixel_values" in inputs: - pooled_output = hf_model.model.get_image_features( - pixel_values=inputs.pixel_values, - ) - else: - pooled_output = hf_model.model.get_text_features( - input_ids=inputs.input_ids, - ) - - if not isinstance(pooled_output, torch.Tensor): - pooled_output = pooled_output.pooler_output - pooled_output = pooled_output.squeeze(0) - all_outputs.append(pooled_output.tolist()) - - hf_outputs = all_outputs - - check_embeddings_close( - embeddings_0_lst=hf_outputs, - embeddings_1_lst=vllm_outputs, - name_0="hf", - name_1="vllm", - ) - @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", ["float"]) -def test_models_text( +def test_models( hf_runner, vllm_runner, image_assets, model: str, dtype: str, ) -> None: - input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] - - _run_test( - hf_runner, - vllm_runner, - input_texts, - input_images, # type: ignore - model, - dtype=dtype, - tokenization_kwargs={ - "padding": "max_length", - "max_length": 64, - }, # siglip2 was trained with this padding setting. - ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_image( - hf_runner, - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - input_texts_images = [ - (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets) + text_images = [None] * len(HF_TEXT_PROMPTS) + images = [asset.pil_image for asset in image_assets] + input_cases = [ + ( + HF_TEXT_PROMPTS, + text_images, + { + "padding": "max_length", + "max_length": 64, + }, + ), + (HF_IMAGE_PROMPTS, images, {}), ] - input_texts = [text for text, _ in input_texts_images] - input_images = [image for _, image in input_texts_images] _run_test( hf_runner, vllm_runner, - input_texts, - input_images, + input_cases, # type: ignore[arg-type] model, dtype=dtype, ) - - -@pytest.mark.parametrize("model", MODELS) -@pytest.mark.parametrize("dtype", ["float"]) -def test_models_text_image_no_crash( - vllm_runner, - image_assets, - model: str, - dtype: str, -) -> None: - texts = [HF_TEXT_PROMPTS[0]] - images = [image_assets[0].pil_image] - - with vllm_runner( - model, - runner="pooling", - dtype=dtype, - enforce_eager=True, - max_model_len=64, - gpu_memory_utilization=0.7, - ) as vllm_model: - with pytest.raises(ValueError, match="not both"): - vllm_model.embed(texts, images=images) - - vllm_model.embed(texts) - vllm_model.embed([""], images=images) diff --git a/tests/models/quantization/test_awq.py b/tests/models/quantization/test_awq.py index 25a63f6bd907..38f60be6bfca 100644 --- a/tests/models/quantization/test_awq.py +++ b/tests/models/quantization/test_awq.py @@ -17,6 +17,15 @@ } ) +IMAGE_SIZE_FACTOR_GROUPS = ( + # Single-scale + (1.0,), + # Single-scale, batched + (1.0, 1.0, 1.0), + # Multi-scale + (0.25, 0.5, 1.0), +) + def run_awq_test( vllm_runner: type[VllmRunner], @@ -24,7 +33,7 @@ def run_awq_test( source_model: str, quant_model: str, *, - size_factors: list[float], + size_factor_groups: tuple[tuple[float, ...], ...], dtype: str, max_tokens: int, num_logprobs: int, @@ -33,11 +42,12 @@ def run_awq_test( ): images = [asset.pil_image for asset in image_assets] - inputs_per_image = [ + inputs_per_image_and_size_group = [ ( [prompt for _ in size_factors], [rescale_image_size(image, factor) for factor in size_factors], ) + for size_factors in size_factor_groups for image, prompt in zip(images, HF_IMAGE_PROMPTS) ] @@ -60,7 +70,7 @@ def run_awq_test( vllm_model.generate_greedy_logprobs( prompts, max_tokens, num_logprobs=num_logprobs, images=images ) - for prompts, images in inputs_per_image + for prompts, images in inputs_per_image_and_size_group ] with vllm_runner( @@ -77,7 +87,7 @@ def run_awq_test( vllm_model.generate_greedy_logprobs( prompts, max_tokens, num_logprobs=num_logprobs, images=images ) - for prompts, images in inputs_per_image + for prompts, images in inputs_per_image_and_size_group ] for source_outputs, quant_outputs in zip( @@ -128,17 +138,6 @@ def test_awq_load( ("source_model", "quant_model"), [("OpenGVLab/InternVL2-2B", "OpenGVLab/InternVL2-2B-AWQ")], ) -@pytest.mark.parametrize( - "size_factors", - [ - # Single-scale - [1.0], - # Single-scale, batched - [1.0, 1.0, 1.0], - # Multi-scale - [0.25, 0.5, 1.0], - ], -) @pytest.mark.parametrize("dtype", ["half"]) @pytest.mark.parametrize("max_tokens", [128]) @pytest.mark.parametrize("num_logprobs", [5]) @@ -148,7 +147,6 @@ def test_awq_models( image_assets, source_model, quant_model, - size_factors, dtype, max_tokens, num_logprobs, @@ -158,7 +156,7 @@ def test_awq_models( image_assets, source_model, quant_model, - size_factors=size_factors, + size_factor_groups=IMAGE_SIZE_FACTOR_GROUPS, dtype=dtype, max_tokens=max_tokens, num_logprobs=num_logprobs, diff --git a/tests/models/registry.py b/tests/models/registry.py index ed1786fdab3f..16604ec7ff83 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -825,7 +825,14 @@ def check_available_online( "Cosmos3EdgeForConditionalGeneration": _HfExamplesInfo( "nvidia/Cosmos3-Edge", max_model_len=4096, - is_available_online=False, + min_transformers_version="5.15", + use_original_num_layers=True, + hf_overrides={ + "text_config": { + "num_hidden_layers": 2, + "hybrid_override_pattern": "*-", + } + }, ), "DeepseekVLV2ForCausalLM": _HfExamplesInfo( "deepseek-ai/deepseek-vl2-tiny", @@ -840,8 +847,7 @@ def check_available_online( "deepseek-ai/DeepSeek-OCR-2", ), "Dots3NoteForCausalLM": _HfExamplesInfo( - "rednote-hilab/dots3.note", - trust_remote_code=True, + "dots-studio/dots3-note-prev", is_available_online=False, ), "UnlimitedOCRForCausalLM": _HfExamplesInfo( @@ -1192,6 +1198,12 @@ def check_available_online( # required by current PrefixLM implementation max_num_batched_tokens=31872, ), + "MuseGlimmerForConditionalGeneration": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + ), + "MuseGlimmerForCausalLM": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + ), "NVLM_D": _HfExamplesInfo("nvidia/NVLM-D-72B", trust_remote_code=True), "Llama_Nemotron_Nano_VL": _HfExamplesInfo( "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1", @@ -1489,6 +1501,18 @@ def check_available_online( max_num_seqs=32, min_transformers_version="4.56.3", # Required for Qwen3Next ), + "MuseGlimmerAssistantModel": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + speculative_model="meta-models/Muse-Glimmer-30B-assistant", + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env + max_num_seqs=32, + ), + "DFlashMuseGlimmerAssistantModel": _HfExamplesInfo( + "meta-models/Muse-Glimmer-30B", + speculative_model="meta-models/Muse-Glimmer-30B-assistant", + max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env + max_num_seqs=32, + ), # [DSpark] "DSparkDraftModel": _HfExamplesInfo( "deepseek-ai/DeepSeek-V4-Pro-DSpark", @@ -1656,9 +1680,8 @@ def check_available_online( is_available_online=False, ), "Dots3NoteMTPModel": _HfExamplesInfo( - "rednote-hilab/dots3.note", - speculative_model="rednote-hilab/dots3.note", - trust_remote_code=True, + "dots-studio/dots3-note-prev", + speculative_model="dots-studio/dots3-note-prev", is_available_online=False, ), "Gemma4MTPModel": _HfExamplesInfo( diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 6eda6075d2f8..70b8b18f76f6 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -143,8 +143,10 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda): @pytest.mark.parametrize( "model_arch,supported", [ - # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today. + # ReplaySSM is opt-in per model. ("NemotronHForCausalLM", True), + ("KimiLinearForCausalLM", not current_platform.is_rocm()), + ("KimiK3ForConditionalGeneration", not current_platform.is_rocm()), ("Mamba2ForCausalLM", False), ("Zamba2ForCausalLM", False), ], diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index 4b6b2796af64..5fd39a5b9d54 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -263,7 +263,16 @@ def test_embed_loading(vllm_runner, model): @pytest.mark.parametrize( "arch", ["TransformersEmbeddingModel", "TransformersForSequenceClassification"] ) -def test_pooling(hf_runner, vllm_runner, example_prompts, arch): +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_pooling( + hf_runner, + vllm_runner, + example_prompts, + arch, + monkeypatch, + use_v2_model_runner, +): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) model = get_model(arch) vllm_kwargs = dict(max_model_len=None, model_impl="transformers") diff --git a/tests/models/utils.py b/tests/models/utils.py index f938a702231d..6bd543182997 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -564,15 +564,20 @@ def _supports_multimodal_for_mm_prefix() -> bool: ) if hasattr(hf_config, "vision_config"): - hf_config.vision_config.update( + vision_config = hf_config.vision_config + vision_config.update( { "num_layers": 1, "num_hidden_layers": 1, } ) + # Keep per-layer metadata consistent with the reduced layer count. + if layer_types := getattr(vision_config, "layer_types", None): + vision_config.update({"layer_types": layer_types[:1]}) + if model_arch in ("Moondream3ForCausalLM", "HfMoondream"): - hf_config.vision_config.update({"enc_n_layers": 1}) + vision_config.update({"enc_n_layers": 1}) # e.g.: ibm-granite/granite-speech-3.3-2b if hasattr(hf_config, "encoder_config"): diff --git a/tests/multimodal/media/test_unprocessable_entity_error.py b/tests/multimodal/media/test_unprocessable_entity_error.py index 7cad42955755..bca2ac5f5e25 100644 --- a/tests/multimodal/media/test_unprocessable_entity_error.py +++ b/tests/multimodal/media/test_unprocessable_entity_error.py @@ -13,7 +13,7 @@ import aiohttp import pytest -from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve import create_error_response from vllm.exceptions import VLLMClientError, VLLMUnprocessableEntityError from vllm.multimodal.media import MediaConnector diff --git a/tests/multimodal/test_video.py b/tests/multimodal/test_video.py index 8e27a5c58f88..17d29f850ea4 100644 --- a/tests/multimodal/test_video.py +++ b/tests/multimodal/test_video.py @@ -24,12 +24,14 @@ Molmo2VideoBackend, PyNvVideoCodecDecoderSlot, PyNvVideoCodecVideoBackend, + PyNvVideoCodecVideoBackendMixin, Qwen2VLVideoBackend, Qwen3VLVideoBackend, VideoBackend, VideoLoader, VideoSourceMetadata, VideoTargetMetadata, + _pynv_decoder_pool, get_video_loader_backend_for_processor, ) from vllm.platforms import current_platform @@ -47,6 +49,27 @@ FAKE_OUTPUT_2 = np.random.rand(NUM_FRAMES, 1280, 720, 3) +@contextmanager +def _fresh_decoder_pool(): + """Reset module-level decoder pool for isolated test runs.""" + pool = _pynv_decoder_pool + old_slots = pool.slots + old_active = pool.active + old_cond = pool.cond + old_max = pool.max_slots + pool.slots = [] + pool.active = 0 + pool.cond = threading.Condition() + pool.max_slots = None + try: + yield pool + finally: + pool.slots = old_slots + pool.active = old_active + pool.cond = old_cond + pool.max_slots = old_max + + @VIDEO_LOADER_REGISTRY.register("test_video_loader_1") class TestVideoLoader1(VideoLoader): @classmethod @@ -206,16 +229,7 @@ def test_pynvvideocodec_corrupted_videos_raise_value_error(): corrupted_video = (ASSETS_DIR / "corrupted.mp4").read_bytes() malformed_video = corrupted_video[:128] - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots - try: - PyNvVideoCodecVideoBackend._decoder_slots = [] - PyNvVideoCodecVideoBackend._active_decoder_slots = 0 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = None - + with _fresh_decoder_pool(): loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND) with pytest.raises( ValueError, @@ -247,11 +261,6 @@ def test_pynvvideocodec_corrupted_videos_raise_value_error(): hw_decoders=1, ) assert frames.shape[0] == 1 - finally: - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots @pytest.mark.parametrize("hw_decoders", [1, 3]) @@ -263,15 +272,7 @@ class FakeSlot: pass create_count = 0 - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots - try: - PyNvVideoCodecVideoBackend._decoder_slots = [] - PyNvVideoCodecVideoBackend._active_decoder_slots = 0 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = None + with _fresh_decoder_pool(): PyNvVideoCodecVideoBackend._configure_decoder_slots(hw_decoders) def fake_create_slot(cls): @@ -309,17 +310,12 @@ def borrow_extra_slot(): assert seen_slots[0] in retained_slots assert create_count == hw_decoders - finally: - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots def test_pynvvideocodec_decoder_slots_are_configured_once( monkeypatch: pytest.MonkeyPatch, ): - monkeypatch.setattr(PyNvVideoCodecVideoBackend, "_max_decoder_slots", None) + monkeypatch.setattr(_pynv_decoder_pool, "max_slots", None) PyNvVideoCodecVideoBackend._configure_decoder_slots(2) PyNvVideoCodecVideoBackend._configure_decoder_slots(2) @@ -358,15 +354,16 @@ def SimpleDecoder(file_path: str, **kwargs): assert slot.source_path is None raise RuntimeError("construct failed") - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots + pool = _pynv_decoder_pool + old_slots = pool.slots + old_active = pool.active + old_cond = pool.cond + old_max = pool.max_slots try: - PyNvVideoCodecVideoBackend._decoder_slots = [slot] - PyNvVideoCodecVideoBackend._active_decoder_slots = 1 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = 1 + pool.slots = [slot] + pool.active = 1 + pool.cond = threading.Condition() + pool.max_slots = 1 with ( pytest.raises(RuntimeError, match="construct failed"), @@ -386,12 +383,12 @@ def SimpleDecoder(file_path: str, **kwargs): assert old_decoder.poisoned assert slot.decoder is None assert slot.source_path is None - assert PyNvVideoCodecVideoBackend._decoder_slots == [slot] + assert pool.slots == [slot] finally: - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots + pool.slots = old_slots + pool.active = old_active + pool.cond = old_cond + pool.max_slots = old_max @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @@ -405,15 +402,15 @@ def test_pynvvideocodec_h200_recovers_after_unsupported_8k(): valid_video = create_long_gop_video(num_frames=2, width=64, height=64) unsupported_video = (ASSETS_DIR / "unsupported_8k_h264.mp4").read_bytes() - old_slots = PyNvVideoCodecVideoBackend._decoder_slots - old_active_slots = PyNvVideoCodecVideoBackend._active_decoder_slots - old_cond = PyNvVideoCodecVideoBackend._decoder_slot_cond - old_max_slots = PyNvVideoCodecVideoBackend._max_decoder_slots + old_slots = _pynv_decoder_pool.slots + old_active = _pynv_decoder_pool.active + old_cond = _pynv_decoder_pool.cond + old_max = _pynv_decoder_pool.max_slots try: - PyNvVideoCodecVideoBackend._decoder_slots = [] - PyNvVideoCodecVideoBackend._active_decoder_slots = 0 - PyNvVideoCodecVideoBackend._decoder_slot_cond = threading.Condition() - PyNvVideoCodecVideoBackend._max_decoder_slots = None + _pynv_decoder_pool.slots = [] + _pynv_decoder_pool.active = 0 + _pynv_decoder_pool.cond = threading.Condition() + _pynv_decoder_pool.max_slots = None loader = VIDEO_LOADER_REGISTRY.load(PYNVVIDEOCODEC_VIDEO_BACKEND) frames_before, _ = loader.load_bytes( @@ -443,12 +440,66 @@ def test_pynvvideocodec_h200_recovers_after_unsupported_8k(): assert frames_after.shape == frames_before.shape finally: - for slot in PyNvVideoCodecVideoBackend._decoder_slots: + for slot in _pynv_decoder_pool.slots: slot.invalidate() - PyNvVideoCodecVideoBackend._decoder_slots = old_slots - PyNvVideoCodecVideoBackend._active_decoder_slots = old_active_slots - PyNvVideoCodecVideoBackend._decoder_slot_cond = old_cond - PyNvVideoCodecVideoBackend._max_decoder_slots = old_max_slots + _pynv_decoder_pool.slots = old_slots + _pynv_decoder_pool.active = old_active + _pynv_decoder_pool.cond = old_cond + _pynv_decoder_pool.max_slots = old_max + + +def test_pynvvideocodec_cross_subclass_shares_single_pool(): + """Regression test for GHSA-j682-9xp5-rrf3. + + Multiple subclasses of PyNvVideoCodecVideoBackendMixin must share the + same process-wide decoder slot limit rather than getting independent + counters via ClassVar shadowing. + """ + + class FakeSlot: + pass + + create_count = 0 + + def fake_create_slot(cls): + nonlocal create_count + create_count += 1 + return FakeSlot() + + with _fresh_decoder_pool() as pool: + pool.max_slots = 2 + + orig_create = PyNvVideoCodecVideoBackendMixin._create_decoder_slot + PyNvVideoCodecVideoBackendMixin._create_decoder_slot = classmethod( + fake_create_slot + ) + try: + with ExitStack() as stack: + stack.enter_context(VideoBackend._borrow_decoder_slot()) + stack.enter_context(Qwen3VLVideoBackend._borrow_decoder_slot()) + assert pool.active == 2 + + blocked = threading.Event() + acquired = threading.Event() + + def try_borrow(): + blocked.set() + with Qwen2VLVideoBackend._borrow_decoder_slot(): + acquired.set() + + t = threading.Thread(target=try_borrow) + t.start() + blocked.wait(timeout=2.0) + assert not acquired.wait(timeout=0.3) + + assert acquired.wait(timeout=2.0) + t.join(timeout=2.0) + assert not t.is_alive() + + assert create_count == 2 + assert len(pool.slots) == 2 + finally: + PyNvVideoCodecVideoBackendMixin._create_decoder_slot = orig_create @pytest.mark.parametrize("hw_decoders", [0, -1, 1.5, True, "2"]) diff --git a/tests/parser/engine/test_engine.py b/tests/parser/engine/test_engine.py index 6da4c428e028..07f545571885 100644 --- a/tests/parser/engine/test_engine.py +++ b/tests/parser/engine/test_engine.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for the streaming parser engine core pipeline.""" +from dataclasses import replace from unittest.mock import MagicMock import pytest @@ -71,6 +72,16 @@ def _think_config() -> ParserEngineConfig: ) +def _token_think_config() -> ParserEngineConfig: + return replace( + _think_config(), + token_id_terminals={ + "THINK_START": "", + "THINK_END": "", + }, + ) + + class TestNonStreaming: def test_plain_text(self): engine = StreamingParserEngine(_hermes_config(), tokenizer=None) @@ -441,6 +452,46 @@ def test_mixed_text_then_real_tool_call(self): assert sum(1 for e in all_events if e.type == EventType.TOOL_CALL_END) == 1 +class TestReasoningTokenCounts: + def test_counts_reasoning_and_excludes_boundaries_and_content(self): + engine = StreamingParserEngine(_token_think_config(), _make_think_tokenizer()) + + events = engine.feed( + "tok1tok2tok3", + [_START_ID, 1, 2, _END_ID, 3], + ) + events.extend(engine.finish()) + + reasoning = [e for e in events if e.type == EventType.REASONING_CHUNK] + assert "".join(e.value for e in reasoning) == "tok1tok2" + assert sum(e.token_count for e in reasoning) == 2 + assert engine.reasoning_token_count == 2 + + def test_counts_tokens_after_final_deferred_start_terminal(self): + engine = StreamingParserEngine(_token_think_config(), _make_think_tokenizer()) + + assert engine.feed("", [_START_ID, 1, 2]) == [] + events = engine.feed("tok1tok2", []) + + reasoning = [e for e in events if e.type == EventType.REASONING_CHUNK] + assert "".join(e.value for e in reasoning) == "tok1tok2" + assert sum(e.token_count for e in reasoning) == 2 + assert engine.reasoning_token_count == 2 + + def test_deferred_reasoning_tokens_stay_before_next_end_terminal(self): + engine = StreamingParserEngine(_token_think_config(), _make_think_tokenizer()) + + assert engine.feed("", [_START_ID, 1, 2]) == [] + events = engine.feed("tok1tok2tok3", [_END_ID, 3]) + + assert ( + sum(e.token_count for e in events if e.type == EventType.REASONING_CHUNK) + == 2 + ) + assert sum(e.token_count for e in events if e.type == EventType.TEXT_CHUNK) == 1 + assert engine.reasoning_token_count == 2 + + def _func_prefix_config() -> ParserEngineConfig: """Config mixing token-ID terminals (TOOL_START/END) with text-only terminals (FUNC_PREFIX) and fallback transitions.""" diff --git a/tests/parser/engine/test_parser_engine.py b/tests/parser/engine/test_parser_engine.py index 8616aac381af..37c9e9734fee 100644 --- a/tests/parser/engine/test_parser_engine.py +++ b/tests/parser/engine/test_parser_engine.py @@ -33,6 +33,7 @@ ParserState, Transition, ) +from vllm.parser.parser_manager import ParserManager # ── Shared test configs ────────────────────────────────────────────── @@ -892,6 +893,129 @@ class _CombinedDelegating(DelegatingParser): tool_parser_cls = _CombinedToolAdapter +def test_parser_manager_uses_shared_engine_directly(monkeypatch): + monkeypatch.setattr( + ParserManager, + "get_reasoning_parser", + classmethod(lambda cls, name: _CombinedReasoningAdapter), + ) + monkeypatch.setattr( + ParserManager, + "get_tool_parser", + classmethod(lambda cls, name, enabled, model: _CombinedToolAdapter), + ) + + parser_cls = ParserManager.get_parser( + tool_parser_name="combined", + reasoning_parser_name="combined", + enable_auto_tools=True, + ) + + assert parser_cls is not None + assert parser_cls is _CombinedTestEngine + parser = parser_cls(make_mock_tokenizer(_VOCAB)) + request = _make_delegating_request() + reasoning, content, _ = parser.parse( + "abc", + request, + model_output_token_ids=[ord("a"), ord("b"), 201, ord("c")], + ) + assert reasoning == "ab" + assert content == "c" + assert parser.count_reasoning_tokens([]) == 2 + + +def test_parser_manager_preserves_reasoning_only_adapter(monkeypatch): + monkeypatch.setattr( + ParserManager, + "get_reasoning_parser", + classmethod(lambda cls, name: _CombinedReasoningAdapter), + ) + monkeypatch.setattr( + ParserManager, + "get_tool_parser", + classmethod(lambda cls, name, enabled, model: None), + ) + + parser_cls = ParserManager.get_parser(reasoning_parser_name="combined") + + assert parser_cls is not None + parser = parser_cls(make_mock_tokenizer(_VOCAB)) + assert parser.reasoning_parser is not None + assert parser.tool_parser is None + reasoning, content, _ = parser.parse( + 'ab{"name":"h","arguments":{}}', + _make_delegating_request(), + model_output_token_ids=[ord("a"), ord("b"), 201], + ) + assert reasoning == "ab" + assert content == '{"name":"h","arguments":{}}' + assert parser.count_reasoning_tokens([ord("a"), ord("b"), 201]) == 2 + + +def test_parser_manager_preserves_tool_only_adapter(monkeypatch): + monkeypatch.setattr( + ParserManager, + "get_reasoning_parser", + classmethod(lambda cls, name: None), + ) + monkeypatch.setattr( + ParserManager, + "get_tool_parser", + classmethod(lambda cls, name, enabled, model: _CombinedToolAdapter), + ) + + parser_cls = ParserManager.get_parser( + tool_parser_name="combined", enable_auto_tools=True + ) + + assert parser_cls is not None + parser = parser_cls(make_mock_tokenizer(_VOCAB)) + assert parser.reasoning_parser is None + assert parser.tool_parser is not None + reasoning, content, tool_calls = parser.parse( + "abc", + _make_delegating_request(), + model_output_token_ids=[200, ord("a"), ord("b"), 201, ord("c")], + ) + assert reasoning is None + assert content == "abc" + assert tool_calls == [] + + +def test_parser_manager_rejects_non_engine_adapter_metadata(): + class TraditionalParser: + pass + + class InvalidEngineAdapter: + _parser_engine_cls = TraditionalParser + + assert ParserManager._get_parser_engine_cls(TraditionalParser) is None + assert ParserManager._get_parser_engine_cls(InvalidEngineAdapter) is None + assert ( + ParserManager._get_parser_engine_cls(_CombinedReasoningAdapter) + is _CombinedTestEngine + ) + + +def test_reasoning_adapter_counts_after_final_non_streaming_parse(): + parser = _CombinedReasoningAdapter(make_mock_tokenizer(_VOCAB)) + request = _make_delegating_request() + token_ids = [ord("a"), ord("b"), 201, ord("c")] + + parser.extract_reasoning_streaming( + "", + "abc", + "abc", + [], + token_ids, + token_ids, + ) + parser.extract_reasoning("abc", request) + + assert parser.count_reasoning_tokens(token_ids) == 2 + + def _make_delegating_request(): req = MagicMock(spec=ChatCompletionRequest) req.tools = [] diff --git a/tests/parser/engine/test_token_id_scanner.py b/tests/parser/engine/test_token_id_scanner.py index 21537784306f..c37a46793aaa 100644 --- a/tests/parser/engine/test_token_id_scanner.py +++ b/tests/parser/engine/test_token_id_scanner.py @@ -163,9 +163,28 @@ def test_empty_delta_text_drops_individual_decode_text(self, tokenizer): assert len(result) == 0 flushed = scanner.flush_pending() - assert len(flushed) == 1 + assert len(flushed) == 2 assert isinstance(flushed[0], PreLexedTerminal) assert flushed[0].terminal == "TOOL_START" + assert isinstance(flushed[1], TextChunk) + assert flushed[1].text == "" + assert flushed[1].token_count == 2 + + def test_deferred_terminal_preserves_trailing_token_count(self, tokenizer): + tokenizer.decode.side_effect = lambda ids: { + CHANNEL_START_ID: CHANNEL_START, + 201: "alpha", + 202: "beta", + }[ids[0]] + scanner = TokenIDScanner({CHANNEL_START_ID: "THINK_START"}, tokenizer) + + assert scanner.scan("", [CHANNEL_START_ID, 201, 202]) == [] + + result = scanner.scan(f"{CHANNEL_START}alphabeta", []) + assert isinstance(result[0], PreLexedTerminal) + assert isinstance(result[1], TextChunk) + assert result[1].text == "alphabeta" + assert result[1].token_count == 2 def test_holdback_before_start_tag(self, scanner): result = scanner.scan( diff --git a/tests/parser/test_harmony.py b/tests/parser/test_harmony.py index 028bbed012dc..aed633cb3037 100644 --- a/tests/parser/test_harmony.py +++ b/tests/parser/test_harmony.py @@ -12,10 +12,9 @@ RenderConversationConfig, Role, ) -from transformers import AutoTokenizer -from xgrammar import Grammar -from xgrammar.testing import _is_grammar_accept_string +from transformers import AutoTokenizer, GenerationConfig +from vllm.config import StructuredOutputsConfig, VllmConfig from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.engine.protocol import FunctionCall from vllm.entrypoints.openai.parser.harmony_utils import ( @@ -25,6 +24,8 @@ from vllm.parser.harmony import HarmonyParser from vllm.parser.parser_manager import ParserManager from vllm.sampling_params import StructuredOutputsParams +from vllm.v1.structured_output.backend_types import StructuredOutputOptions +from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend REASONING_MODEL_NAME = "openai/gpt-oss-20b" @@ -34,6 +35,26 @@ def gpt_oss_tokenizer(): return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME) +@pytest.fixture(scope="module") +def gpt_oss_stop_token_ids() -> set[int]: + eos_token_id = GenerationConfig.from_pretrained(REASONING_MODEL_NAME).eos_token_id + if isinstance(eos_token_id, int): + return {eos_token_id} + return set(eos_token_id) + + +@pytest.fixture(scope="module") +def xgrammar_backend(gpt_oss_tokenizer) -> XgrammarBackend: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig(backend="xgrammar") + ) + return XgrammarBackend( + vllm_config, + tokenizer=gpt_oss_tokenizer, + vocab_size=len(gpt_oss_tokenizer), + ) + + @pytest.fixture def harmony_parser(gpt_oss_tokenizer): parser_cls = ParserManager.get_parser( @@ -1036,21 +1057,25 @@ def _assert_structured_outputs_admission( cls, adjusted_request: ChatCompletionRequest | ResponsesRequest, expected_admission: Sequence[str], + xgrammar_backend: XgrammarBackend, + stop_token_ids: set[int], ) -> None: structured_outputs = adjusted_request.structured_outputs assert structured_outputs is not None assert structured_outputs.structural_tag is not None assert structured_outputs.all_non_structural_tag_constraints_none() - grammar = Grammar.from_structural_tag(structured_outputs.structural_tag) + grammar = xgrammar_backend.compile_grammar( + StructuredOutputOptions.STRUCTURAL_TAG, + structured_outputs.structural_tag, + stop_token_ids=stop_token_ids, + ) expected_admission_set = set(expected_admission) for sample_name in cls.ADMISSION_SAMPLES: - admitted = _is_grammar_accept_string( - grammar, - getattr(cls, sample_name), - require_termination=False, - ) + tokens = encode_output(getattr(cls, sample_name)) + accepted = grammar.validate_tokens(tokens) + admitted = accepted == tokens should_admit = sample_name in expected_admission_set assert admitted is should_admit, ( f"Expected structured_outputs admission for {sample_name} " @@ -1183,6 +1208,8 @@ def _assert_structured_outputs_admission( def test_adjust_request( self, harmony_parser, + xgrammar_backend, + gpt_oss_stop_token_ids, request_kind, request_kwargs, expected_admission, @@ -1193,4 +1220,6 @@ def test_adjust_request( self._assert_structured_outputs_admission( adjusted_request, expected_admission, + xgrammar_backend, + gpt_oss_stop_token_ids, ) diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py index f4880abcb5f3..04a365b48dc1 100644 --- a/tests/quantization/test_turboquant.py +++ b/tests/quantization/test_turboquant.py @@ -281,7 +281,10 @@ class TestTurboQuantKVCacheSpec: @pytest.mark.parametrize("preset", ALL_PRESETS) def test_kv_cache_spec_sets_kv_quant_mode(self, preset): from vllm.model_executor.layers.attention.attention import Attention - from vllm.v1.kv_cache_interface import KVQuantMode, TQFullAttentionSpec + from vllm.v1.attention.backends.turboquant_attn import ( + TurboQuantAttentionBackend, + ) + from vllm.v1.kv_cache_interface import FullAttentionSpec layer = SimpleNamespace( attn_type="decoder", @@ -294,10 +297,16 @@ def test_kv_cache_spec_sets_kv_quant_mode(self, preset): ) vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32)) + # The layer builds an unpacked spec; the worker's spec-collection + # loop applies TQ slot packing via the backend's customize_spec hook. spec = Attention.get_kv_cache_spec(layer, vllm_config) + assert isinstance(spec, FullAttentionSpec) + assert spec.kv_quant_mode.is_turboquant + assert spec.state_content_bytes is None - assert isinstance(spec, TQFullAttentionSpec) - assert spec.kv_quant_mode == KVQuantMode.TURBOQUANT + spec = TurboQuantAttentionBackend.customize_spec(spec) + expected_slot = TurboQuantConfig.from_cache_dtype(preset, 128).slot_size_aligned + assert spec.state_content_bytes == expected_slot class TestTurboQuantWorkspaceReservation: @@ -333,15 +342,15 @@ def _fake_vllm_config( @staticmethod def _fake_kv_cache_spec(): - from vllm.v1.kv_cache_interface import TQFullAttentionSpec + from vllm.v1.kv_cache_interface import FullAttentionSpec - return TQFullAttentionSpec( + return FullAttentionSpec( block_size=32, num_kv_heads=4, head_size=128, head_size_v=128, dtype=torch.uint8, - tq_slot_size=102, + state_content_bytes=102, ) def test_metadata_builder_reserves_decode_and_continuation_prefill_workspace( diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py index 84bc40afe83e..0c1fc890092f 100644 --- a/tests/reasoning/test_cohere_command_reasoning_parser.py +++ b/tests/reasoning/test_cohere_command_reasoning_parser.py @@ -5,7 +5,7 @@ import json from collections import UserDict -from dataclasses import dataclass, field +from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -34,20 +34,12 @@ from vllm.sampling_params import StructuredOutputsParams -@dataclass -class ExpectedToolCall: - id: str - name: str - arguments: dict - - @dataclass class ReasoningCase: parser_cls: Any model_output: str expected_reasoning: str | None expected_content: str | None - expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) REASONING_CASES = [ @@ -67,9 +59,6 @@ class ReasoningCase: {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} ] <|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ], ), id="cmd3-single_tool_call", ), @@ -89,9 +78,6 @@ class ReasoningCase: {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} ] <|END_ACTION|>""", - expected_tool_calls=[ - ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), - ], ), id="cmd4-single_tool_call", ), @@ -236,29 +222,8 @@ def test_streaming(self, tokenizer, case: ReasoningCase): assert reasoning == case.expected_reasoning content = "".join(content_parts) if content_parts else None - if case.expected_tool_calls: - assert content is None or content == "" - else: - assert content == case.expected_content - - accumulated: dict[int, dict] = {} - for d in tool_call_deltas: - idx = d["index"] - if idx not in accumulated: - accumulated[idx] = {"id": "", "name": "", "arguments": ""} - if d["id"]: - accumulated[idx]["id"] = d["id"] - if d["name"]: - accumulated[idx]["name"] = d["name"] - if d["arguments"]: - accumulated[idx]["arguments"] += d["arguments"] - - assert len(accumulated) == len(case.expected_tool_calls) - for i, expected_tc in enumerate(case.expected_tool_calls): - tc = accumulated[i] - assert tc["id"] == expected_tc.id - assert tc["name"] == expected_tc.name - assert json.loads(tc["arguments"]) == expected_tc.arguments + assert content == case.expected_content + assert tool_call_deltas == [] class TestIsReasoningEnd: diff --git a/tests/renderers/test_process_multi_modal_uuids.py b/tests/renderers/test_process_multi_modal_uuids.py index 9f3440fc2e19..c811630c859f 100644 --- a/tests/renderers/test_process_multi_modal_uuids.py +++ b/tests/renderers/test_process_multi_modal_uuids.py @@ -8,7 +8,7 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset from vllm.config import CacheConfig, ModelConfig, VllmConfig -from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve import create_error_response from vllm.multimodal.parse import parse_mm_uuids from vllm.renderers.hf import HfRenderer from vllm.tokenizers.registry import cached_tokenizer_from_config diff --git a/tests/standalone_tests/python_only_compile.sh b/tests/standalone_tests/python_only_compile.sh index 14f563ae6fbb..7036dd4314f2 100644 --- a/tests/standalone_tests/python_only_compile.sh +++ b/tests/standalone_tests/python_only_compile.sh @@ -4,70 +4,160 @@ set -e -# ROCm CI runs this script inside `run-amd-test.sh` where /vllm-workspace often has no .git -# (wheel artifact layout). The wrapper passes CI_STANDALONE_MERGE_BASE from the agent checkout. merge_base_commit="" -if [[ -n "${CI_STANDALONE_MERGE_BASE:-}" ]]; then - merge_base_commit="${CI_STANDALONE_MERGE_BASE}" -elif merge_base_commit="$(git -C /vllm-workspace merge-base HEAD origin/main 2>/dev/null)"; then - : -elif merge_base_commit="$(git merge-base HEAD origin/main 2>/dev/null)"; then - : -else - echo "ERROR: need a git checkout or CI_STANDALONE_MERGE_BASE to resolve wheels.vllm.ai commit." >&2 - exit 1 +rocm_wheel="" +is_rocm=0 +_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')" +if [[ "${_vllm_target_lower}" == "rocm" || -n "${ROCM_PATH:-}" || -d /opt/rocm ]] \ + || command -v rocminfo >/dev/null 2>&1; then + is_rocm=1 fi +unset -v _vllm_target_lower -echo "INFO: current merge base commit with main: $merge_base_commit" -if git show --oneline -s "$merge_base_commit" 2>/dev/null; then - : -else - echo "INFO: git show unavailable in this environment; using SHA above for precompiled metadata." +if [[ "${is_rocm}" == "1" ]]; then + # Native CI passes the verified wheel artifact explicitly. Legacy ROCm + # images carry the same-build wheel in /opt/vllm-wheels. + if [[ -n "${VLLM_PRECOMPILED_WHEEL_LOCATION:-}" ]]; then + rocm_wheel="${VLLM_PRECOMPILED_WHEEL_LOCATION}" + if [[ ! -f "${rocm_wheel}" || "$(basename "${rocm_wheel}")" != vllm-*.whl ]]; then + echo "ERROR: invalid ROCm wheel location: ${rocm_wheel}" >&2 + exit 1 + fi + rocm_wheel="$(realpath -- "${rocm_wheel}")" + elif [[ -d /opt/vllm-wheels ]]; then + shopt -s nullglob + rocm_wheels=(/opt/vllm-wheels/vllm-*.whl) + shopt -u nullglob + if [[ "${#rocm_wheels[@]}" -ne 1 ]]; then + echo "ERROR: expected exactly one vLLM wheel in /opt/vllm-wheels, found ${#rocm_wheels[@]}." >&2 + exit 1 + fi + rocm_wheel="${rocm_wheels[0]}" + fi fi -# test whether the metadata.json url is valid, retry each 3 minutes up to 5 times -# this avoids cumbersome error messages & manual retries in case the precompiled wheel -# for the given commit is still being built in the release pipeline -meta_json_url="https://wheels.vllm.ai/$merge_base_commit/vllm/metadata.json" -echo "INFO: will use metadata.json from $meta_json_url" - -for i in {1..5}; do - echo "Checking metadata.json URL (attempt $i)..." - if curl --fail "$meta_json_url" > metadata.json; then - echo "INFO: metadata.json URL is valid." - # check whether it is valid json by python (printed to stdout) - if python3 -m json.tool metadata.json; then - echo "INFO: metadata.json is valid JSON. Proceeding with the check." - # check whether there is an object in the json matching: - # "package_name": "vllm", and "platform_tag" matches the current architecture - # see `determine_wheel_url` in setup.py for more details - if python3 -c "import platform as p,json as j,sys as s; d = j.load(open('metadata.json')); \ - s.exit(int(not any(o.get('package_name') == 'vllm' and p.machine() in o.get('platform_tag') \ - for o in d)))" 2>/dev/null; then - echo "INFO: metadata.json contains a pre-compiled wheel for the current architecture." - break - else - echo "WARN: metadata.json does not have a pre-compiled wheel for the current architecture." +if [[ -n "${rocm_wheel}" ]]; then + echo "INFO: using same-build ROCm wheel: ${rocm_wheel}" +else + # Some CI images do not include .git under /vllm-workspace. Their wrapper + # passes CI_STANDALONE_MERGE_BASE from the agent checkout. + if [[ -n "${CI_STANDALONE_MERGE_BASE:-}" ]]; then + merge_base_commit="${CI_STANDALONE_MERGE_BASE}" + elif merge_base_commit="$(git -C /vllm-workspace merge-base HEAD origin/main 2>/dev/null)"; then + : + elif merge_base_commit="$(git merge-base HEAD origin/main 2>/dev/null)"; then + : + else + echo "ERROR: need a git checkout or CI_STANDALONE_MERGE_BASE to resolve wheels.vllm.ai commit." >&2 + exit 1 + fi + + echo "INFO: current merge base commit with main: $merge_base_commit" + if git show --oneline -s "$merge_base_commit" 2>/dev/null; then + : + else + echo "INFO: git show unavailable in this environment; using SHA above for precompiled metadata." + fi + + # Test whether the metadata.json URL is valid, retry each 5 minutes up to 5 times. + # This avoids manual retries while a new main-branch wheel is still publishing. + if [[ "${is_rocm}" == "1" ]]; then + _rocm_env_variant="$(python3 - <<'PY' +import ctypes +import os +from pathlib import Path + + +def get_rocm_version() -> str | None: + rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH") or "/opt/rocm" + try: + librocm_core = Path(rocm_home) / "lib" / "librocm-core.so" + if not librocm_core.is_file(): + return None + librocm = ctypes.CDLL(str(librocm_core)) + get_rocm_core_version = librocm.getROCmVersion + major = ctypes.c_uint32() + minor = ctypes.c_uint32() + patch = ctypes.c_uint32() + if get_rocm_core_version( + ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch) + ) == 0: + return f"{major.value}.{minor.value}.{patch.value}" + except Exception: + return None + return None + + +version = get_rocm_version() +if version: + print(f"rocm{version.replace('.', '')}", end="") +PY +)" + _available_variants="$(curl -sf "https://wheels.vllm.ai/rocm/${merge_base_commit}/" \ + | grep -oP 'rocm\d+' | sort -u | tr '\n' ' ' || true)" + if [[ -n "${VLLM_PRECOMPILED_WHEEL_VARIANT:-}" ]]; then + _rocm_variant="${VLLM_PRECOMPILED_WHEEL_VARIANT}" + if [[ -n "${_rocm_env_variant}" && "${_rocm_variant}" != "${_rocm_env_variant}" ]]; then + echo "ERROR: VLLM_PRECOMPILED_WHEEL_VARIANT=${_rocm_variant} does not match detected environment ROCm variant ${_rocm_env_variant}" >&2 + exit 1 fi else - echo "CRITICAL: metadata.json exists but is not valid JSON, please do report in #sig-ci channel!" - echo "INFO: metadata.json content:" - cat metadata.json + _rocm_variant="${_rocm_env_variant}" + fi + if [[ -z "${_rocm_variant}" ]]; then + echo "ERROR: Could not detect ROCm variant from the environment for commit ${merge_base_commit}" >&2 exit 1 fi - fi - # failure handling & retry logic - if [ "$i" -eq 5 ]; then - echo "ERROR: metadata is still not available after 5 attempts." - echo "ERROR: Please check whether the precompiled wheel for commit $merge_base_commit is available." - echo " NOTE: If $merge_base_commit is a new commit on main, maybe try again after its release pipeline finishes." - echo " NOTE: If it fails, please report in #sig-ci channel." - exit 1 + if [[ -z "${_available_variants}" ]] \ + || [[ " ${_available_variants} " != *" ${_rocm_variant} "* ]]; then + echo "ERROR: Environment ROCm variant '${_rocm_variant}' is not published for commit ${merge_base_commit} (available:${_available_variants:-none})" >&2 + exit 1 + fi + meta_json_url="https://wheels.vllm.ai/rocm/${merge_base_commit}/${_rocm_variant}/vllm/metadata.json" + unset -v _rocm_env_variant _available_variants _rocm_variant else - echo "WARNING: metadata is not available. Retrying after 5 minutes..." - sleep 300 + meta_json_url="https://wheels.vllm.ai/${merge_base_commit}/vllm/metadata.json" fi -done + echo "INFO: will use metadata.json from ${meta_json_url}" + + for i in {1..5}; do + echo "Checking metadata.json URL (attempt $i)..." + if curl --fail "$meta_json_url" > metadata.json; then + echo "INFO: metadata.json URL is valid." + # check whether it is valid json by python (printed to stdout) + if python3 -m json.tool metadata.json; then + echo "INFO: metadata.json is valid JSON. Proceeding with the check." + # check whether there is an object in the json matching: + # "package_name": "vllm", and "platform_tag" matches the current architecture + # see `determine_wheel_url` in setup.py for more details + if python3 -c "import platform as p,json as j,sys as s; d = j.load(open('metadata.json')); \ + s.exit(int(not any(o.get('package_name') == 'vllm' and p.machine() in o.get('platform_tag') \ + for o in d)))" 2>/dev/null; then + echo "INFO: metadata.json contains a pre-compiled wheel for the current architecture." + break + else + echo "WARN: metadata.json does not have a pre-compiled wheel for the current architecture." + fi + else + echo "CRITICAL: metadata.json exists but is not valid JSON, please do report in #sig-ci channel!" + echo "INFO: metadata.json content:" + cat metadata.json + exit 1 + fi + fi + # failure handling & retry logic + if [ "$i" -eq 5 ]; then + echo "ERROR: metadata is still not available after 5 attempts." + echo "ERROR: Please check whether the precompiled wheel for commit $merge_base_commit is available." + echo " NOTE: If $merge_base_commit is a new commit on main, maybe try again after its release pipeline finishes." + echo " NOTE: If it fails, please report in #sig-ci channel." + exit 1 + else + echo "WARNING: metadata is not available. Retrying after 5 minutes..." + sleep 300 + fi + done +fi set -x @@ -87,16 +177,17 @@ fi apt remove --purge build-essential -y apt autoremove -y +rm -f /tmp/changed.file echo 'import os; os.system("touch /tmp/changed.file")' >> vllm/__init__.py # ROCm CI uses setuptools develop for editable installs (see Dockerfile.rocm and run-amd-test.sh). -_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')" -if [[ "${_vllm_target_lower}" == "rocm" ]]; then - VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 python3 setup.py develop +if [[ -n "${rocm_wheel}" ]]; then + VLLM_PRECOMPILED_WHEEL_LOCATION="${rocm_wheel}" VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps +elif [[ "${is_rocm}" == "1" ]]; then + VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 python3 setup.py develop --no-deps else - VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . + VLLM_PRECOMPILED_WHEEL_COMMIT=$merge_base_commit VLLM_USE_PRECOMPILED=1 pip3 install -vvv -e . fi -unset -v _vllm_target_lower # Run the script python3 -c 'import vllm' diff --git a/tests/test_cmake_utils.py b/tests/test_cmake_utils.py index d0673bc462eb..dd0c6261d349 100644 --- a/tests/test_cmake_utils.py +++ b/tests/test_cmake_utils.py @@ -1,10 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import shutil import subprocess +import sys from pathlib import Path +def _get_cmake_bin() -> str: + cmake = shutil.which("cmake") + if cmake: + return cmake + venv_cmake = Path(sys.executable).parent / "cmake" + if venv_cmake.is_file(): + return str(venv_cmake) + return "cmake" + + def test_exact_family_arch_precedes_generic_family_fallback(tmp_path: Path): repo_root = Path(__file__).parents[1] script = tmp_path / "test_cuda_archs.cmake" @@ -20,7 +32,7 @@ def test_exact_family_arch_precedes_generic_family_fallback(tmp_path: Path): """ ) - subprocess.run(["cmake", "-P", script], check=True) + subprocess.run([_get_cmake_bin(), "-P", script], check=True) def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( @@ -45,4 +57,26 @@ def test_extract_archs_prefers_sass_target_over_corrupted_virtual_arch( """ ) - subprocess.run(["cmake", "-P", script], check=True) + subprocess.run([_get_cmake_bin(), "-P", script], check=True) + + +def test_clear_cuda_gencode_flags(tmp_path: Path): + repo_root = Path(__file__).parents[1] + script = tmp_path / "test_clear_flags.cmake" + script.write_text( + f""" +cmake_minimum_required(VERSION 3.26) +include("{repo_root / "cmake" / "utils.cmake"}") +set(CMAKE_CUDA_FLAGS "-Wall -gencode arch=compute_80,code=sm_80") +clear_cuda_gencode_flags(CUDA_ARCH_FLAGS) +if(NOT "${{CMAKE_CUDA_FLAGS}}" STREQUAL "-Wall ") + message(FATAL_ERROR "Expected '-Wall ', got '${{CMAKE_CUDA_FLAGS}}'") +endif() +if(NOT "${{CUDA_ARCH_FLAGS}}" STREQUAL "-gencode arch=compute_80,code=sm_80") + message(FATAL_ERROR "Expected '-gencode arch=compute_80,code=sm_80', " + "got '${{CUDA_ARCH_FLAGS}}'") +endif() +""" + ) + + subprocess.run([_get_cmake_bin(), "-P", script], check=True) diff --git a/tests/test_config.py b/tests/test_config.py index 93bcdbc258d6..a2797e52126a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -29,6 +29,7 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.config.kernel import IrOpPriorityConfig from vllm.config.load import LoadConfig +from vllm.config.mamba import MambaBackendEnum from vllm.config.utils import get_field from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel from vllm.platforms import current_platform @@ -37,6 +38,52 @@ DEVICE_TYPE = current_platform.device_type +def test_kda_recoverssm_derivation_is_revalidated(): + config = SimpleNamespace( + cache_config=SimpleNamespace( + use_replayssm=True, + use_kda_recoverssm=False, + mamba_cache_mode="none", + ), + num_speculative_tokens=3, + model_config=SimpleNamespace( + supports_replayssm=True, + architecture="KimiLinearForCausalLM", + ), + mamba_config=SimpleNamespace( + backend=MambaBackendEnum.TRITON, + enable_stochastic_rounding=False, + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + kv_transfer_config=None, + use_v2_model_runner=True, + ) + + VllmConfig.validate_mamba_cached_kernel(config) + assert config.cache_config.use_replayssm + assert config.cache_config.use_kda_recoverssm + + config.cache_config.mamba_cache_mode = "align" + VllmConfig.validate_mamba_cached_kernel(config) + config.use_v2_model_runner = False + with pytest.raises(ValueError, match="VLLM_USE_V2_MODEL_RUNNER=1"): + VllmConfig.validate_mamba_cached_kernel(config) + config.use_v2_model_runner = True + config.cache_config.mamba_cache_mode = "all" + with pytest.raises(ValueError, match="only none and align"): + VllmConfig.validate_mamba_cached_kernel(config) + config.cache_config.mamba_cache_mode = "none" + + config.model_config.architecture = "NemotronHForCausalLM" + with pytest.raises(ValueError, match="only supported for Kimi-K3 KDA"): + VllmConfig.validate_mamba_cached_kernel(config) + + config.model_config.architecture = "KimiLinearForCausalLM" + config.parallel_config.pipeline_parallel_size = 2 + with pytest.raises(ValueError, match="pipeline_parallel_size=1"): + VllmConfig.validate_mamba_cached_kernel(config) + + def test_compile_config_repr_succeeds(): # setup: VllmBackend mutates the config object config = VllmConfig() @@ -66,41 +113,18 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected -@pytest.mark.parametrize( - "cudagraph_mode", - [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE], -) -def test_deepseek_v4_rejects_mrv1_piecewise_cudagraph(cudagraph_mode): - config = SimpleNamespace( - use_v2_model_runner=False, - model_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]), - compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), - ) - - with pytest.raises(ValueError, match="DeepSeek V4 does not support PIECEWISE"): - VllmConfig._validate_mrv1_piecewise_cudagraph(config) - - -@pytest.mark.parametrize( - ("use_v2_model_runner", "architecture", "cudagraph_mode"), - [ - (True, "DeepseekV4ForCausalLM", CUDAGraphMode.PIECEWISE), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.NONE), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL), - (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY), - (False, "LlamaForCausalLM", CUDAGraphMode.PIECEWISE), - ], -) -def test_mrv1_piecewise_cudagraph_allowed( - use_v2_model_runner, architecture, cudagraph_mode -): - config = SimpleNamespace( - use_v2_model_runner=use_v2_model_runner, - model_config=SimpleNamespace(architectures=[architecture]), - compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), - ) +def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch): + """ROCm keeps DeepSeek V4 on MRV1, which is still faster there.""" + from vllm.config.vllm import default_v2_model_runner_architectures + from vllm.platforms import current_platform - VllmConfig._validate_mrv1_piecewise_cudagraph(config) + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + # The lookup is lru_cached against a fixed platform. + default_v2_model_runner_architectures.cache_clear() + try: + assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures() + finally: + default_v2_model_runner_architectures.cache_clear() @pytest.mark.parametrize( @@ -330,10 +354,20 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( ), ], ) -def test_is_default_v2_model_runner_model(model_config, expected): +def test_is_default_v2_model_runner_model(model_config, expected, monkeypatch): + from vllm.config.vllm import default_v2_model_runner_architectures + from vllm.platforms import current_platform + + # The expectations below are the platform-independent defaults; ROCm's + # DeepSeek V4 carve-out is covered by test_rocm_defaults_deepseek_v4_to_mrv1. + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + default_v2_model_runner_architectures.cache_clear() config = SimpleNamespace(model_config=model_config) - assert VllmConfig._is_default_v2_model_runner_model(config) is expected + try: + assert VllmConfig._is_default_v2_model_runner_model(config) is expected + finally: + default_v2_model_runner_architectures.cache_clear() @pytest.mark.skip_global_cleanup diff --git a/tests/test_envs.py b/tests/test_envs.py index 5e0363e33a11..3d214fdbff81 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -251,6 +251,21 @@ def get_choices(): env_func() +def test_gdn_decode_kernel_env(monkeypatch: pytest.MonkeyPatch): + env_func = environment_variables["VLLM_GDN_DECODE_KERNEL"] + monkeypatch.delenv("VLLM_GDN_DECODE_KERNEL", raising=False) + assert env_func() == "cuda" + + for value in ("cuda", "triton"): + monkeypatch.setenv("VLLM_GDN_DECODE_KERNEL", value) + assert env_func() == value + + for value in ("fused", "invalid"): + monkeypatch.setenv("VLLM_GDN_DECODE_KERNEL", value) + with pytest.raises(ValueError, match="VLLM_GDN_DECODE_KERNEL"): + env_func() + + class TestEnvListWithChoices: """Test cases for env_list_with_choices function.""" diff --git a/tests/test_zen_cpu_platform_detection.py b/tests/test_zen_cpu_platform_detection.py index a1798d2b52a3..2857b1ac114c 100644 --- a/tests/test_zen_cpu_platform_detection.py +++ b/tests/test_zen_cpu_platform_detection.py @@ -2,7 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest.mock import mock_open, patch -from vllm.platforms import _is_amd_zen_cpu +import pytest + +from vllm.platforms import _is_amd_zen_cpu, resolve_current_platform_cls_qualname def test_is_amd_zen_cpu_detects_amd_with_avx512(): @@ -35,3 +37,23 @@ def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512(): def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing(): with patch("os.path.exists", return_value=False): assert not _is_amd_zen_cpu() + + +def test_cpu_target_selects_cpu_platform_from_non_cpu_wheel( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("VLLM_TARGET_DEVICE", "cpu") + + with ( + patch("vllm.platforms.vllm_version_matches_substr") as version_matches, + patch("vllm.platforms._is_amd_zen_cpu", return_value=False), + patch("vllm.platforms.rocm_platform_plugin") as rocm_plugin, + ): + assert ( + resolve_current_platform_cls_qualname() == "vllm.platforms.cpu.CpuPlatform" + ) + + # An explicit target does not depend on the installed wheel's version + # suffix or host accelerators (a native CI job can reuse a ROCm wheel). + version_matches.assert_not_called() + rocm_plugin.assert_not_called() diff --git a/tests/tool_parsers/test_cohere_command_tool_parser.py b/tests/tool_parsers/test_cohere_command_tool_parser.py new file mode 100644 index 000000000000..f102474b6289 --- /dev/null +++ b/tests/tool_parsers/test_cohere_command_tool_parser.py @@ -0,0 +1,510 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.tool_parsers.cohere_command_tool_parser import ( + CohereCommand3ToolParser, + CohereCommand4ToolParser, +) + + +@dataclass +class ExpectedToolCall: + id: str + name: str + arguments: dict + + +@dataclass +class ToolCallCase: + parser_cls: Any + model_output: str + expected_tool_calls: list[ExpectedToolCall] = field(default_factory=list) + expected_reasoning: str | None = None + expected_content: str | None = None + + +TOOL_CALL_CASES = [ + pytest.param( + ToolCallCase( + parser_cls=CohereCommand3ToolParser, + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_tool_calls=[ + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), + ], + expected_reasoning="i will call foo with query1", + ), + id="cmd3-single_tool_call", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand4ToolParser, + model_output="""\ +<|START_THINKING|> i will call foo with query1<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}} +] +<|END_ACTION|>""", + expected_tool_calls=[ + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), + ], + expected_reasoning="i will call foo with query1", + ), + id="cmd4-single_tool_call", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand3ToolParser, + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: 🌈", + expected_content="foo bar", + ), + id="cmd3-citations_no_tool_calls", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand4ToolParser, + model_output="""\ +<|START_THINKING|>This is a rainbow emoji: 🌈<|END_THINKING|> +<|START_RESPONSE|>foo bar<|END_RESPONSE|>""", + expected_reasoning="This is a rainbow emoji: 🌈", + expected_content="foo bar", + ), + id="cmd4-citations_no_tool_calls", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand3ToolParser, + model_output="""\ +<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, + {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} +] +<|END_ACTION|>""", + expected_tool_calls=[ + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), + ExpectedToolCall(id="1", name="bar", arguments={"x": 42}), + ], + expected_reasoning="first I think about foo", + ), + id="cmd3-multiple_tool_calls", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand4ToolParser, + model_output="""\ +<|START_THINKING|>first I think about foo<|END_THINKING|><|START_ACTION|> +[ + {"tool_call_id": "0", "tool_name": "foo", "parameters": {"query": "query1"}}, + {"tool_call_id": "1", "tool_name": "bar", "parameters": {"x": 42}} +] +<|END_ACTION|>""", + expected_tool_calls=[ + ExpectedToolCall(id="0", name="foo", arguments={"query": "query1"}), + ExpectedToolCall(id="1", name="bar", arguments={"x": 42}), + ], + expected_reasoning="first I think about foo", + ), + id="cmd4-multiple_tool_calls", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand3ToolParser, + model_output="""\ +<|START_THINKING|>just think, no response<|END_THINKING|>""", + expected_reasoning="just think, no response", + ), + id="cmd3-reasoning_only", + ), + pytest.param( + ToolCallCase( + parser_cls=CohereCommand4ToolParser, + model_output="""\ +<|START_THINKING|>just think, no response<|END_THINKING|>""", + expected_reasoning="just think, no response", + ), + id="cmd4-reasoning_only", + ), +] + + +class MockCohereTokenizer: + """Minimal byte-level stand-in for the Cohere tokenizer.""" + + def get_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return list(text.encode("utf-8")) + + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: + return bytes(ids).decode("utf-8", errors="replace") + + +@pytest.fixture(scope="module") +def tokenizer() -> MockCohereTokenizer: + return MockCohereTokenizer() + + +@pytest.fixture +def request_obj() -> ChatCompletionRequest: + return ChatCompletionRequest(messages=[], model="test-model") + + +REPLACEMENT_CHAR = "\ufffd" + + +def _token_deltas(tokenizer: MockCohereTokenizer, text: str) -> list[str]: + """Decode per-token string deltas, buffering incomplete multi-byte chars.""" + ids = tokenizer.encode(text, add_special_tokens=False) + deltas: list[str] = [] + prev = "" + for i in range(1, len(ids) + 1): + current = tokenizer.decode(ids[:i], skip_special_tokens=False) + if current.endswith(REPLACEMENT_CHAR): + continue + delta = current[len(prev) :] + if delta: + deltas.append(delta) + prev = current + return deltas + + +@dataclass +class StreamingResult: + tool_calls: dict[int, dict] + reasoning: str | None + content: str | None + + +def _run_streaming_over_deltas( + parser, + deltas: list[str], + request_obj: ChatCompletionRequest, +) -> StreamingResult: + accumulated: dict[int, dict] = {} + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + previous_text = "" + + for token_str in deltas: + current_text = previous_text + token_str + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=token_str, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=request_obj, + ) + if delta is not None: + if delta.reasoning is not None: + reasoning_parts.append(delta.reasoning) + if delta.content is not None: + content_parts.append(delta.content) + for tc in delta.tool_calls: + idx = tc.index + if idx not in accumulated: + accumulated[idx] = {"id": "", "name": "", "arguments": ""} + if tc.id: + accumulated[idx]["id"] = tc.id + if tc.function and tc.function.name: + accumulated[idx]["name"] = tc.function.name + if tc.function and tc.function.arguments: + accumulated[idx]["arguments"] += tc.function.arguments + previous_text = current_text + + return StreamingResult( + tool_calls=accumulated, + reasoning="".join(reasoning_parts) if reasoning_parts else None, + content="".join(content_parts) if content_parts else None, + ) + + +def _run_streaming( + parser, + tokenizer: MockCohereTokenizer, + model_output: str, + request_obj: ChatCompletionRequest, +) -> StreamingResult: + return _run_streaming_over_deltas( + parser, + _token_deltas(tokenizer, model_output), + request_obj, + ) + + +@pytest.mark.parametrize("case", TOOL_CALL_CASES) +class TestExtractToolCalls: + def test_streaming( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + case: ToolCallCase, + ): + parser = case.parser_cls(tokenizer) + streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) + + assert len(streamed.tool_calls) == len(case.expected_tool_calls) + + for i, expected_tc in enumerate(case.expected_tool_calls): + tc = streamed.tool_calls[i] + assert tc["id"] == expected_tc.id + assert tc["name"] == expected_tc.name + assert json.loads(tc["arguments"]) == expected_tc.arguments + + def test_streaming_reasoning( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + case: ToolCallCase, + ): + parser = case.parser_cls(tokenizer) + streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) + + assert streamed.reasoning == case.expected_reasoning + + def test_streaming_content( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + case: ToolCallCase, + ): + parser = case.parser_cls(tokenizer) + streamed = _run_streaming(parser, tokenizer, case.model_output, request_obj) + + assert streamed.content == case.expected_content + + def test_nonstreaming( + self, + request_obj: ChatCompletionRequest, + tokenizer: MockCohereTokenizer, + case: ToolCallCase, + ): + parser = case.parser_cls(tokenizer) + result = parser.extract_tool_calls(case.model_output, request_obj) + + assert result.tools_called == (len(case.expected_tool_calls) > 0) + assert len(result.tool_calls) == len(case.expected_tool_calls) + + for actual_tc, expected_tc in zip(result.tool_calls, case.expected_tool_calls): + assert actual_tc.type == "function" + assert actual_tc.function.name == expected_tc.name + assert json.loads(actual_tc.function.arguments) == expected_tc.arguments + + def test_streaming_nonstreaming_agree( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + case: ToolCallCase, + ): + parser_streaming = case.parser_cls(tokenizer) + parser_nonstreaming = case.parser_cls(tokenizer) + + streamed = _run_streaming( + parser_streaming, + tokenizer, + case.model_output, + request_obj, + ) + result = parser_nonstreaming.extract_tool_calls( + case.model_output, + request_obj, + ) + + assert len(streamed.tool_calls) == len(result.tool_calls) + + for i, actual_tc in enumerate(result.tool_calls): + assert streamed.tool_calls[i]["name"] == actual_tc.function.name + assert json.loads(streamed.tool_calls[i]["arguments"]) == json.loads( + actual_tc.function.arguments + ) + + +SPECIAL_TOKEN_MARKERS = ( + "<|START_THINKING|>", + "<|END_THINKING|>", + "<|START_RESPONSE|>", + "<|END_RESPONSE|>", + "<|START_ACTION|>", + "<|END_ACTION|>", + "<|START_TEXT|>", + "<|END_TEXT|>", +) + + +def _multi_token_deltas( + tokenizer: MockCohereTokenizer, + text: str, + chunk_size: int, +) -> list[str]: + ids = tokenizer.encode(text, add_special_tokens=False) + deltas: list[str] = [] + prev = "" + i = 0 + while i < len(ids): + end = min(len(ids), i + chunk_size) + current = tokenizer.decode(ids[:end], skip_special_tokens=False) + i = end + if current.endswith(REPLACEMENT_CHAR): + continue + delta = current[len(prev) :] + if delta: + deltas.append(delta) + prev = current + return deltas + + +class TestSpeculativeDecodingMultiTokenDelta: + MODEL_OUTPUT = ( + "<|START_THINKING|> i will call foo with query1<|END_THINKING|>" + "<|START_ACTION|>\n" + '[\n {"tool_call_id": "0", "tool_name": "foo", ' + '"parameters": {"query": "query1"}}\n]\n' + "<|END_ACTION|>" + ) + + @pytest.mark.parametrize( + "parser_cls", + [CohereCommand3ToolParser, CohereCommand4ToolParser], + ids=["cmd3", "cmd4"], + ) + @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) + def test_no_special_token_leak_in_streaming_deltas( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + parser_cls, + chunk_size: int, + ): + parser = parser_cls(tokenizer) + chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size) + + previous_text = "" + for token_str in chunked: + current_text = previous_text + token_str + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=token_str, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=request_obj, + ) + previous_text = current_text + if delta is None: + continue + + fields: list[tuple[str, str | None]] = [ + ("reasoning", delta.reasoning), + ("content", delta.content), + ] + for tc in delta.tool_calls or []: + if tc.function: + fields.append(("tool_call.name", tc.function.name)) + fields.append(("tool_call.arguments", tc.function.arguments)) + + for marker in SPECIAL_TOKEN_MARKERS: + for field_name, value in fields: + assert value is None or marker not in value, ( + f"special token {marker!r} leaked into {field_name} " + f"with chunk_size={chunk_size} delta={delta!r}" + ) + + @pytest.mark.parametrize( + "parser_cls", + [CohereCommand3ToolParser, CohereCommand4ToolParser], + ids=["cmd3", "cmd4"], + ) + @pytest.mark.parametrize("chunk_size", [2, 3, 4, 6]) + def test_multi_token_chunks_still_produce_correct_tool_call( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + parser_cls, + chunk_size: int, + ): + parser = parser_cls(tokenizer) + chunked = _multi_token_deltas(tokenizer, self.MODEL_OUTPUT, chunk_size) + streamed = _run_streaming_over_deltas(parser, chunked, request_obj) + + assert len(streamed.tool_calls) == 1 + tc = streamed.tool_calls[0] + assert tc["id"] == "0" + assert tc["name"] == "foo" + assert json.loads(tc["arguments"]) == {"query": "query1"} + + +class TestStreamingDeltaShape: + @pytest.mark.parametrize( + "parser_cls", + [CohereCommand3ToolParser, CohereCommand4ToolParser], + ids=["cmd3", "cmd4"], + ) + def test_reasoning_and_tool_calls_are_separate_deltas( + self, + tokenizer: MockCohereTokenizer, + request_obj: ChatCompletionRequest, + parser_cls, + ): + parser = parser_cls(tokenizer) + model_output = ( + "<|START_THINKING|> i will call foo with query1<|END_THINKING|>" + "<|START_ACTION|>\n" + '[\n {"tool_call_id": "0", "tool_name": "foo", ' + '"parameters": {"query": "query1"}}\n]\n' + "<|END_ACTION|>" + ) + + token_strings = _token_deltas(tokenizer, model_output) + previous_text = "" + saw_reasoning = False + saw_tool_call = False + + for token_str in token_strings: + current_text = previous_text + token_str + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=token_str, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=request_obj, + ) + if delta is not None: + populated = [ + delta.content is not None, + delta.reasoning is not None, + bool(delta.tool_calls), + ] + assert sum(populated) == 1, ( + "A single streaming delta must carry exactly one of " + f"content/reasoning/tool_calls, got {delta!r}" + ) + if delta.reasoning is not None: + saw_reasoning = True + if delta.tool_calls: + saw_tool_call = True + previous_text = current_text + + assert saw_reasoning, "expected at least one reasoning delta" + assert saw_tool_call, "expected at least one tool-call delta" diff --git a/tests/tool_use/test_muse_glimmer.py b/tests/tool_use/test_muse_glimmer.py new file mode 100644 index 000000000000..77f99a47ec38 --- /dev/null +++ b/tests/tool_use/test_muse_glimmer.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the MuseGlimmer ATEM tool parser and reasoning parser. + +MuseGlimmer writes every turn as a sequence of channel-scoped messages rather +than JSON, so the two parsers are tested together: the reasoning parser strips +the reasoning span and forwards the remaining channels as content, and the tool +parser reads ATEM markup out of those channels. + +Four areas, in order: + + 1. non-streaming tool-call extraction, including channel scoping (an + ```` echoed inside reasoning must never become a call); + 2. the reasoning -> tool-parser handoff, which regressed once by returning + ``content=None`` and starving the tool parser; + 3. streaming, where markers routinely straddle chunk boundaries, plus + truncation isolation for an unterminated ``to=self`` block; + 4. tool-name normalization against the tools registered on the request. + +These drive the parsers directly and need no checkpoint. The tests that require +a real tokenizer live in ``test_muse_glimmer_parse_delta.py``. +""" + +import json +from types import SimpleNamespace + +import pytest + +from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser +from vllm.tool_parsers.muse_glimmer_tool_parser import MuseGlimmerToolParser + +R: MuseGlimmerReasoningParser +T: MuseGlimmerToolParser + + +@pytest.fixture(autouse=True) +def _fresh_parsers(): + """Give each test request-scoped parser state through the real constructors.""" + global R, T + R = MuseGlimmerReasoningParser(object()) + T = MuseGlimmerToolParser(object()) + + +# Any framing token that must NEVER appear in surfaced reasoning/content. +_FRAMING = [ + "<|start|>", + "<|message|>", + "<|eom|>", + "<|eot|>", + "to=self", + "to=user", + "to=read.read", + "assistant to={name}<|message|>" + f'\n\n' + f'Paris\n' + f"\n" + ) + + +# ---------------------------------------------------------------- tool calls + + +def test_single_tool_call_after_reasoning(): + raw = ( + "to=self<|message|>Let me check the weather.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '\n\n' + 'Paris\n' + 'celsius\n' + "\n<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert out.tools_called and len(out.tool_calls) == 1 + assert out.tool_calls[0].function.name == "weather.get" + assert json.loads(out.tool_calls[0].function.arguments) == { + "city": "Paris", + "units": "celsius", + } + + +def test_parallel_calls_across_eom_boundaries(): + raw = ( + "<|start|>assistant to=math.add<|message|>" + '\n\n' + '1\n' + '2\n' + "\n<|eom|>" + "<|start|>assistant to=math.mul<|message|>" + '\n\n' + '3\n' + '4\n' + "\n<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert out.tools_called and len(out.tool_calls) == 2, len(out.tool_calls) + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"] + # JSON-typed values decode to ints + assert json.loads(out.tool_calls[0].function.arguments) == {"a": 1, "b": 2} + + +def test_echoed_invoke_in_reasoning_is_not_parsed(): + """Channel scoping: an invoke quoted inside reasoning is not a call.""" + raw = ( + 'to=self<|message|>I could call ' + '1 ' + "but I will not.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert not out.tools_called, "channel scoping failed -- echoed invoke parsed!" + assert out.content == "The answer is 42.", repr(out.content) + + +def test_plain_answer_yields_no_tool_calls(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, "to=user<|message|>Just a plain answer.<|eot|>", None + ) + assert not out.tools_called + + +def test_json_object_array_and_bool_params_decode(): + raw = ( + "<|start|>assistant to=api.call<|message|>" + '\n\n' + '{"nested": [1, 2, 3]}\n' + 'true\n' + "\n<|eot|>" + ) + out = MuseGlimmerToolParser.extract_tool_calls(T, raw, None) + assert json.loads(out.tool_calls[0].function.arguments) == { + "payload": {"nested": [1, 2, 3]}, + "flag": True, + } + + +# ------------------------------------------------- reasoning -> tool handoff + + +def test_reasoning_to_toolcall_handoff(): + """The regression: content=None here starved the tool parser.""" + raw = ( + " to=self<|message|>Let me call the tool.<|eom|>" + "<|start|>assistant to=weather.get<|message|>" + '\n\n' + 'Paris\n' + "\n" + ) + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None) + assert reasoning == "Let me call the tool.", repr(reasoning) + assert content is not None and "thinking<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" + ) + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None) + assert reasoning == "thinking", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert not MuseGlimmerToolParser.extract_tool_calls(T, content, None).tools_called + + +def test_plain_content_without_framing_passes_through(): + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning( + R, "Just a direct answer.", None + ) + assert reasoning is None and content == "Just a direct answer.", ( + reasoning, + content, + ) + + +def test_reasoning_then_parallel_calls(): + raw = ( + " to=self<|message|>need two calls<|eom|>" + "<|start|>assistant to=math.add<|message|>" + '\n\n' + '1\n\n' + "<|eom|>" + "<|start|>assistant to=math.mul<|message|>" + '\n\n' + '3\n\n' + "<|eot|>" + ) + reasoning, content = MuseGlimmerReasoningParser.extract_reasoning(R, raw, None) + assert reasoning == "need two calls", repr(reasoning) + out = MuseGlimmerToolParser.extract_tool_calls(T, content, None) + assert [t.function.name for t in out.tool_calls] == ["math.add", "math.mul"], ( + out.tool_calls + ) + + +# ----------------------------------------------------------------- streaming + + +def _stream(raw: str, chunk: int): + """Feed ``raw`` incrementally in ``chunk``-char steps through BOTH streaming + parsers; return (reasoning, content, tool_calls).""" + reasoning, content, toolcalls = [], [], [] + prev = "" + i = 0 + while i < len(raw): + cur = raw[: i + chunk] + delta = cur[len(prev) :] + dm = MuseGlimmerReasoningParser.extract_reasoning_streaming( + R, prev, cur, delta, [], [], [] + ) + if dm is not None: + if getattr(dm, "reasoning", None): + reasoning.append(dm.reasoning) + content_delta = getattr(dm, "content", None) + # Tool-channel content is an internal handoff to the tool parser, + # not client-visible content from the unified parser. + if content_delta and "" not in content_delta: + content.append(content_delta) + dt = MuseGlimmerToolParser.extract_tool_calls_streaming( + T, prev, cur, delta, [], [], [], _FakeReq() + ) + if dt is not None and dt.tool_calls: + toolcalls.extend(dt.tool_calls) + prev = cur + i += chunk + return "".join(reasoning), "".join(content), toolcalls + + +def _fn_of(tc): + fn = tc.function + if isinstance(fn, dict): + return fn.get("name"), fn.get("arguments") + return fn.name, fn.arguments + + +RAW_TOOLCALL = ( + " to=self<|message|>I should read the hostname file to answer.<|eom|>" + "<|start|>assistant to=read.read<|message|>" + '\n\n' + '/etc/hostname\n' + "\n" +) + +RAW_ANSWER = ( + " to=self<|message|>Think about it.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>" +) + +# NO closing <|eom|> -> truncated CoT +RAW_TRUNCATED = ( + " to=self<|message|>Maybe I should call " + '\n\n' + '/etc/hostname\n' + "\n but wait" +) + + +def _check_toolcall_stream(chunk): + reasoning, content, tcs = _stream(RAW_TOOLCALL, chunk) + # (a) no framing token leaks into reasoning or content + for f in _FRAMING: + assert f not in reasoning, ( + f"framing {f!r} leaked into reasoning (chunk={chunk})" + ) + assert f not in content, f"framing {f!r} leaked into content (chunk={chunk})" + # (b) exactly one tool_call with correct name + args + assert len(tcs) == 1, f"expected 1 tool_call, got {len(tcs)} (chunk={chunk})" + name, args = _fn_of(tcs[0]) + assert name == "read.read", name + assert json.loads(args) == {"path": "/etc/hostname"}, args + assert tcs[0].index == 0 and tcs[0].type == "function" and tcs[0].id + # (c) reasoning captured separately and clean + assert reasoning == "I should read the hostname file to answer.", repr(reasoning) + assert content == "", repr(content) + + +def test_streaming_toolcall_chunk3(): + _check_toolcall_stream(3) + + +def test_streaming_toolcall_charwise(): + # worst case: markers arrive one char at a time (mid-marker deltas) + _check_toolcall_stream(1) + + +def test_streaming_toolcall_bigchunks(): + _check_toolcall_stream(17) + + +def test_streaming_reasoning_then_content(): + reasoning, content, tcs = _stream(RAW_ANSWER, 3) + for f in _FRAMING: + assert f not in reasoning and f not in content, f + assert reasoning == "Think about it.", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert tcs == [] + + +def test_truncated_cot_no_toolcall_nonstreaming(): + out = MuseGlimmerToolParser.extract_tool_calls(T, RAW_TRUNCATED, _FakeReq()) + assert not out.tools_called and out.tool_calls == [] + # partial reasoning must still be recovered by the reasoning parser + reasoning, _ = MuseGlimmerReasoningParser.extract_reasoning( + R, RAW_TRUNCATED, _FakeReq() + ) + assert reasoning and "Maybe I should call" in reasoning, repr(reasoning) + + +def test_truncated_cot_no_toolcall_streaming(): + _, _, tcs = _stream(RAW_TRUNCATED, 3) + assert tcs == [], f"truncated CoT invoke leaked as streaming tool call: {tcs}" + + +# ------------------------------------------------------- name normalization +# +# MuseGlimmer emits `get_weather.get_weather` for a bare-registered +# `get_weather`, and `weather.get` verbatim for a namespaced one. The parser +# normalizes against the tools actually registered on the request. + + +def test_doubled_bare_name_collapses(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("get_weather.get_weather"), _req("get_weather") + ) + assert out.tools_called and out.tool_calls[0].function.name == "get_weather", ( + out.tool_calls[0].function.name + ) + + +def test_namespaced_name_preserved(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("weather.get"), _req("weather.get") + ) + assert out.tool_calls[0].function.name == "weather.get" + + +def test_unregistered_namespace_is_preserved(): + # Suffix-only matching can silently dispatch a tool from the wrong namespace. + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("foo.get_weather"), _req("get_weather") + ) + assert out.tool_calls[0].function.name == "foo.get_weather" + + +def test_trailing_segment_ambiguous_left_alone(): + # two registered tools share leaf 'get' -> ambiguous -> do NOT rewrite + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("x.get"), _req("weather.get", "time.get") + ) + assert out.tool_calls[0].function.name == "x.get" + + +def test_no_registered_tools_passthrough(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("get_weather.get_weather"), None + ) + assert out.tool_calls[0].function.name == "get_weather.get_weather" + + +def test_exact_match_kept(): + out = MuseGlimmerToolParser.extract_tool_calls( + T, _call("get_weather"), _req("get_weather") + ) + assert out.tool_calls[0].function.name == "get_weather" diff --git a/tests/tool_use/test_muse_glimmer_parse_delta.py b/tests/tool_use/test_muse_glimmer_parse_delta.py new file mode 100644 index 000000000000..805a46e13bdb --- /dev/null +++ b/tests/tool_use/test_muse_glimmer_parse_delta.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""parse_delta-level regression tests for the MuseGlimmer parsers. + +These drive the UNIFIED parser-engine streaming API the serving layer actually +uses (``DelegatingParser.parse_delta``), NOT the ``extract_*_streaming`` methods +directly. This is the coverage that was missing when a live no-tools streaming +request leaked raw harmony framing into ``content``: the phase machine marked +``reasoning_ended=True`` at the prompt boundary (because ``is_reasoning_end`` on +a prompt with no ``to=self`` wrongly returned True), skipping the reasoning +phase for the whole generation. + +Requires a real MuseGlimmer tokenizer; skipped if the checkpoint is unavailable. +""" + +import json +import os + +import pytest + +CKPT = os.environ.get("MUSE_GLIMMER_CKPT", "") + +pytestmark = pytest.mark.skipif( + not os.path.isdir(CKPT), reason=f"MuseGlimmer checkpoint not found at {CKPT}" +) + +_FRAMING = [ + "<|start|>", + "<|message|>", + "<|eom|>", + "<|eot|>", + "to=self", + "to=user", + "Let me think step by step about the sum.<|eom|>" + "<|start|>assistant to=user<|message|>The answer is 42.<|eot|>", + ) + _assert_no_framing(content) + assert reasoning == "Let me think step by step about the sum.", repr(reasoning) + assert content == "The answer is 42.", repr(content) + assert tools == [] + + +def test_parse_delta_content_only(parser_cls, tok): + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=user<|message|>Just a direct answer.<|eot|>", + ) + _assert_no_framing(content) + assert content == "Just a direct answer.", repr(content) + assert tools == [] + + +def test_parse_delta_tool_call(parser_cls, tok): + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>I should read the hostname.<|eom|>" + "<|start|>assistant to=read.read<|message|>" + '\n\n' + '/etc/hostname\n' + "\n", + ) + _assert_no_framing(content) + assert reasoning == "I should read the hostname.", repr(reasoning) + assert len(tools) == 1, tools + idx, name, args = tools[0] + assert idx == 0 and name == "read.read" + assert json.loads(args) == {"path": "/etc/hostname"} + + +def test_parse_delta_truncated_cot_no_toolcall(parser_cls, tok): + # Contemplated invoke inside an unterminated to=self block: no tool call, + # no framing leak into content, partial reasoning recovered. + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>Maybe I should call " + '\n\n' + '/etc/hostname\n' + "\n but wait", + ) + _assert_no_framing(content) + assert tools == [], f"contemplated invoke leaked as tool call: {tools}" + assert "Maybe I should call" in reasoning, repr(reasoning) + + +def test_parse_delta_reasoning_suppressed_when_not_requested(parser_cls, tok): + class _NoReasonReq: + tools = None + tool_choice = None + include_reasoning = False + + reasoning, content, tools = _drive( + parser_cls, + tok, + " to=self<|message|>secret thoughts<|eom|>" + "<|start|>assistant to=user<|message|>Public answer.<|eot|>", + req=_NoReasonReq(), + ) + assert reasoning == "", repr(reasoning) # suppressed + assert content == "Public answer.", repr(content) + _assert_no_framing(content) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/transformers_utils/test_dspark_mla_config.py b/tests/transformers_utils/test_dspark_mla_config.py index f43d668b73cd..b81f06e14388 100644 --- a/tests/transformers_utils/test_dspark_mla_config.py +++ b/tests/transformers_utils/test_dspark_mla_config.py @@ -140,26 +140,3 @@ def test_dspark_mla_speculative_config_preserves_architecture(tmp_path): assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"] assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark" assert speculative_config.draft_model_config.use_mla - - -def test_dspark_mla_rejects_decode_context_parallelism(tmp_path): - target_path = tmp_path / "target" - draft_path = tmp_path / "draft" - _write_target_config(target_path) - _write_dspark_config(draft_path) - target_config = ModelConfig( - model=str(target_path), tokenizer_mode="skip", max_model_len=32768 - ) - - with pytest.raises(ValueError, match="does not currently support decode context"): - SpeculativeConfig( - model=str(draft_path), - method="dspark", - num_speculative_tokens=8, - target_model_config=target_config, - target_parallel_config=ParallelConfig( - tensor_parallel_size=2, - decode_context_parallel_size=2, - distributed_executor_backend="external_launcher", - ), - ) diff --git a/tests/transformers_utils/test_muse_glimmer_config.py b/tests/transformers_utils/test_muse_glimmer_config.py new file mode 100644 index 000000000000..ce9d4a819e22 --- /dev/null +++ b/tests/transformers_utils/test_muse_glimmer_config.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""MuseGlimmer config normalization. + +Two schema pairs have to converge on the same model, and each has produced a +silent-wrong-output bug: + + 1. flat (legacy converter) vs nested (canonical) ``config.json``. A FLAT + config once deserialized to an all-default text config, silently dropping + every checkpoint value. + 2. native vs modular attention config. The modular HF text_config OMITS + use_qk_norm / use_attn_output_gate (read as None) and ships a PRE-FOLDED + qk_scale_factor (43.784/sqrt(128)=3.87). Missing flags must read as True + (MuseGlimmer always applies QK-norm + output gate) and the query pre-scale + must normalize so native and modular land on the same number. +""" + +import math +from typing import Any + +from vllm.model_executor.models.muse_glimmer import ( + _muse_glimmer_query_prescale, + _muse_glimmer_use_attn_output_gate, + _muse_glimmer_use_qk_norm, +) +from vllm.transformers_utils.configs.muse_glimmer import MuseGlimmerConfig + +# A representative FLAT config (Ruan rl_v1/hf shape), trimmed. +FLAT: dict[str, Any] = { + "architectures": ["MuseGlimmerForCausalLM"], + "model_type": "muse_glimmer", + "has_vision": True, + "bos_token_id": 200000, + "eos_token_id": 200001, + "vocab_size": 202048, + "hidden_size": 6656, + "intermediate_size": 19968, + "num_hidden_layers": 52, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 128, + "hidden_act": "silu", + "max_position_embeddings": 16384, + "rms_norm_eps": 1e-5, + "post_norm_eps": 1e-8, + "qk_scale_factor": 43.7840518911, + "use_qk_norm": True, + "use_attn_output_gate": True, + "output_multiplier": 0.19611613513818404, + "output_soft_cap_temp": 20.0, + "normalize_tok_embeddings": True, + "rope_theta": 500000.0, + "sliding_window": 2048, + "patch_token_id": 200092, + "vision_latent_dim": 1536, + "vision_heads": 16, + "vision_layers": 50, + "vision_output_dim": 6144, + "vision_patch_size": 14, + "vision_patch_temporal": 2, + "vision_adapter_dim": 4096, + "vision_pos_emb_grid_h": 32, + "vision_pos_emb_grid_w": 32, +} + +NESTED: dict[str, Any] = { + "architectures": ["MuseGlimmerForCausalLM"], + "model_type": "muse_glimmer", + "image_token_id": 200092, + "text_config": { + "model_type": "muse_glimmer_text", + "vocab_size": 202048, + "hidden_size": 6656, + "num_hidden_layers": 52, + "hidden_activation": "silu", + "final_logit_softcapping": 20.0, + "qk_scale_factor": 43.7840518911, + "rope_parameters": {"rope_type": "default", "rope_theta": 500000.0}, + }, + "vision_config": { + "model_type": "muse_glimmer_vision", + "hidden_size": 1536, + "num_hidden_layers": 50, + }, +} + +HEAD_DIM = 128 +SQRT_HD = math.sqrt(HEAD_DIM) +NATIVE = 43.7840518911 +FOLDED = NATIVE / SQRT_HD # 3.8700... + + +class Cfg: + """Attention-config stand-in carrying only the fields under test.""" + + def __init__(self, **kw): + self.head_dim = HEAD_DIM + for k, v in kw.items(): + setattr(self, k, v) + + +# ------------------------------------------------------------ flat vs nested + + +def test_flat_config_values_respected(): + c = MuseGlimmerConfig(**FLAT) + t = c.text_config + assert t.hidden_size == 6656 + assert t.num_hidden_layers == 52 + assert t.vocab_size == 202048 + assert t.head_dim == 128 + assert t.hidden_activation == "silu" # renamed from hidden_act + assert t.final_logit_softcapping == 20.0 # renamed from output_soft_cap_temp + assert abs(t.output_multiplier - 0.19611613513818404) < 1e-12 + assert abs(t.qk_scale_factor - 43.7840518911) < 1e-9 + assert t.rope_parameters["rope_theta"] == 500000.0 # from flat rope_theta + # vision hoisted + renamed + assert c.vision_config.hidden_size == 1536 + assert c.vision_config.num_hidden_layers == 50 + assert c.vision_config.output_dim == 6144 + # flat patch_token_id -> image_token_id + assert c.image_token_id == 200092 + + +def test_flat_config_no_silent_default(): + # The regression: a non-default value MUST be honored, not silently dropped. + flat = dict(FLAT) + flat["hidden_size"] = 4096 + flat["num_hidden_layers"] = 40 + c = MuseGlimmerConfig(**flat) + assert c.text_config.hidden_size == 4096, "flat hidden_size silently ignored!" + assert c.text_config.num_hidden_layers == 40, "flat num_hidden_layers ignored!" + + +def test_nested_config_unchanged(): + c = MuseGlimmerConfig(**NESTED) + assert c.text_config.hidden_size == 6656 + assert c.text_config.hidden_activation == "silu" + assert c.text_config.final_logit_softcapping == 20.0 + assert c.vision_config.hidden_size == 1536 + assert c.vision_config.num_hidden_layers == 50 + assert c.image_token_id == 200092 + + +# -------------------------------------------------------- native vs modular + + +def test_qk_norm_missing_defaults_true(): + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=None)) is True # modular + assert _muse_glimmer_use_qk_norm(Cfg()) is True # absent + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=True)) is True # native + assert _muse_glimmer_use_qk_norm(Cfg(use_qk_norm=False)) is False # explicit off + + +def test_output_gate_missing_defaults_true(): + assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=None)) is True + assert _muse_glimmer_use_attn_output_gate(Cfg()) is True + assert _muse_glimmer_use_attn_output_gate(Cfg(use_attn_output_gate=False)) is False + + +def test_query_prescale_native_and_modular_converge(): + # Both schemas must yield the SAME final scale_query_by (~3.87). + assert ( + abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=NATIVE)) - FOLDED) < 1e-9 + ) + assert ( + abs(_muse_glimmer_query_prescale(Cfg(qk_scale_factor=FOLDED)) - FOLDED) < 1e-9 + ) + + +def test_query_prescale_explicit_wins(): + c = Cfg(scale_query_by=FOLDED, qk_scale_factor=NATIVE) + assert abs(_muse_glimmer_query_prescale(c) - FOLDED) < 1e-9 diff --git a/tests/utils_/test_serial_utils.py b/tests/utils_/test_serial_utils.py index 85661657d50f..8dd0229b6218 100644 --- a/tests/utils_/test_serial_utils.py +++ b/tests/utils_/test_serial_utils.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import io + +import numpy as np +import pybase64 import pytest import torch @@ -12,6 +16,7 @@ Endianness, MmMetadataDType, binary2tensor, + numpy2base64, tensor2binary, ) @@ -19,6 +24,17 @@ INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys()) +def test_numpy2base64_round_trip(): + array = np.arange(24, dtype=np.uint8).reshape(2, 3, 4) + + decoded = np.load( + io.BytesIO(pybase64.b64decode(numpy2base64(array))), + allow_pickle=False, + ) + + np.testing.assert_array_equal(decoded, array) + + def _build_integer_tensor( embed_dtype: MmMetadataDType, shape: tuple[int, ...] ) -> torch.Tensor: diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index 3d9ab6df8d84..cb35864b70d1 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -848,7 +848,7 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype): AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, reason="FlashInfer is not available.", ) -def test_flashinfer_native_prefill_with_sinks(): +def test_flashinfer_native_prefill_with_sinks(default_vllm_config): if not ( current_platform.is_cuda() and current_platform.is_device_capability_family(120) ): diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py index acc38919b802..f3a46bbb1900 100644 --- a/tests/v1/attention/test_flashinfer_mla_dcp.py +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -45,6 +45,7 @@ def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): impl.bmm1_scale = 1.0 impl.bmm2_scale = 1.0 impl.need_to_return_lse_for_decode = True + impl.dcp_world_size = 2 impl.num_heads = 6 impl.qk_nope_head_dim = 128 impl.kv_lora_rank = 512 diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py index 3a7677e7511f..12a8cf5cbc39 100644 --- a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py +++ b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py @@ -7,6 +7,9 @@ import torch from vllm.config import set_current_vllm_config +from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( + _required_sm120_sparse_topk, +) from vllm.platforms.interface import DeviceCapability from vllm.utils import flashinfer as fi_utils from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( @@ -52,3 +55,33 @@ def test_v32_glm_sm120_backend_accepts_glm_block_size( ) assert invalid_reasons == [] + + +def test_sm120_dsv4_capability_checks_exact_dispatch_shape(monkeypatch) -> None: + fake_module = SimpleNamespace( + _DECODE_DSV4_DISPATCH=frozenset({(32, 128), (32, 192)}) + ) + monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True) + monkeypatch.setattr(fi_utils, "_get_submodule", lambda _name: fake_module) + fi_utils.has_flashinfer_sparse_mla_sm120_config.cache_clear() + + assert fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 128) + assert fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 192) + assert not fi_utils.has_flashinfer_sparse_mla_sm120_config(32, 256) + assert not fi_utils.has_flashinfer_sparse_mla_sm120_config(16, 192) + + fi_utils.has_flashinfer_sparse_mla_sm120_config.cache_clear() + + +def test_sm120_dsv4_required_topk_tracks_dspark_width() -> None: + causal = SimpleNamespace( + attention_config=SimpleNamespace(use_non_causal=False), + speculative_config=SimpleNamespace(num_speculative_tokens=5), + ) + dspark = SimpleNamespace( + attention_config=SimpleNamespace(use_non_causal=True), + speculative_config=SimpleNamespace(num_speculative_tokens=5), + ) + + assert _required_sm120_sparse_topk(causal, 128) == 128 + assert _required_sm120_sparse_topk(dspark, 128) == 192 diff --git a/tests/v1/attention/test_group_head_counts.py b/tests/v1/attention/test_group_head_counts.py new file mode 100644 index 000000000000..b799bccc29d6 --- /dev/null +++ b/tests/v1/attention/test_group_head_counts.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Scheduler metadata sizes a scratchpad from the query head count, so it must +come from the builder's own group: the model-wide ``get_num_attention_heads()`` +is wrong for models that vary it per layer (e.g. Laguna), and too small a +scratchpad is indexed past its end. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.cpu_attn import ( + CPUAttentionBackendImpl, + CPUAttentionMetadataBuilder, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cpu(), reason="CPU attention backend" +) + +# Laguna's shape: 48 query heads model-wide, 64 on its sliding layers, both +# against 8 KV heads. +MODEL_WIDE_NUM_HEADS = 48 +NUM_KV_HEADS = 8 + + +def _layers(layer_num_heads: list[int]): + """Stand-in attention layers, one per head count, as one attention group.""" + return { + f"layer_{i}": SimpleNamespace( + impl=MagicMock( + spec=CPUAttentionBackendImpl, + num_heads=num_heads, + sliding_window=None, + ) + ) + for i, num_heads in enumerate(layer_num_heads) + } + + +def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder: + layers = _layers(layer_num_heads) + vllm_config = MagicMock() + vllm_config.model_config.dtype = torch.bfloat16 + vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS + vllm_config.cache_config.block_size = 16 + vllm_config.cache_config.cache_dtype = "auto" + kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64) + + with ( + patch( + "vllm.v1.attention.backends.utils.get_layers_from_vllm_config", + return_value=layers, + ), + patch( + "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config", + return_value=layers, + ), + ): + return CPUAttentionMetadataBuilder( + kv_cache_spec=kv_cache_spec, + layer_names=list(layers), + vllm_config=vllm_config, + device=torch.device("cpu"), + ) + + +@pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16]) +def test_num_heads_comes_from_the_group(group_num_heads): + """The group's own count wins, even when it is not the model-wide one.""" + builder = _build([group_num_heads, group_num_heads]) + assert builder.num_heads == group_num_heads + + +def test_mixed_head_counts_in_one_group_are_rejected(): + """Grouping guarantees uniformity; a mixed group means that broke.""" + with pytest.raises(AssertionError, match="share num_heads"): + _build([MODEL_WIDE_NUM_HEADS, 64]) diff --git a/tests/v1/attention/test_indexer_native_next_n.py b/tests/v1/attention/test_indexer_native_next_n.py new file mode 100644 index 000000000000..3ec14fe62ef9 --- /dev/null +++ b/tests/v1/attention/test_indexer_native_next_n.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which next_n the DSA indexer decode path may hand to DeepGEMM unflattened. + +Getting this wrong is not a slow path but a crash: `fp8_fp4_paged_mqa_logits` +asserts both that the architecture implements the requested `next_n` and that +the schedule metadata was sized for the matching slot count. +""" + +import pytest + +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import _paged_mqa_logits_schedule_slots +from vllm.v1.attention.backends.mla import indexer + +NUM_SMS = 114 # H100 PCIe + + +def _set_arch(monkeypatch, family: int, *, cuda: bool = True, deep_gemm: bool = True): + monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda) + monkeypatch.setattr( + current_platform, + "is_device_capability_family", + lambda capability, device_id=0: capability // 10 == family, + ) + monkeypatch.setattr(indexer, "has_deep_gemm", lambda: deep_gemm) + + +@pytest.mark.parametrize( + "family,expected_native", + [ + # SM90 gained next_n=4 (MTP=3) via 2-CTA multicast, but never 3. + (9, {1, 2, 4}), + # SM100 schedules any next_n with multi-atom tiles. + (10, {1, 2, 3, 4, 5, 8}), + # SM120 advertises multi-atom too but is unvalidated on hardware, so + # it stays on the conservative gate. Loosen it only with measurements. + (12, {1, 2}), + ], +) +def test_native_decode_gate_per_architecture(monkeypatch, family, expected_native): + _set_arch(monkeypatch, family) + for next_n in (1, 2, 3, 4, 5, 8): + assert indexer._supports_native_decode(next_n) == (next_n in expected_native), ( + f"family={family} next_n={next_n}" + ) + + +@pytest.mark.parametrize( + "cuda,deep_gemm", [(False, True), (True, False), (False, False)] +) +def test_native_decode_gate_without_deepgemm(monkeypatch, cuda, deep_gemm): + """Without the DeepGEMM kernels only the shapes every backend handles.""" + _set_arch(monkeypatch, 9, cuda=cuda, deep_gemm=deep_gemm) + assert [indexer._supports_native_decode(n) for n in (1, 2, 3, 4)] == [ + True, + True, + False, + False, + ] + + +def test_sm90_next_n_4_halves_the_schedule_slots(monkeypatch): + """SM90 next_n=4 runs one scheduler task per 2-CTA cluster, not per SM.""" + _set_arch(monkeypatch, 9) + assert _paged_mqa_logits_schedule_slots(NUM_SMS, 4) == NUM_SMS // 2 + for next_n in (1, 2, 3): + assert _paged_mqa_logits_schedule_slots(NUM_SMS, next_n) == NUM_SMS + + +@pytest.mark.parametrize("family", [10, 12]) +def test_multicast_is_sm90_only(monkeypatch, family): + _set_arch(monkeypatch, family) + for next_n in (1, 2, 3, 4): + assert _paged_mqa_logits_schedule_slots(NUM_SMS, next_n) == NUM_SMS diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cd291937241a..72773adcb674 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -837,9 +837,86 @@ class _AttnMeta: def test_tokenspeed_mla_noncausal_capability(): builder = tokenspeed_mla_module.TokenspeedMLAMetadataBuilder assert builder.supports_non_causal_multi_token_decode + assert builder.supports_non_causal_multi_token_dcp assert tokenspeed_mla_module.TokenspeedMLABackend.supports_non_causal() +def test_flashinfer_mla_dcp_multi_token_decode_uses_per_query_bounds(monkeypatch): + flashinfer_mla_module = pytest.importorskip( + "vllm.v1.attention.backends.mla.flashinfer_mla" + ) + + decode_call = None + + def fake_decode(**kwargs): + nonlocal decode_call + decode_call = kwargs + query = kwargs["query"] + output = torch.empty(*query.shape[:-1], 512, dtype=torch.bfloat16) + lse = torch.empty(query.shape[0], query.shape[-2], dtype=torch.float32) + return output, lse + + monkeypatch.setattr( + flashinfer_mla_module, + "trtllm_batch_decode_with_kv_cache_mla", + fake_decode, + ) + monkeypatch.setattr( + flashinfer_mla_module, + "_get_workspace_buffer", + lambda return_lse: torch.empty(1, dtype=torch.int8), + ) + + impl = object.__new__(flashinfer_mla_module.FlashInferMLAImpl) + impl.dcp_world_size = 2 + impl.dcp_rank = 1 + impl.cp_kv_cache_interleave_size = 1 + impl.need_to_return_lse_for_decode = True + impl.kv_lora_rank = 512 + impl.qk_nope_head_dim = 128 + impl.qk_rope_head_dim = 64 + impl.bmm1_scale = 1.0 + impl.bmm2_scale = 1.0 + + block_table = torch.tensor([[1], [2]], dtype=torch.int32) + metadata = SimpleNamespace( + num_decodes=2, + num_decode_tokens=6, + max_seq_len=7, + causal=True, + decode=SimpleNamespace( + block_table=block_table, + seq_lens=torch.tensor([5, 6], dtype=torch.int32), + dcp_tot_seq_lens=torch.tensor([10, 13], dtype=torch.int32), + flattened_block_table=None, + flattened_seq_lens=None, + query_len=0, + ), + ) + query = torch.empty(6, 2, 576, dtype=torch.bfloat16) + kv_cache = torch.empty(3, 16, 576, dtype=torch.bfloat16) + + output, lse = impl.forward_mqa( + query, + kv_cache, + metadata, + SimpleNamespace(), + ) + + assert output.shape == (6, 2, 512) + assert lse is not None + assert lse.shape == (6, 2) + assert decode_call is not None + assert decode_call["query"].shape == (6, 1, 2, 576) + torch.testing.assert_close( + decode_call["seq_lens"], + torch.tensor([4, 4, 5, 5, 6, 6], dtype=torch.int32), + ) + torch.testing.assert_close( + decode_call["block_tables"], block_table.repeat_interleave(3, dim=0) + ) + + @pytest.mark.parametrize( ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"), [ 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 97199a15ab49..d237ce9f345b 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 @@ -261,6 +261,80 @@ def test_hybrid_mamba_align_partial_hash_hit(): assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8 +def test_eagle_group_registers_unaligned_tail_under_partial_hash_hits(): + """An EAGLE group must not re-floor what partial hash hits leaves un-floored. + + ``cache_blocks`` decides once how far a request may be registered, and with + fine-grained partial hash hits that bound is the raw token count. The EAGLE + branch then re-derives its own bound for the lookahead block; if it rounds + down to ``scheduler_block_size`` again, everything between the last aligned + boundary and the tail stops being registered -- ``(n % scheduler_block_size) + - manager.block_size`` tokens per call, which is most of a segment whenever + the group's own block is much smaller than the scheduler block. + """ + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=40, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + coordinator = manager.coordinator + assert coordinator.enable_partial_hash_hits + # The full-attention group is the EAGLE one, and its block is smaller than + # the scheduler block -- the geometry where re-flooring loses tokens. + eagle_manager = coordinator.single_type_managers[0] + eagle_manager.use_eagle = True + assert eagle_manager.block_size < coordinator.scheduler_block_size + + # Deliberately not a multiple of the scheduler block, so the two bounds + # differ: floor(22/8)*8 + 2 = 18 against 22. + num_tokens = coordinator.scheduler_block_size * 2 + hash_block_size * 3 + req = make_request("0", list(range(num_tokens)), hash_block_size, sha256) + + recorded: list[int] = [] + for single_type_manager in coordinator.single_type_managers: + original = single_type_manager.cache_blocks + + def spy(request, num_tokens_to_cache, *args, _orig=original, **kwargs): + recorded.append(num_tokens_to_cache) + return _orig(request, num_tokens_to_cache, *args, **kwargs) + + single_type_manager.cache_blocks = spy + + # allocate_slots caches on the way out, so this exercises the real path. + computed_blocks, num_computed, _ = manager.get_computed_blocks(req) + assert manager.allocate_slots(req, num_tokens, num_computed, computed_blocks) + + # Every group, EAGLE or not, may register the whole unaligned tail. + assert recorded == [num_tokens] * len(coordinator.single_type_managers) + + def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): hash_block_size = 2 block_size = 2 * hash_block_size diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 88d17e9acdc7..40906f621369 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -88,6 +88,7 @@ def _make_groups(n_c4, n_c128, n_swa): def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None): config = MagicMock() config.cache_config.num_gpu_blocks_override = None + config.cache_config.prefix_cache_retention_interval = 0 config.kv_transfer_config = None if kv_connector_extra_config is not None: config.kv_transfer_config = MagicMock() @@ -322,6 +323,7 @@ def test_hma_attention_groups_keep_default_backing(self): ) assert config.num_blocks == 32 + assert config.prefix_cache_retention_interval == 0 assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32 assert config.kv_cache_tensors == [ KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]), diff --git a/tests/v1/core/test_encoder_cache_manager.py b/tests/v1/core/test_encoder_cache_manager.py index e225666f8443..e56bcbf5c63a 100644 --- a/tests/v1/core/test_encoder_cache_manager.py +++ b/tests/v1/core/test_encoder_cache_manager.py @@ -167,6 +167,25 @@ def test_get_freed_mm_hashes_clears_freed_list(): assert manager.get_freed_mm_hashes() == [] +def test_reallocated_hash_is_not_reported_as_freed(): + manager = EncoderCacheManager(cache_size=8) + req_a = MockRequest("reqA", ["a"], [4]) + req_b = MockRequest("reqB", ["b"], [4]) + req_c = MockRequest("reqC", ["c"], [4]) + + manager.allocate(req_a, 0) + manager.allocate(req_b, 0) + manager.free(req_a) + manager.free(req_b) + + assert manager.can_allocate(req_c, 0, int(1e9), 0) + manager.allocate(req_c, 0) + assert manager.can_allocate(req_a, 0, int(1e9), 0) + manager.allocate(req_a, 0) + + assert manager.get_freed_mm_hashes() == ["b"] + + def test_schedule_request_multi_images_respect_space_limit(): manager = EncoderCacheManager(cache_size=10) req = MockRequest("reqA", ["a", "b"], [5, 6]) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 33430ce85063..bfe4996b1aa8 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -816,6 +816,7 @@ def test_metrics_empty_stats(): def test_get_kv_cache_configs_multiple_workers(): model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() same_kv_cache_specs = [ @@ -1173,6 +1174,7 @@ def test_get_kv_cache_configs_multiple_workers(): def test_get_kv_cache_configs_pp_sharding(asymmetric_memory): model_config = ModelConfig(max_model_len=512) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.prefix_cache_retention_interval = None ref_kv_cache_spec = new_kv_cache_spec() pp_kv_cache_specs = [ @@ -1702,6 +1704,7 @@ def test_get_kv_cache_config_one_worker(): # pass max_model_len to pass check_enough_kv_cache_memory model_config = ModelConfig(max_model_len=16) vllm_config = VllmConfig(model_config=model_config) + vllm_config.cache_config.prefix_cache_retention_interval = None mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2 # all layers are full attention -> single group @@ -2015,6 +2018,7 @@ def test_get_kv_cache_config_one_worker(): def test_get_kv_cache_configs_attention_free(): kv_cache_specs: dict[str, KVCacheSpec] = {} vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16)) + vllm_config.cache_config.prefix_cache_retention_interval = None kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0]) assert kv_cache_configs == [ KVCacheConfig( @@ -2766,21 +2770,20 @@ def test_unify_kv_cache_page_size_padding_requires_backend_support(): kv_cache_utils.unify_kv_cache_spec_page_size(specs) -def test_unpadded_page_size_without_quant_matches_real_page(): - # Without quantization the offload transfer width is just the raw page. - spec = new_kv_cache_spec() - assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes - assert spec.page_size_bytes == spec.unpadded_page_size_bytes - - def test_unpadded_page_size_includes_per_token_head_scales(): # Per-token-head quant carries inline fp32 scales that are carved from the - # raw KV allocation, so they must be budgeted into the offload width. - spec = new_kv_cache_spec( - dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD + # raw KV allocation, so they must be budgeted into the offload width. The + # packing is published by the owning backend's customize_spec hook. + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend + + dense = new_kv_cache_spec(dtype=torch.uint8) + spec = TritonAttentionBackend.customize_spec( + new_kv_cache_spec( + dtype=torch.uint8, kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD + ) ) scales = 2 * spec.block_size * spec.num_kv_heads * 4 - assert spec.unpadded_page_size_bytes == spec.real_page_size_bytes + scales + assert spec.unpadded_page_size_bytes == dense.unpadded_page_size_bytes + scales assert spec.page_size_bytes == spec.unpadded_page_size_bytes diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 13ed7c7b9d8b..6cd117b1fa4a 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -4,6 +4,7 @@ import copy from collections.abc import Callable +from dataclasses import replace from math import lcm from types import SimpleNamespace @@ -110,6 +111,11 @@ def make_kv_cache_manager(kv_cache_config: KVCacheConfig, **kwargs) -> KVCacheMa "scheduler_block_size", lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)), ) + if "retention_interval" in kwargs: + kv_cache_config = replace( + kv_cache_config, + prefix_cache_retention_interval=kwargs.pop("retention_interval"), + ) return KVCacheManager(kv_cache_config, **kwargs) @@ -3119,9 +3125,8 @@ def test_hybrid_cache_blocks_clamped_to_lcm(): ) -def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch): +def test_hybrid_local_kv_retention_interval_aligns_in_manager(): """Verify fixed intervals retain sparse tails plus the latest replay tail.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3153,6 +3158,7 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch): max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=64, ) # The SWA manager uses the configured 64-token interval (a multiple of the @@ -3184,18 +3190,15 @@ def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch): "interval, expected_match", [ # scheduler_block_size is 32 (= lcm(4*8, 8)); 33 is not a multiple of it. - ("33", "multiple of scheduler_block_size"), + (33, "multiple of scheduler_block_size"), # A negative multiple (-32 % 32 == 0) must still be rejected explicitly, # otherwise it would pass the modulo check and silently degrade to dense. - ("-32", "non-negative"), + (-32, "non-negative"), ], ) -def test_hybrid_local_kv_retention_interval_rejects_invalid( - monkeypatch, interval, expected_match -): +def test_hybrid_local_kv_retention_interval_rejects_invalid(interval, expected_match): """A retention interval that is negative or not a multiple of scheduler_block_size errors out at construction time.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", interval) block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3228,12 +3231,36 @@ def test_hybrid_local_kv_retention_interval_rejects_invalid( max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=interval, ) -def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch): +def test_zero_retention_is_ignored_for_full_attention(): + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10) + manager = make_kv_cache_manager( + kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=16, + retention_interval=0, + ) + assert manager.coordinator.retention_interval == 0 + + +def test_positive_retention_rejects_full_attention(): + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10) + with pytest.raises(ValueError, match="no sliding-window or Mamba"): + make_kv_cache_manager( + kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=16, + retention_interval=16, + ) + + +def test_hybrid_local_kv_retention_interval_survives_recycling(): """Verify retained local checkpoints are reused after block recycling.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "1024") hash_block_size = 4 kv_cache_config = KVCacheConfig( num_blocks=800, @@ -3286,6 +3313,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch): max_model_len=4096, enable_caching=True, hash_block_size=hash_block_size, + retention_interval=1024, ) def fill_request(request_id: str, token_offset: int) -> list[int]: @@ -3314,9 +3342,8 @@ def fill_request(request_id: str, token_offset: int) -> list[int]: assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256] -def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatch): +def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(): """Verify latest-only retention reuses only the replayable prompt boundary.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3348,6 +3375,7 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatc max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=0, ) token_ids = [i for i in range(16) for _ in range(block_size)] @@ -3388,14 +3416,13 @@ def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatc assert len(computed_blocks.blocks[1]) == 0 -def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch): +def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): """Verify MTP/EAGLE SWA retention keeps the extra proof block. EAGLE/MTP lookup matches one additional local block after the returned prefix and then drops it. Sparse retention must therefore cache the normal local tail at the latest replay boundary plus one extra SWA block. """ - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3428,6 +3455,7 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch): max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=0, use_eagle=True, ) @@ -3824,12 +3852,11 @@ def test_cache_hit_local_and_external_two_groups_preempt_and_reallocate(): assert manager.get_blocks("test").get_block_ids() != ([], []) -def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(monkeypatch): - """Default path (no retention): freeing an SWA request must place its +def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(): + """Dense retention: freeing an SWA request must place its uncached scratch blocks at the front of the free queue (recycled first) and keep its cached checkpoint blocks at the back (retained for prefix hits). This split is always-on, independent of the retention interval.""" - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) block_size = 8 kv_cache_config = KVCacheConfig( num_blocks=100, @@ -3937,13 +3964,14 @@ def _make_pure_swa_manager(block_size, sliding_window, num_blocks=100, **kwargs) ) -def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch): +def test_pure_swa_retention_interval_caches_sparse_tails(): """Sparse retention must work for a pure-SWA single-group model, not just hybrid models: only the per-interval tails plus the latest replay tail are cached, and a replay still hits the latest replayable boundary.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64") block_size = 16 - manager = _make_pure_swa_manager(block_size, sliding_window=block_size) + manager = _make_pure_swa_manager( + block_size, sliding_window=block_size, retention_interval=64 + ) assert type(manager.coordinator).__name__ == "UnitaryKVCacheCoordinator" token_ids = [i for i in range(16) for _ in range(block_size)] @@ -3976,11 +4004,12 @@ def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch): assert num_computed == 240 -def test_pure_swa_retention_latest_only(monkeypatch): +def test_pure_swa_retention_latest_only(): """`=0` on a pure-SWA model keeps only the latest replay tail.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 16 - manager = _make_pure_swa_manager(block_size, sliding_window=block_size) + manager = _make_pure_swa_manager( + block_size, sliding_window=block_size, retention_interval=0 + ) token_ids = [i for i in range(16) for _ in range(block_size)] req = make_request("0", token_ids, block_size, sha256) @@ -4008,10 +4037,9 @@ def test_pure_swa_retention_latest_only(monkeypatch): assert num_computed == 240 -def test_pure_swa_retention_dense_default_caches_all(monkeypatch): - """With retention unset, a pure-SWA model must keep the dense behavior: +def test_pure_swa_dense_retention_caches_all(): + """With retention set to ``None``, a pure-SWA model keeps dense behavior: every block boundary is a potential hit, so all blocks are cached.""" - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) block_size = 16 manager = _make_pure_swa_manager(block_size, sliding_window=block_size) @@ -4037,7 +4065,7 @@ def test_pure_swa_retention_dense_default_caches_all(monkeypatch): def test_mamba_reachable_block_mask_sparsifies_retention(): - """Mamba state-snapshot retention: with VLLM_PREFIX_CACHE_RETENTION_INTERVAL + """Mamba state-snapshot retention: with a configured retention interval, the manager keeps one cached state per interval-sized segment (plus the latest replay boundary) instead of a snapshot per block, which is what lets a small attention block_size avoid Mamba dominating the KV pool.""" @@ -4063,7 +4091,7 @@ def retained(retention_interval, num_prompt_tokens=256, end_block=16): ) return None if m is None else {i for i, v in enumerate(m) if v} - # Dense default (None) -> no mask, every block cached (unchanged behavior). + # Dense retention (None) -> no mask, every block cached. assert retained(None) is None # interval == block_size -> every block is a boundary -> stays dense. assert retained(block_size) is None @@ -4111,7 +4139,7 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, 100) == {5, 14} # Coexists with segment tails (interval 64 -> {3,7,11,15} + replay 14). assert retained(64, 96) == {3, 5, 7, 11, 14, 15} - # Dense default ignores the hint (nothing to sparsify). + # Dense retention ignores the hint (nothing to sparsify). assert retained(None, 96) is None # Out-of-range boundary is a no-op (only replay 14 remains). assert retained(0, 16 * block_size * 2) == {14} @@ -4120,14 +4148,13 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, None) == {14} -def test_mamba_shared_prefix_survives_zero_retention(monkeypatch): +def test_mamba_shared_prefix_survives_zero_retention(): """Manager-level check of the full wiring: a pinned shared-prefix boundary (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style detection) keeps its Mamba state block cached under - ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``, which otherwise retains only the + ``prefix_cache_retention_interval=0``, which otherwise retains only the end-of-prompt replay boundary. Without this, a shared prefix (junction before ``num_prompt``) would be recomputed by every sharing request.""" - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0") block_size = 16 # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. @@ -4140,6 +4167,7 @@ def cached_mamba_blocks(shared_prefix_boundary): max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=0, ) req = make_request("r", token_ids, block_size, sha256) req.shared_prefix_boundary = shared_prefix_boundary @@ -4163,24 +4191,21 @@ def cached_mamba_blocks(shared_prefix_boundary): assert cached_mamba_blocks(96) == {5, 14} -def test_mamba_shared_prefix_reuse_under_zero_retention(monkeypatch): +def test_mamba_shared_prefix_reuse_under_zero_retention(): """Full cross-request Marconi flow: a partial shared prefix cached by the detecting request must stay reusable by a later request under - ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``. Without the pin the junction is + ``prefix_cache_retention_interval=0``. Without the pin the junction is masked out and the later request misses; with it (and under dense) the reuse is preserved.""" block_size = 16 def last_req_hit(retention, pin): - if retention is None: - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) - else: - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", str(retention)) manager = make_kv_cache_manager( _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=retention, ) shared = [7 for _ in range(2 * block_size)] # 2-block shared prefix @@ -4261,23 +4286,20 @@ def retained(retention, boundary, window, end_block=16): assert retained(0, 0, block_size) == {14} -def test_swa_shared_prefix_reuse_under_zero_retention(monkeypatch): +def test_swa_shared_prefix_reuse_under_zero_retention(): """SWA cross-request analog: a partial shared prefix's sliding-window tail - must stay reusable under ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0``. Without + must stay reusable under ``prefix_cache_retention_interval=0``. Without the pin the junction window is masked out and a later request misses; with it (and under dense) reuse is preserved.""" block_size = 16 def last_req_hit(retention, pin): - if retention is None: - monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False) - else: - monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", str(retention)) manager = make_kv_cache_manager( _make_hybrid_kv_cache_config(block_size, 200, ["full", "sliding_window"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, + retention_interval=retention, ) shared = [7 for _ in range(4 * block_size)] # 4-block shared prefix diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index c8c6fc85480a..c0932d975c18 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -3701,6 +3701,7 @@ def test_mamba_align_eagle_schedules_encoder_at_boundary(): ) scheduler.need_mamba_block_aligned_split = True scheduler.use_eagle = True + scheduler.num_prefill_lookahead = 1 scheduler.max_num_encoder_input_tokens = 2048 scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=2048) @@ -5341,8 +5342,9 @@ def test_free_encoder_inputs_defers_for_eagle_lookahead(): worker-side token-embedding fallback is only a backstop.""" scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") # create_scheduler only builds ngram spec configs; force the eagle path that - # _free_encoder_inputs keys off (self.use_eagle). + # _free_encoder_inputs keys off (its read-ahead deferral). scheduler.use_eagle = True + scheduler.num_prefill_lookahead = 1 mm_positions = [[PlaceholderRange(offset=50, length=100)]] request = create_requests( num_requests=1, diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index f2a54c55c809..a01214a40d29 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -475,7 +475,9 @@ def _make_video_mm_kwargs( # --------------------------------------------------------------------------- -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) class TestEncoderCudaGraphCaptureReplay: def setup_method(self): self.device = torch.device("cuda:0") @@ -758,7 +760,9 @@ def test_video_model_returns_video_for_video_kwargs(self): _VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" +) class TestEncoderCudaGraphVideoReplay: def setup_method(self): self.device = torch.device("cuda:0") diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 4381e2c8d1fd..d8b42b9abfc4 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager @@ -873,6 +874,10 @@ def get_mamba_prefix_cache_step_configs( def _run_mamba_prefix_cache_mrv1( monkeypatch: pytest.MonkeyPatch, async_scheduling: bool ): + # This test patches the V1 model runner, so pin V1 explicitly: MoE/hybrid + # models like Qwen3-Next now default to the V2 runner. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + envs.disable_envs_cache() global async_scheduling_mode async_scheduling_mode = async_scheduling run_ref_mamba_state_in_subprocess() @@ -977,7 +982,10 @@ def temporal_states(model_state, block_tables, kv_cache_config): yield forward_context[layer_name].kv_cache[-1], block_table def temporal_block(temporal_state, block_table, col): - return temporal_state[int(block_table[0, col].item())] + # Resolving the block id for assertions is a deliberate D2H. + with gpu_sync_allowed(): + block_id = int(block_table[0, col].item()) + return temporal_state[block_id] def wrapped_preprocess_state( self: MambaHybridModelState, @@ -1001,19 +1009,24 @@ def wrapped_preprocess_state( self, input_batch, block_tables, kv_cache_config, num_computed_tokens ) if cur_step_action is not None: - req_idx = int(input_batch.idx_mapping[0].item()) - src_col = int(self._mamba_src_col_gpu[req_idx].item()) - off = int(self._mamba_src_off_gpu[req_idx].item()) - dst = int(self._mamba_state_idx_gpu[req_idx].item()) + # Reading the GPU-side copy state back to assert on it is a + # deliberate D2H. + with gpu_sync_allowed(): + req_idx = int(input_batch.idx_mapping[0].item()) + src_col = int(self._mamba_src_col_gpu[req_idx].item()) + off = int(self._mamba_src_off_gpu[req_idx].item()) + dst = int(self._mamba_state_idx_gpu[req_idx].item()) actual = (-1, -1) if src_col < 0 or src_col == dst else (src_col + off, dst) assert actual == expected, ( f"V2 align preprocess copy: expected={expected}, " f"actual={actual}, {cur_step_action=}" ) - for temporal, bt, src_state in snapshots: - torch.testing.assert_close( - temporal_block(temporal, bt, expected[1]), src_state - ) + # Comparing device tensors for the assertion is a deliberate D2H. + with gpu_sync_allowed(): + for temporal, bt, src_state in snapshots: + torch.testing.assert_close( + temporal_block(temporal, bt, expected[1]), src_state + ) return ret def wrapped_postprocess_state( @@ -1044,10 +1057,12 @@ def wrapped_postprocess_state( ret = original_postprocess_state( self, idx_mapping, num_sampled, num_computed_tokens ) - for temporal, bt, src_state in snapshots: - torch.testing.assert_close( - temporal_block(temporal, bt, expected[1]), src_state - ) + # Comparing device tensors for the assertion is a deliberate D2H. + with gpu_sync_allowed(): + for temporal, bt, src_state in snapshots: + torch.testing.assert_close( + temporal_block(temporal, bt, expected[1]), src_state + ) return ret def wrapped_execute_model( @@ -1068,10 +1083,10 @@ def wrapped_execute_model( ret = original_execute_model(self, scheduler_output, *args, **kwargs) if cur_step_action is not None and self.execute_model_state is not None: input_batch = self.execute_model_state.input_batch - assert ( - cur_step_action.num_computed_tokens_start - == input_batch.positions[input_batch.query_start_loc[0]].item() - ) + # Reading positions back to assert on them is a deliberate D2H. + with gpu_sync_allowed(): + start_pos = input_batch.positions[input_batch.query_start_loc[0]].item() + assert cur_step_action.num_computed_tokens_start == start_pos return ret def fake_sample( diff --git a/tests/v1/ec_connector/unit/test_ec_output_aggregator.py b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py new file mode 100644 index 000000000000..4d5b7dae1037 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_ec_output_aggregator.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ECOutputAggregator.""" + +import pytest + +from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator +from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + KVConnectorOutput, + ModelRunnerOutput, +) + +pytestmark = pytest.mark.cpu_test + + +class FakeWorkerMeta(ECConnectorWorkerMetadata): + """Records merge order. `aggregate` returns a new object, as the base class + declares: an aggregator discarding the return value would lose the merge. + """ + + def __init__(self, saves: list[str]): + self.saves = saves + + def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta": + return FakeWorkerMeta(self.saves + other.saves) + + +def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=[], req_id_to_index={}, ec_connector_output=ec_output + ) + + +def test_aggregate_folds_every_rank_onto_output_rank(): + """EC work done on any rank reaches the scheduler via output_rank's output. + + The middle rank reports no worker metadata: it must neither seed nor clobber + the accumulator. + """ + outputs = [ + _worker_output( + ECConnectorOutput( + finished_sending={"mm0"}, + ec_connector_worker_meta=FakeWorkerMeta(["mm0"]), + ) + ), + _worker_output(ECConnectorOutput(finished_recving={"mm1"})), + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm2"])) + ), + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=2) + + assert result is outputs[2] + assert result.ec_connector_output.finished_sending == {"mm0"} + assert result.ec_connector_output.finished_recving == {"mm1"} + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"] + + +def test_aggregate_leaves_no_ec_output_when_no_worker_reported(): + """Empty per-worker reports must not reach the scheduler as an empty object.""" + outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())] + + result = ECOutputAggregator().aggregate(outputs, output_rank=0) + + assert result is outputs[0] + assert result.ec_connector_output is None + assert ECOutputAggregator().aggregate([None], output_rank=0) is None + + +def test_aggregate_does_not_write_through_the_shared_empty_output(): + """A rank with nothing to report yields the shared empty output singleton. + + Folding another rank's metadata onto it must not write through to the + module-level object, which every later step would then carry. + """ + outputs = [ + _worker_output( + ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"])) + ), + EMPTY_MODEL_RUNNER_OUTPUT, + ] + + result = ECOutputAggregator().aggregate(outputs, output_rank=1) + + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None + assert result is not EMPTY_MODEL_RUNNER_OUTPUT + assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"] + + +def test_chaining_with_kv_aggregator_preserves_both_outputs(): + """MultiprocExecutor chains both aggregators and keeps only the last result, + so each must merge onto the same output_rank output rather than replace it. + """ + outputs = [ + _worker_output(ECConnectorOutput(finished_sending={"mm0"})), + _worker_output(None), + ] + outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7}) + + result = None + for aggregator in ( + KVOutputAggregator(expected_finished_count=1), + ECOutputAggregator(), + ): + result = aggregator.aggregate(outputs, output_rank=1) + + assert result is outputs[1] + assert result.kv_connector_output.invalid_block_ids == {7} + assert result.ec_connector_output.finished_sending == {"mm0"} diff --git a/tests/v1/ec_connector/unit/test_worker_ec_connector.py b/tests/v1/ec_connector/unit/test_worker_ec_connector.py new file mode 100644 index 000000000000..3dcad1e50ae5 --- /dev/null +++ b/tests/v1/ec_connector/unit/test_worker_ec_connector.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the V2 GPU model runner's EC connector wrapper.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.distributed.ec_transfer.ec_connector.base import ( + ECConnectorBase, + ECConnectorMetadata, +) +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT +from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector + +pytestmark = pytest.mark.cpu_test + +WORKER_META = object() + + +def _scheduler_output() -> SimpleNamespace: + return SimpleNamespace( + ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset() + ) + + +def _connector( + encoder_cache: dict | None = None, + is_producer: bool = True, + is_consumer: bool = False, +) -> tuple[ActiveECConnector, MagicMock]: + fake = MagicMock(spec=ECConnectorBase) + fake.is_producer = is_producer + fake.is_consumer = is_consumer + fake.get_finished.return_value = (None, None) + fake.build_connector_worker_meta.return_value = WORKER_META + with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake): + return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake + + +@pytest.mark.parametrize( + ("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)] +) +def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer): + """An ec_both node is also a producer: it must offload what it just computed.""" + encoder_cache = {"mm_old": None} + connector, fake = _connector(encoder_cache, is_producer, is_consumer) + + with connector.maybe_get_output(_scheduler_output()): + encoder_cache["mm_new"] = None + + saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list] + assert saved == (["mm_new"] if is_producer else []) + assert fake.start_load_caches.called == is_consumer + + +def test_worker_meta_is_reported_on_context_exit(): + """Reported in the finally block, so is_empty() sees it only after the exit.""" + connector, fake = _connector() + + with connector.maybe_get_output(_scheduler_output()) as output: + assert output.ec_connector_worker_meta is None + + assert output.ec_connector_worker_meta is WORKER_META + assert fake.clear_connector_metadata.called + + +def test_no_forward_reports_without_running_the_model(): + connector, _ = _connector() + + output = connector.no_forward(_scheduler_output()) + + assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META + + empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output()) + assert empty is EMPTY_MODEL_RUNNER_OUTPUT diff --git a/tests/v1/engine/test_engine_args.py b/tests/v1/engine/test_engine_args.py index c03c244bb1f8..b1ba1adcfd0a 100644 --- a/tests/v1/engine/test_engine_args.py +++ b/tests/v1/engine/test_engine_args.py @@ -18,6 +18,7 @@ def test_prefix_caching_from_cli(): assert vllm_config.cache_config.enable_prefix_caching, ( "V1 turns on prefix caching by default." ) + assert vllm_config.cache_config.prefix_cache_retention_interval == 0 # Turn it off possible with flag. args = parser.parse_args(["--no-enable-prefix-caching"]) @@ -47,6 +48,10 @@ def test_prefix_caching_from_cli(): with pytest.raises(ArgumentError): args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"]) + args = parser.parse_args(["--prefix-cache-retention-interval", "64"]) + vllm_config = EngineArgs.from_cli_args(args=args).create_engine_config() + assert vllm_config.cache_config.prefix_cache_retention_interval == 64 + @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed") def test_prefix_caching_xxhash_from_cli(): @@ -99,3 +104,20 @@ def test_mm_prefix_lm_raises_batched_tokens_floor(): vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) assert vllm_config.scheduler_config.max_num_batched_tokens >= 2496 + + +def test_data_parallel_start_rank_zero_infers_hybrid_lb(): + """An explicit --data-parallel-start-rank 0 must be treated the same as + any other explicit start rank when inferring hybrid LB mode, not as + "unset" (regression test for a truthiness-vs-`is not None` bug). + """ + engine_args = EngineArgs( + model="facebook/opt-125m", + data_parallel_size=4, + data_parallel_size_local=2, + data_parallel_start_rank=0, + ) + vllm_config = engine_args.create_engine_config(UsageContext.OPENAI_API_SERVER) + + assert vllm_config.parallel_config.data_parallel_hybrid_lb is True + assert vllm_config.parallel_config.data_parallel_rank == 0 diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 64adf7a8b3c8..74f7b803b890 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -1066,6 +1066,7 @@ def test_kv_cache_events( model=model_name, enforce_eager=True, enable_prefix_caching=True, + prefix_cache_retention_interval=None, block_size=block_size, ) engine_args.kv_events_config = publisher_config @@ -1292,6 +1293,7 @@ def create_mock_executor(vllm_config): mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}] mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024] mock_executor.initialize_from_config.return_value = None + mock_executor.supports_draft_weight_updates.return_value = False return mock_executor diff --git a/tests/v1/executor/test_executor.py b/tests/v1/executor/test_executor.py index c529c3204d50..a21a8dafd328 100644 --- a/tests/v1/executor/test_executor.py +++ b/tests/v1/executor/test_executor.py @@ -9,6 +9,7 @@ import pytest +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.sampling_params import SamplingParams @@ -98,6 +99,7 @@ def collective_rpc( non_block: bool = False, unique_reply_rank: int | None = None, kv_output_aggregator: KVOutputAggregator = None, + ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any | list[Any] | Future[Any | list[Any]]: # Drop marker to show that this was run with open(".marker", "w"): @@ -110,6 +112,7 @@ def collective_rpc( non_block, unique_reply_rank, kv_output_aggregator, + ec_output_aggregator, ) diff --git a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py index 0c0f9f1f8998..e5cbd977b6d8 100644 --- a/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py +++ b/tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py @@ -55,6 +55,7 @@ def _run_engine_core_handshake( class _FakeScheduler: def __init__(self, **kwargs: Any) -> None: self.connector = connector + self.ec_connector = None def get_kv_connector(self) -> KVConnectorBase_V1: return connector diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 6c8af7fb73c1..32457f13d525 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -228,12 +228,11 @@ def _fake_thread_init(*args, **kwargs): block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=hs, can_save=True, + store_job_id=1, ) - send_thread.add_stored_request("r0") - # Put the request in the queue so task_done() doesn't underflow. - send_thread.request_queue.put(save_req) - req = send_thread.request_queue.get() - send_thread._handle_request(req) + # add_request also queues the job, so task_done() doesn't underflow. + send_thread.add_request(save_req) + send_thread._handle_request(send_thread.request_queue.get()) # Point worker.store at the dict store (the worker constructor captured # the MagicMock; replace with the real dict store for lookup). @@ -499,15 +498,15 @@ def batch_put_from_multi_buffers(self, keys, addrs, sizes, *a, **k): can_save=True, num_prompt_tokens=12, partial_tail_offloads=[(1, 7, 12)], + store_job_id=1, ) req.current_event = event - send.add_stored_request("r1") - send.request_queue.put(req) + send.add_request(req) send._handle_request(send.request_queue.get()) assert send.request_queue.qsize() == 0 assert store._data - assert send.stored_requests["r1"] == 0 + assert send.stored_requests["r1"] == set() event.synchronize.assert_called_once() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 13960c40340e..998a8408e651 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -5,12 +5,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( LoadSpec, + MooncakeStoreWorkerMetadata, ReqMeta, RequestTracker, ) from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler import ( MooncakeStoreScheduler, ) +from vllm.v1.core.block_pool import BlockPool def _make_bare_scheduler( @@ -27,6 +29,12 @@ def _make_bare_scheduler( scheduler._unfinished_request_ids = {"req-0"} scheduler._unfinished_requests = {} scheduler._request_trackers = {} + scheduler._gpu_block_pool = BlockPool( + num_gpu_blocks=64, enable_caching=True, hash_block_size=hash_block_size + ) + scheduler._num_workers = 1 + scheduler._next_store_job_id = 0 + scheduler._pinned_saves = {} return scheduler @@ -64,6 +72,14 @@ def _make_preemption_scheduler_output(): ) +def _make_worker_output(completed_saves: dict[int, int]) -> SimpleNamespace: + return SimpleNamespace( + kv_connector_worker_meta=MooncakeStoreWorkerMetadata( + completed_saves=completed_saves + ) + ) + + def _add_unfinished_request( scheduler: MooncakeStoreScheduler, *, @@ -133,7 +149,7 @@ def test_cached_request_without_spec_decode_keeps_current_step_save_overlap(): assert tracker.num_saved_tokens == 48 -def test_preemption_resets_tracker_before_request_finished(): +def test_preemption_resets_tracker(): scheduler = _make_bare_scheduler() _add_unfinished_request( scheduler, @@ -152,8 +168,6 @@ def test_preemption_resets_tracker_before_request_finished(): assert tracker.token_ids is None assert tracker.has_pending_offload is False assert tracker.prefill_end_tokens == 0 - request = SimpleNamespace(request_id="req-0") - assert scheduler.request_finished(request, ([0, 1],)) == (False, None) def test_preemption_clears_stale_load_state(): @@ -215,7 +229,7 @@ def test_pending_load_does_not_co_queue_save(): # enqueue a save in the same scheduling step. Co-queuing both produces a # recv+send pair for the same req_id, and the scheduler's # _update_from_kv_xfer_finished then trips `assert req_id in self.requests` - # when both completions land for the delay-freed request. + # when a completion lands for a request it has already dropped. scheduler = _make_bare_scheduler() _make_pending_load_unfinished_request( scheduler, @@ -238,9 +252,7 @@ def test_pending_load_does_not_co_queue_save(): # Load is still issued as planned. assert req_meta.load_spec is not None assert req_meta.load_spec.can_load is True - # And the tracker's saved-tokens watermark stays at 0 so request_finished - # later sees `num_saved_tokens <= 0` and frees immediately rather than - # waiting for a finished_sending that will never come. + # And the save watermark does not advance for a save that was never queued. tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 @@ -661,30 +673,34 @@ def test_disabled_lookup_reports_no_hit_without_querying_client(): assert scheduler.load_specs == {} -def test_pending_partial_tail_emits_offload_only_reqmeta(): - # A sub-block prompt never produces a block-aligned save, so the partial- - # tail offload arriving this step is emitted as an offload-only ReqMeta - # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so - # the worker skips the normal save). Pending-offload state delays the free - # without advancing the normal-save watermark before the put succeeds. - scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) +def _add_pending_partial_tail_request( + scheduler: MooncakeStoreScheduler, + *, + num_tokens: int, + block_hashes: list[bytes], + block_ids: tuple[list[int], ...], +) -> SimpleNamespace: + """Register a sub-block request and return the step that offloads its tail. + + The CoW block holding the tail is block 7, which the core deliberately keeps + out of the request's block table. + """ request = SimpleNamespace( - all_token_ids=list(range(12)), - block_hashes=[b"h0", b"h1", b"h2"], + all_token_ids=list(range(num_tokens)), + block_hashes=block_hashes, num_output_placeholders=0, num_prompt_tokens=12, ) - scheduler._unfinished_requests["req-0"] = (request, ([0],)) + scheduler._unfinished_requests["req-0"] = (request, block_ids) scheduler._request_trackers["req-0"] = RequestTracker( req_id="req-0", - token_len=12, - allocated_block_ids=([0],), + token_len=num_tokens, + allocated_block_ids=block_ids, num_saved_tokens=0, - token_ids=list(range(12)), - prefill_end_tokens=12, + token_ids=list(range(num_tokens)), + prefill_end_tokens=num_tokens, ) - - out = SimpleNamespace( + return SimpleNamespace( finished_req_ids=set(), preempted_req_ids=set(), scheduled_new_reqs=[], @@ -699,6 +715,21 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): partial_tail_offloads={"req-0": [(1, 7, 12)]}, ) + +def test_pending_partial_tail_emits_offload_only_reqmeta(): + # A sub-block prompt never produces a block-aligned save, so the partial- + # tail offload arriving this step is emitted as an offload-only ReqMeta + # (can_save=True so it takes the normal enqueue path, token_len_chunk=0 so + # the worker skips the normal save), without advancing the normal-save + # watermark before the put succeeds. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + out = _add_pending_partial_tail_request( + scheduler, + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + block_ids=([0],), + ) + meta = scheduler.build_connector_meta(out) assert len(meta.requests) == 1 @@ -712,42 +743,16 @@ def test_pending_partial_tail_emits_offload_only_reqmeta(): tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 0 assert tracker.has_pending_offload is True - request = SimpleNamespace(request_id="req-0") - assert scheduler.request_finished(request, ([0],)) == (True, None) def test_resumed_partial_tail_uses_handoff_boundary(): scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) - request = SimpleNamespace( - all_token_ids=list(range(20)), + # Resumption replays prompt + previously generated tokens. + out = _add_pending_partial_tail_request( + scheduler, + num_tokens=20, block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4"], - num_output_placeholders=0, - num_prompt_tokens=12, - ) - scheduler._unfinished_requests["req-0"] = (request, ([0, 1],)) - scheduler._request_trackers["req-0"] = RequestTracker( - req_id="req-0", - token_len=20, - allocated_block_ids=([0, 1],), - num_saved_tokens=0, - token_ids=list(range(20)), - # Resumption replays prompt + previously generated tokens. - prefill_end_tokens=20, - ) - - out = SimpleNamespace( - finished_req_ids=set(), - preempted_req_ids=set(), - scheduled_new_reqs=[], - scheduled_cached_reqs=SimpleNamespace( - req_ids=[], - new_block_ids=[], - num_computed_tokens=[], - resumed_req_ids=set(), - ), - num_scheduled_tokens={}, - scheduled_spec_decode_tokens={}, - partial_tail_offloads={"req-0": [(1, 7, 12)]}, + block_ids=([0, 1],), ) meta = scheduler.build_connector_meta(out) @@ -791,3 +796,69 @@ def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary(): tracker = scheduler._request_trackers["req-0"] assert tracker.num_saved_tokens == 48 assert tracker.has_pending_offload is True + + +def test_partial_tail_cow_block_is_referenced_for_the_job(): + # The CoW block a partial-tail offload reads is deliberately kept out of the + # request block table, so it is absent from ReqMeta.block_ids. The worker + # DMAs out of it just as asynchronously, so it needs its own reference. + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True) + out = _add_pending_partial_tail_request( + scheduler, + num_tokens=12, + block_hashes=[b"h0", b"h1", b"h2"], + block_ids=([0],), + ) + pool = scheduler._gpu_block_pool + + meta = scheduler.build_connector_meta(out) + + store_job_id = meta.requests[0].store_job_id + # It leads the list, as in `pop_blocks_for_free`, so that the reversed free + # puts it last in eviction priority. + assert scheduler._pinned_saves[store_job_id][0] == [7, 0] + assert pool.blocks[7].ref_cnt == 1 + + scheduler.update_connector_output(_make_worker_output({store_job_id: 1})) + assert pool.blocks[7].ref_cnt == 0 + + +def test_store_job_blocks_are_released_once_every_rank_reports(): + # Every rank DMAs the job's blocks on its own, so the reference can only be + # dropped once the last of them reports. Until then the engine has to keep + # stepping: a completion only reaches the scheduler as worker metadata + # attached to a step, and a finishing request no longer defers its own free. + scheduler = _make_bare_scheduler() + scheduler._num_workers = 2 + _add_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[b"h0", b"h1", b"h2"], + prefill_end_tokens=48, + ) + pool = scheduler._gpu_block_pool + assert scheduler.has_pending_push_work() is False + + meta = scheduler.build_connector_meta( + _make_scheduler_output(scheduled_spec_tokens=None) + ) + store_job_id = meta.requests[0].store_job_id + assert pool.blocks[2].ref_cnt == 1 + assert scheduler.has_pending_push_work() is True + + scheduler.update_connector_output(_make_worker_output({store_job_id: 1})) + assert pool.blocks[2].ref_cnt == 1 + assert scheduler.has_pending_push_work() is True + + scheduler.update_connector_output(_make_worker_output({store_job_id: 1})) + assert pool.blocks[2].ref_cnt == 0 + assert scheduler.has_pending_push_work() is False + + +def test_worker_metadata_aggregates_completions_across_ranks(): + # The engine merges each rank's metadata before the scheduler sees it, so a + # job that every rank finished in one step arrives as a single count. + merged = MooncakeStoreWorkerMetadata(completed_saves={1: 1}).aggregate( + MooncakeStoreWorkerMetadata(completed_saves={1: 1, 2: 1}) + ) + assert merged.completed_saves == {1: 2, 2: 1} diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index 034297100a7a..50d2ee16149e 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib +import itertools import json import logging import math @@ -156,6 +158,17 @@ def _make_load_req( ) +_TEST_SAVE_SEQ = itertools.count(1) + + +def _run_store_req(thread, req_meta: ReqMeta) -> None: + """Register, enqueue and run a store job the way the worker does.""" + if req_meta.store_job_id is None: + req_meta.store_job_id = next(_TEST_SAVE_SEQ) + thread.add_request(req_meta) + thread._handle_request(thread.request_queue.get()) + + def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta: return ReqMeta( req_id=req_id, @@ -234,7 +247,9 @@ def _make_vllm_config( ) -def _make_kv_cache_config(*, block_size: int = 16) -> object: +def _make_kv_cache_config( + *, block_size: int = 16, prefix_cache_retention_interval: int | None = 0 +) -> object: """Minimal single-group KVCacheConfig for topology tests.""" from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -249,6 +264,7 @@ def _make_kv_cache_config(*, block_size: int = 16) -> object: num_blocks=10, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)], + prefix_cache_retention_interval=prefix_cache_retention_interval, ) @@ -385,27 +401,23 @@ def test_store_sending_thread_skips_request_during_cpu_pressure(): ] thread = _make_store_sending_thread(store) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert thread._store_pressure_active is True assert "req-a" in thread._skip_store_requests assert store.batch_put_from_multi_buffers.call_count == 1 - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"])) + _run_store_req(thread, _make_store_req("req-a", [b"a2", b"a3"])) assert store.batch_put_from_multi_buffers.call_count == 1 - thread.add_stored_request("req-b") - thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"])) + _run_store_req(thread, _make_store_req("req-b", [b"b0", b"b1"])) assert thread._store_pressure_active is False assert "req-a" not in thread._skip_store_requests assert store.batch_put_from_multi_buffers.call_count == 2 - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"])) + _run_store_req(thread, _make_store_req("req-a", [b"a4", b"a5"])) assert store.batch_put_from_multi_buffers.call_count == 3 @@ -418,8 +430,7 @@ def test_store_sending_thread_records_mooncake_metrics(): stats = MooncakeStoreConnectorStats() thread._record_operation_cb = stats.record_operation - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert len(stats.data["save_exists"]) == 1 assert stats.data["save_exists"][0]["num_keys"] == 2 @@ -497,16 +508,16 @@ def test_store_sending_thread_delta_saves_only_new_full_attention_chunks(): store.batch_put_from_multi_buffers.return_value = [256, 256] thread = _make_store_sending_thread(store) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 32 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -523,16 +534,16 @@ def test_store_sending_thread_delta_strides_with_local_phase(): store.batch_put_from_multi_buffers.return_value = [256] thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 16 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -550,15 +561,15 @@ def test_tp_sharded_group_saves_every_block_on_every_rank(): thread = _make_store_sending_thread(store, tp_rank=0, put_step=2) thread.group_put_steps = [1] - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -576,15 +587,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): thread._store_pressure_active = True thread._skip_store_requests.add("req-a") - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=16, block_ids=([0],), block_hashes=[b"a0"], can_save=True, - ) + ), ) store.batch_is_exist.assert_not_called() @@ -595,15 +606,15 @@ def test_store_sending_thread_retries_skipped_range_after_pressure(): # The next batch resumes from offset 0, re-covering the chunk skipped under # pressure (chunk 0) rather than losing it. - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -703,13 +714,11 @@ def test_partial_tail_offload_honors_active_pressure_gate(): thread = _make_partial_tail_send_thread(store) thread._store_pressure_active = True thread._skip_store_requests.add("req-a") - thread.add_stored_request("req-a") - thread._handle_request(_make_partial_tail_req([1, 2, 3])) + _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) store.batch_is_exist.assert_not_called() store.batch_put_from_multi_buffers.assert_not_called() - assert thread.stored_requests["req-a"] == 0 def test_partial_tail_put_failure_activates_pressure_gate(): @@ -717,17 +726,14 @@ def test_partial_tail_put_failure_activates_pressure_gate(): store.batch_is_exist.side_effect = lambda keys: [0] * len(keys) store.batch_put_from_multi_buffers.return_value = [256, -200, 256] thread = _make_partial_tail_send_thread(store) - thread.add_stored_request("req-a") - thread._handle_request(_make_partial_tail_req([1, 2, 3])) + _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) assert thread._store_pressure_active is True assert thread._skip_store_requests == {"req-a"} assert thread._saved_offset.get("req-a", 0) == 0 - assert thread.stored_requests["req-a"] == 0 - thread.add_stored_request("req-a") - thread._handle_request(_make_partial_tail_req([1, 2, 3])) + _run_store_req(thread, _make_partial_tail_req([1, 2, 3])) assert store.batch_put_from_multi_buffers.call_count == 1 @@ -737,16 +743,16 @@ def test_store_sending_thread_delta_start_rank_saves_second_local_chunk(): store.batch_put_from_multi_buffers.return_value = [256, 256] thread = _make_store_sending_thread(store, tp_rank=1, put_step=2) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 16 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3],), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -790,16 +796,16 @@ def test_store_sending_thread_delta_saves_only_new_masked_chunks(): token_databases=[db_full, db_masked], ) - thread.add_stored_request("req-a") thread._saved_offset["req-a"] = 32 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=64, block_ids=([0, 1, 2, 3], [0, 1, 2, 3]), block_hashes=[b"a0", b"a1", b"a2", b"a3"], can_save=True, - ) + ), ) keys = store.batch_is_exist.call_args.args[0] @@ -845,15 +851,15 @@ def test_store_sending_thread_prepares_missing_chunks_once_per_group(): coord=coord, token_databases=[db0, db1], ) - thread.add_stored_request("req-a") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="req-a", token_len_chunk=48, block_ids=([0, 1, 2], [2, 1, 0]), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - ) + ), ) db0.prepare_value.assert_not_called() @@ -881,46 +887,71 @@ def test_store_sending_thread_only_skips_on_no_available_handle(): ] thread = _make_store_sending_thread(store) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert thread._store_pressure_active is False assert "req-a" not in thread._skip_store_requests assert store.batch_put_from_multi_buffers.call_count == 1 - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"])) + _run_store_req(thread, _make_store_req("req-a", [b"a2", b"a3"])) assert store.batch_put_from_multi_buffers.call_count == 2 -def test_store_sending_thread_releases_pin_on_batch_is_exist_failure(): - # `batch_is_exist` raising must still decrement `stored_requests` so the - # scheduler can drop `delay_free_blocks` and release the pinned GPU blocks. +@pytest.mark.parametrize( + "failing_call", ["batch_is_exist", "batch_put_from_multi_buffers"] +) +def test_store_sending_thread_reports_job_when_store_raises(failing_call): + # A store that blows up must still report its job, or the scheduler keeps + # the job's GPU block references for the rest of the run. store = MagicMock() - store.batch_is_exist.side_effect = RuntimeError("mooncake down") + store.batch_is_exist.return_value = [0, 0] + getattr(store, failing_call).side_effect = RuntimeError("mooncake down") thread = _make_store_sending_thread(store) + req = _make_store_req("req-a", [b"a0", b"a1"]) - thread.add_stored_request("req-a") - with pytest.raises(RuntimeError): - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + with contextlib.suppress(RuntimeError): + _run_store_req(thread, req) - assert thread.stored_requests["req-a"] == 0 - store.batch_put_from_multi_buffers.assert_not_called() + assert thread.take_completed_saves() == {req.store_job_id: 1} -def test_store_sending_thread_releases_pin_on_batch_put_failure(): - # `batch_put_from_multi_buffers` raising is logged (not re-raised), and the - # pin must still be released through the finally block. - store = MagicMock() - store.batch_is_exist.return_value = [0, 0] - store.batch_put_from_multi_buffers.side_effect = RuntimeError("rdma error") - thread = _make_store_sending_thread(store) +def test_store_sending_thread_reports_job_when_the_preamble_raises(): + # The report has to survive a failure before the first store call too, so + # every dequeue leaves through the same exit. + thread = _make_store_sending_thread(MagicMock()) + req = _make_store_req("req-a", [b"a0", b"a1"]) + req.token_len_chunk = None # type: ignore[assignment] + + with contextlib.suppress(TypeError): + _run_store_req(thread, req) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + assert thread.take_completed_saves() == {req.store_job_id: 1} + thread.request_queue.task_done.assert_called_once() - assert thread.stored_requests["req-a"] == 0 + +def test_stale_store_job_cannot_touch_a_reused_request_id(): + # A preempted request resumes under its original id, so a job left over from + # the retired generation carries a req_id that now belongs to a live one. + thread = _make_store_sending_thread(MagicMock()) + stale = _make_store_req("req-a", [b"a0", b"a1"]) + stale.store_job_id = 1 + thread.add_request(stale) + thread._record_saved(stale, 32) + thread.delete_finished_stored_request("req-a") + + live = _make_store_req("req-a", [b"a0", b"a1"]) + live.store_job_id = 2 + thread.add_request(live) + + thread.finish_store_job(stale) + thread._record_saved(stale, 64) + thread._mark_request_skipped_for_pressure(stale) + + assert thread.is_live_store_job(live) + assert not thread.is_live_store_job(stale) + assert thread._saved_offset.get("req-a") is None + assert "req-a" not in thread._skip_store_requests def test_store_recving_thread_reports_failed_block_ids(): @@ -991,8 +1022,7 @@ def test_store_sending_thread_passes_replicate_config_when_preferred_segment_set replicate_config = SimpleNamespace(preferred_segment="10.0.0.7:50053") thread = _make_store_sending_thread(store, replicate_config=replicate_config) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 call_args = store.batch_put_from_multi_buffers.call_args.args @@ -1010,8 +1040,7 @@ def test_store_sending_thread_passes_default_replicate_config_when_no_preferred_ replicate_config = SimpleNamespace() thread = _make_store_sending_thread(store, replicate_config=replicate_config) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 call_args = store.batch_put_from_multi_buffers.call_args.args @@ -1067,8 +1096,7 @@ def test_store_sending_thread_sets_group_ids_when_enabled(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args @@ -1092,8 +1120,7 @@ def test_store_sending_thread_leaves_group_ids_unchanged_when_flag_disabled(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config @@ -1112,8 +1139,7 @@ def test_store_sending_thread_leaves_group_ids_unchanged_when_unsupported(): supports_group_ids=False, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) assert store.batch_put_from_multi_buffers.call_count == 1 assert store.batch_put_from_multi_buffers.call_args.args[3] is replicate_config @@ -1180,8 +1206,7 @@ def test_store_sending_thread_group_id_excludes_physical_sharding(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1210,8 +1235,7 @@ def test_store_sending_thread_multiple_segments_share_logical_group_id(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1260,8 +1284,7 @@ def test_store_sending_thread_group_ids_share_across_kv_cache_groups(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_multi_group_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_multi_group_store_req("req-a", [b"a0", b"a1"])) keys, addrs, sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == [ @@ -1296,8 +1319,7 @@ def test_store_sending_thread_group_ids_follow_missing_key_filter(): supports_group_ids=True, ) - thread.add_stored_request("req-a") - thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"])) + _run_store_req(thread, _make_store_req("req-a", [b"a0", b"a1"])) keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args assert keys == ["test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131"] @@ -1575,6 +1597,7 @@ def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch): "mlx5_0", "10.0.0.7:50051", ) + assert w.coord.retention_interval == 0 def test_requester_worker_init_prefers_local_hostname_override( @@ -1866,15 +1889,15 @@ def test_store_sending_thread_clamps_token_len_to_lcm(): # token_len_chunk=33 clamps to 32 → 2 chunks (not 3 with a partial 1-token chunk). thread = _make_store_sending_thread(store) - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=33, block_ids=([0, 1, 2],), block_hashes=[b"a0", b"a1", b"a2"], can_save=True, - ) + ), ) keys = store.batch_put_from_multi_buffers.call_args.args[0] @@ -1905,20 +1928,19 @@ def test_store_sending_thread_skips_when_token_len_below_lcm(): store, coord=coord, token_databases=[db], block_size=64 ) - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=32, block_ids=([0, 1],), block_hashes=[b"a0", b"a1"], can_save=True, - ) + ), ) store.batch_is_exist.assert_not_called() store.batch_put_from_multi_buffers.assert_not_called() - assert thread.stored_requests["r0"] == 0 def test_store_sending_thread_only_stores_swa_blocks_in_window(): @@ -1982,15 +2004,15 @@ def test_store_sending_thread_only_stores_swa_blocks_in_window(): ) hs = [bytes([i + 1]) * 4 for i in range(8)] - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=64, block_ids=([0, 1], list(range(8))), block_hashes=hs, can_save=True, - ) + ), ) keys = store.batch_put_from_multi_buffers.call_args.args[0] @@ -2057,16 +2079,16 @@ def test_store_sending_thread_delta_saves_only_new_swa_boundary_chunks(): ) hs = [bytes([i + 1]) * 4 for i in range(8)] - thread.add_stored_request("r0") thread._saved_offset["r0"] = 32 - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=64, block_ids=([0, 1], list(range(8))), block_hashes=hs, can_save=True, - ) + ), ) keys = store.batch_put_from_multi_buffers.call_args.args[0] @@ -2128,8 +2150,8 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): thread.enable_kv_event = True hs = [bytes([i + 1]) * 4 for i in range(4)] - thread.add_stored_request("r0") - thread._handle_request( + _run_store_req( + thread, ReqMeta( req_id="r0", token_len_chunk=32, @@ -2137,7 +2159,7 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata(): block_hashes=hs, can_save=True, token_ids=list(range(32)), - ) + ), ) full_event, swa_event = thread.get_kv_events() diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 244594d58e83..a37b12f62a11 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -433,6 +433,9 @@ def test_kv_transfer_handshake(dist_init): ) ) assert delay + # Pull connector advertises its transfer mode in kv_transfer_params so + # an external router can distinguish it from a push producer. + assert kv_connector_metadata["transfer_mode"] == "pull" # Decode connector will be able to create handshake with the prefill connector. decode_connector = NixlConnector( @@ -2906,6 +2909,42 @@ def test_speculative_attention_backend_not_in_compatibility_hash(): assert local_hash == remote_hash +@pytest.mark.skip_global_cleanup +def test_transfer_mode_changes_compatibility_hash(): + # push (WRITE) and pull (READ) connectors use incompatible transfer + # protocols, so their compatibility hashes must differ; identical modes + # must match. The default mode is pull. + config = create_vllm_config() + + pull_hash = compute_nixl_compatibility_hash( + config, "FLASH_ATTN", False, transfer_mode="pull" + ) + push_hash = compute_nixl_compatibility_hash( + config, "FLASH_ATTN", False, transfer_mode="push" + ) + + assert pull_hash != push_hash + assert pull_hash == compute_nixl_compatibility_hash( + config, "FLASH_ATTN", False, transfer_mode="pull" + ) + assert compute_nixl_compatibility_hash(config, "FLASH_ATTN", False) == pull_hash + + +@pytest.mark.skip_global_cleanup +def test_scheduler_advertises_transfer_mode(): + # Each scheduler advertises its transfer mode in kv_transfer_params so an + # external router can route pull (READ) vs push (WRITE) producers. + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import ( + NixlPullConnectorScheduler, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import ( + NixlPushConnectorScheduler, + ) + + assert NixlPullConnectorScheduler._TRANSFER_MODE == "pull" + assert NixlPushConnectorScheduler._TRANSFER_MODE == "push" + + @pytest.mark.parametrize( "mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat", [ diff --git a/tests/v1/sample/test_thinking_budget_state.py b/tests/v1/sample/test_thinking_budget_state.py new file mode 100644 index 000000000000..943734e8eefd --- /dev/null +++ b/tests/v1/sample/test_thinking_budget_state.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ThinkingBudgetStateHolder batch index moves.""" + +import torch + +from vllm.sampling_params import SamplingParams +from vllm.v1.sample.logits_processor.interface import ( + BatchUpdate, + MoveDirectionality, +) +from vllm.v1.sample.thinking_budget_state import ThinkingBudgetStateHolder + + +class _MockReasoningConfig: + reasoning_start_token_ids = [151667] + reasoning_end_token_ids = [151668] + + +def _make_holder() -> ThinkingBudgetStateHolder: + return ThinkingBudgetStateHolder( + _MockReasoningConfig(), + 8, + 0, + torch.device("cpu"), + False, + ) + + +def test_swap_budgeted_with_unbudgeted_clears_empty_side(): + """Asymmetric SWAP must not leave the empty index sharing state.""" + h = _make_holder() + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=[ + (0, SamplingParams(thinking_token_budget=5), None, []), + (1, SamplingParams(), None, []), + ], + moved=(), + ) + ) + assert list(h._state.keys()) == [0] + budget_state = h._state[0] + + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=(), + moved=[(0, 1, MoveDirectionality.SWAP)], + ) + ) + assert list(h._state.keys()) == [1] + assert h._state[1] is budget_state + assert h._state[1]["thinking_token_budget"] == 5 + + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=(), + moved=[(0, 1, MoveDirectionality.SWAP)], + ) + ) + assert list(h._state.keys()) == [0] + assert h._state[0] is budget_state + + +def test_swap_exchanges_two_budgeted_states(): + h = _make_holder() + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=[ + (0, SamplingParams(thinking_token_budget=3), None, []), + (1, SamplingParams(thinking_token_budget=7), None, []), + ], + moved=(), + ) + ) + b0 = h._state[0]["thinking_token_budget"] + b1 = h._state[1]["thinking_token_budget"] + h.sync_batch( + BatchUpdate( + batch_size=2, + removed=(), + added=(), + moved=[(0, 1, MoveDirectionality.SWAP)], + ) + ) + assert h._state[0]["thinking_token_budget"] == b1 + assert h._state[1]["thinking_token_budget"] == b0 diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 5fc88de9984a..99708f50f8f2 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/tests/v1/spec_decode/test_adaptive_verification.py @@ -9,6 +9,7 @@ from vllm.v1.worker.gpu.spec_decode.adaptive_verification import ( AdaptiveVerificationManager, ) +from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping def make_manager( @@ -168,3 +169,51 @@ def test_zero_budget_rebuilds_cpu_cu_num_logits(): assert cu_num_logits_np.dtype == scheduled_cu_num_logits.dtype # The prefill keeps its scheduled tokens; only drafts are dropped. assert np.array_equal(compacted, np.array([1, 1, 40], dtype=np.int32)) + + +def test_zero_budget_keeps_one_grammar_row_per_scheduled_draft(): + # The scheduler sizes the grammar bitmask from the *scheduled* drafts + # (len(drafts) + 1 rows per request), but a zero budget rewrites + # cu_num_logits_np to bonus-only. Deriving the bitmask -> logits mapping + # from those rewritten offsets drops rows and trips the + # `num_masks == len(mapping)` assert in apply_grammar_bitmask. + manager = make_manager( + np.array([[0.9, 0.9], [0.9, 0.9], [1.0, 1.0]], dtype=np.float32), + np.ones(64), + ) + manager.req_states.req_id_to_index["prefill"] = 2 + manager.req_states.num_computed_tokens_np = np.zeros(3, dtype=np.int32) + manager.req_states.prefill_len.np = np.array([0, 0, 60], dtype=np.int32) + manager._max_total_logits = 2 # < 3 requests * 1 bonus token + + scheduled_spec_decode_tokens = {"low": [1, 2], "high": [3, 4]} + manager.get_num_tokens( + {"low": 3, "high": 3, "prefill": 40}, scheduled_spec_decode_tokens + ) + assert manager._batch_budget[2] == 0 + + req_ids = ["low", "high", "prefill"] + num_draft_tokens_per_req = np.array([2, 2, 0], dtype=np.int32) + _, cu_num_logits_np = manager.compact_batch( + num_draft_tokens_per_req, + np.array([3, 3, 40], dtype=np.int32), + np.array([0, 3, 6, 7], dtype=np.int32), + ) + + mask_stride = manager.num_speculative_steps + manager.num_bonus_tokens + mapping = _build_grammar_mapping( + req_ids, + req_ids, + cu_num_logits_np, + num_draft_tokens_per_req, + manager.num_bonus_tokens, + mask_stride, + ) + + num_bitmask_rows = sum( + len(scheduled_spec_decode_tokens.get(req_id, ())) + 1 for req_id in req_ids + ) + assert len(mapping) == num_bitmask_rows + # (request, position) keys, so the kernel can mask rows the compacted + # device layout no longer has room for. + assert mapping == [0, 1, 2, 3, 4, 5, 6] diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py new file mode 100644 index 000000000000..16d6d8e516df --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( + prepare_dflash_inputs, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a CUDA device" +) + + +def _run_prepare( + *, + target_positions: list[int], + block_table_values: list[int], + cp_rank: int = 0, + cp_size: int = 1, + cp_interleave: int = 1, +): + device = torch.device("cuda") + max_num_reqs = 4 + max_num_tokens = 16 + num_speculative_steps = 3 + + input_buffers = SimpleNamespace( + input_ids=torch.full((max_num_tokens,), -1, dtype=torch.int32, device=device), + positions=torch.full((max_num_tokens,), -1, dtype=torch.int64, device=device), + query_start_loc=torch.full( + (max_num_reqs + 1,), -1, dtype=torch.int32, device=device + ), + seq_lens=torch.full((max_num_reqs,), -1, dtype=torch.int32, device=device), + ) + input_batch = SimpleNamespace( + num_reqs=1, + num_scheduled_tokens=np.array([4], dtype=np.int32), + positions=torch.tensor(target_positions, dtype=torch.int64, device=device), + query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device), + idx_mapping=torch.tensor([2], dtype=torch.int32, device=device), + ) + query_slot_mapping = torch.full( + (max_num_tokens,), -2, dtype=torch.int64, device=device + ) + context_positions = torch.full( + (max_num_tokens,), -1, dtype=torch.int64, device=device + ) + context_slot_mapping = torch.full( + (max_num_tokens,), -2, dtype=torch.int64, device=device + ) + sample_indices = torch.full( + (max_num_reqs * num_speculative_steps,), + -1, + dtype=torch.int64, + device=device, + ) + sample_pos = torch.full_like(sample_indices, -1) + sample_idx_mapping = torch.full( + sample_indices.shape, -1, dtype=torch.int32, device=device + ) + temperature = torch.zeros(max_num_reqs, dtype=torch.float32, device=device) + seeds = torch.zeros(max_num_reqs, dtype=torch.int64, device=device) + input_temperature = torch.tensor( + [0.0, 0.0, 1.0, 0.0], dtype=torch.float32, device=device + ) + input_seeds = torch.tensor([0, 0, 17, 0], dtype=torch.int64, device=device) + last_sampled = torch.tensor([0, 0, 99, 0], dtype=torch.int64, device=device) + next_prefill_tokens = torch.zeros_like(last_sampled) + block_table = torch.tensor([block_table_values], dtype=torch.int32, device=device) + + prepare_dflash_inputs( + input_buffers, + query_slot_mapping, + context_positions, + context_slot_mapping, + sample_indices, + sample_pos, + sample_idx_mapping, + temperature, + seeds, + input_batch, + torch.tensor([1], dtype=torch.int32, device=device), + torch.tensor([2], dtype=torch.int32, device=device), + last_sampled, + next_prefill_tokens, + input_temperature, + input_seeds, + block_table, + 4, + cp_rank, + cp_size, + cp_interleave, + 123, + num_speculative_steps, + num_speculative_steps, + max_num_reqs, + max_num_tokens, + 128, + sample_from_anchor=True, + ) + torch.accelerator.synchronize() + return SimpleNamespace( + input_buffers=input_buffers, + query_slot_mapping=query_slot_mapping.cpu(), + context_positions=context_positions.cpu(), + context_slot_mapping=context_slot_mapping.cpu(), + sample_indices=sample_indices.cpu(), + sample_pos=sample_pos.cpu(), + sample_idx_mapping=sample_idx_mapping.cpu(), + temperature=temperature.cpu(), + seeds=seeds.cpu(), + ) + + +def test_prepare_dflash_inputs_excludes_rejected_context_suffix(): + # Positions 10/11 use physical block 7. Rejected positions 12/13 would use + # block 8, but must be PAD context rather than contaminating draft KV. + out = _run_prepare( + target_positions=[10, 11, 12, 13], + block_table_values=[0, 0, 7, 8, 9, 10, 11, 12], + ) + + assert out.context_positions[:4].tolist() == [10, 11, 0, 0] + assert out.context_slot_mapping[:4].tolist() == [30, 31, PAD_SLOT_ID, PAD_SLOT_ID] + + # The replacement query starts immediately after the two valid rows and + # advances from the last accepted position (11). + assert out.input_buffers.input_ids[:3].cpu().tolist() == [99, 123, 123] + assert out.input_buffers.positions[:3].cpu().tolist() == [12, 13, 14] + assert out.query_slot_mapping[:3].tolist() == [32, 33, 34] + assert out.sample_indices[:3].tolist() == [0, 1, 2] + assert out.sample_pos[:3].tolist() == [13, 14, 15] + assert out.sample_idx_mapping[:3].tolist() == [2, 2, 2] + assert out.temperature[2].item() == 1.0 + assert out.seeds[2].item() == 17 + + +def test_prepare_dflash_inputs_excludes_rejected_context_suffix_with_dcp(): + out = _run_prepare( + target_positions=[10, 11, 12, 13], + block_table_values=[0, 7, 8, 9], + cp_rank=1, + cp_size=2, + cp_interleave=2, + ) + + assert out.context_positions[:4].tolist() == [10, 11, 0, 0] + assert out.context_slot_mapping[:4].tolist() == [28, 29, PAD_SLOT_ID, PAD_SLOT_ID] + assert out.query_slot_mapping[:3].tolist() == [PAD_SLOT_ID, PAD_SLOT_ID, 30] + + +def test_prepare_dflash_inputs_never_writes_the_null_block(): + # The valid context uses logical block 0 and the replacement query uses + # logical block 1. Both map to the null block and must remain unwritable. + out = _run_prepare( + target_positions=[2, 3, 4, 5], + block_table_values=[0, 0, 7, 8, 9, 10, 11, 12], + ) + + assert out.context_slot_mapping[:4].tolist() == [ + PAD_SLOT_ID, + PAD_SLOT_ID, + PAD_SLOT_ID, + PAD_SLOT_ID, + ] + assert out.query_slot_mapping[:3].tolist() == [ + PAD_SLOT_ID, + PAD_SLOT_ID, + PAD_SLOT_ID, + ] diff --git a/tests/v1/structured_output/test_validation.py b/tests/v1/structured_output/test_validation.py index 7ea60cc6609b..af417bc03e6f 100644 --- a/tests/v1/structured_output/test_validation.py +++ b/tests/v1/structured_output/test_validation.py @@ -5,7 +5,7 @@ import pytest from vllm.config import StructuredOutputsConfig -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import VLLMClientError, VLLMValidationError from vllm.sampling_params import SamplingParams, StructuredOutputsParams pytestmark = pytest.mark.cpu_test @@ -70,3 +70,93 @@ def test_degenerate_structured_outputs_rejected(structured_outputs, match): StructuredOutputsConfig(), tokenizer=object(), ) + + +@pytest.mark.parametrize( + "regex", + [ + "\x00", # a lone leading NUL + "\x00\x01\x02\x1f", # a NUL followed by other control chars + "[0-9]\x00", # an embedded NUL + ], +) +def test_regex_with_nul_byte_rejected(regex): + """A NUL byte is never meaningful in a structured-outputs regex and is not + handled by xgrammar's native regex converter. It must be rejected at request + validation in every backend mode (a clean 400), instead of reaching that + native code or silently falling back to another backend in the default + 'auto' mode.""" + params = SamplingParams(structured_outputs=StructuredOutputsParams(regex=regex)) + + # Rejected before backend selection, so it is a 400 even in 'auto' mode + # (which would otherwise catch the error and fall back to another backend). + with pytest.raises(VLLMValidationError, match="NUL"): + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(), + tokenizer=object(), + ) + + # The xgrammar backend also rejects it directly (defense in depth), before + # the pattern reaches the native from_regex call. + from vllm.v1.structured_output.backend_xgrammar import validate_xgrammar_grammar + + with pytest.raises(ValueError, match="NUL"): + validate_xgrammar_grammar(params) + + +INVALID_JSON_SCHEMA = {"type": "object", "properties": {"name": {"type": "str"}}} + + +@pytest.mark.parametrize( + "backend, structured_outputs", + [ + ("xgrammar", StructuredOutputsParams(json=INVALID_JSON_SCHEMA)), + ("outlines", StructuredOutputsParams(json=INVALID_JSON_SCHEMA)), + ("auto", StructuredOutputsParams(json=INVALID_JSON_SCHEMA)), + ("auto", StructuredOutputsParams(json='{"type": ')), + ("xgrammar", StructuredOutputsParams(grammar="not a grammar")), + ("guidance", StructuredOutputsParams(grammar="not a grammar")), + ("lm-format-enforcer", StructuredOutputsParams(grammar="not a grammar")), + ("outlines", StructuredOutputsParams(regex="(")), + ("guidance", StructuredOutputsParams(structural_tag='{"nope": 1}')), + ], +) +def test_unsupported_grammar_is_a_client_error(backend, structured_outputs): + """Only `VLLMClientError` survives `AsyncLLM.generate` untouched; anything else + is wrapped in `EngineGenerateError` and served as a 500 instead of a 400.""" + params = SamplingParams(structured_outputs=structured_outputs) + with pytest.raises(VLLMClientError): + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend=backend), + tokenizer=object(), + ) + + +@pytest.mark.parametrize( + "schema, expected_backend", + [ + # multipleOf is unsupported by xgrammar, patternProperties also by guidance. + ( + { + "type": "object", + "properties": {"n": {"type": "integer", "multipleOf": 2}}, + }, + "guidance", + ), + ( + {"type": "object", "patternProperties": {"^a": {"type": "string"}}}, + "outlines", + ), + ], +) +def test_auto_backend_falls_back_on_unsupported_schema(schema, expected_backend): + """`auto` falls back on rejection, so it must catch what the validators raise.""" + params = SamplingParams(structured_outputs=StructuredOutputsParams(json=schema)) + params._validate_structured_outputs( + _StubModelConfig(is_diffusion=False), + StructuredOutputsConfig(backend="auto"), + tokenizer=object(), + ) + assert params.structured_outputs._backend == expected_backend diff --git a/tests/v1/test_kv_cache_spec_registry.py b/tests/v1/test_kv_cache_spec_registry.py index 63834719c942..58c23ce285b1 100644 --- a/tests/v1/test_kv_cache_spec_registry.py +++ b/tests/v1/test_kv_cache_spec_registry.py @@ -34,7 +34,6 @@ SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, - TQFullAttentionSpec, UniformTypeKVCacheSpecs, get_kv_cache_spec_kind, ) @@ -85,7 +84,6 @@ def max_memory_usage_bytes(self, _) -> int: spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = { FullAttentionSpec: FullAttentionManager, - TQFullAttentionSpec: FullAttentionManager, MLAAttentionSpec: FullAttentionManager, HiddenStateCacheSpec: FullAttentionManager, SlidingWindowSpec: SlidingWindowManager, @@ -98,7 +96,6 @@ def max_memory_usage_bytes(self, _) -> int: spec_uniform_base_map: dict[type[KVCacheSpec], type[KVCacheSpec]] = { FullAttentionSpec: FullAttentionSpec, - TQFullAttentionSpec: FullAttentionSpec, MLAAttentionSpec: FullAttentionSpec, HiddenStateCacheSpec: FullAttentionSpec, SlidingWindowSpec: SlidingWindowSpec, @@ -113,13 +110,6 @@ def max_memory_usage_bytes(self, _) -> int: FullAttentionSpec: dict( block_size=64, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 ), - TQFullAttentionSpec: dict( - block_size=64, - num_kv_heads=8, - head_size=128, - dtype=torch.bfloat16, - tq_slot_size=256, - ), MLAAttentionSpec: dict( block_size=64, num_kv_heads=1, head_size=128, dtype=torch.bfloat16 ), @@ -262,7 +252,6 @@ def test_builtin_specs_are_uniform_with_same_spec_type(self, spec_cls): def test_full_attention_family_specs_are_uniform(self): specs = [ make_spec(FullAttentionSpec), - make_spec(TQFullAttentionSpec), make_spec(MLAAttentionSpec), make_spec(HiddenStateCacheSpec), make_spec(SinkFullAttentionSpec), diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 64632bb25c40..d93ae96e5438 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -5,10 +5,19 @@ import numpy as np import torch -from vllm.v1.outputs import LogprobsLists, LogprobsTensors +from vllm.platforms import current_platform +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + LogprobsLists, + LogprobsTensors, + ModelRunnerOutput, +) from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.sample.output import SamplingMaskTensors +DEVICE_TYPE = current_platform.device_type + def test_logprobs_tensors_cat(): first = LogprobsTensors( @@ -71,9 +80,9 @@ def test_sampling_mask_tensors_from_logits(): [3.0, 4.0, float("-inf")], [float("-inf"), 5.0, 6.0], ], - device="cuda", + device=DEVICE_TYPE, ), - num_sampled_tokens=torch.tensor([1, 0, 1], device="cuda"), + num_sampled_tokens=torch.tensor([1, 0, 1], device=DEVICE_TYPE), ) result = tensors.tolists(np.array([1, 0, 1])) @@ -85,9 +94,11 @@ def test_sampling_mask_tensors_from_logits(): def test_sampling_mask_matches_processed_top_k_top_p_support(): processed_logits = apply_top_k_top_p( - logits=torch.tensor([[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device="cuda"), - k=torch.tensor([3], device="cuda"), - p=torch.tensor([0.9], device="cuda"), + logits=torch.tensor( + [[6.0, 5.0, 4.0, 4.0, 4.0, 2.0, 1.0, 0.0]], device=DEVICE_TYPE + ), + k=torch.tensor([3], device=DEVICE_TYPE), + p=torch.tensor([0.9], device=DEVICE_TYPE), ) expected_token_ids = ( torch.isfinite(processed_logits[0]).nonzero().flatten().tolist() @@ -96,7 +107,7 @@ def test_sampling_mask_matches_processed_top_k_top_p_support(): tensors = SamplingMaskTensors.from_logits( processed_logits, - num_sampled_tokens=torch.tensor([1], device="cuda"), + num_sampled_tokens=torch.tensor([1], device=DEVICE_TYPE), ) result = tensors.tolists(np.array([1])) @@ -195,3 +206,14 @@ def test_slice_all_requests(self): assert len(sliced.logprob_token_ids) == 9 # All tokens assert sliced.logprob_token_ids == self.logprobsLists.logprob_token_ids assert sliced.cu_num_generated_tokens is None + + +def test_with_ec_conn_output_copies_shared_empty_output(): + """The shared empty output is copied, never written to.""" + ec_output = ECConnectorOutput(finished_sending={"mm_hash"}) + + result = ModelRunnerOutput.with_ec_conn_output(EMPTY_MODEL_RUNNER_OUTPUT, ec_output) + + assert result is not EMPTY_MODEL_RUNNER_OUTPUT + assert result.ec_connector_output is ec_output + assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None diff --git a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py index 37f1b0737f1a..d5a698433c0f 100644 --- a/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py +++ b/tests/v1/worker/test_dsv4_packed_zeroer_geometry.py @@ -93,6 +93,7 @@ def test_dsv4_packed_zeroer_geometry(): cache_dtype_str="fp8_ds_mla", alignment=576, model_version="deepseek_v4", + state_content_bytes=584, # >576 to allocate room for scales at back of page ) for _ in layer_names ] diff --git a/tests/v1/worker/test_gpu_bad_words.py b/tests/v1/worker/test_gpu_bad_words.py new file mode 100644 index 000000000000..150c4ed00711 --- /dev/null +++ b/tests/v1/worker/test_gpu_bad_words.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip( + "CUDA required for Model Runner V2 bad words tests", + allow_module_level=True, + ) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.bad_words import BadWordsState +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 + +# Committed tokens: prompt [5], output [10, 11]. Draft tokens: [12, 13]. +# The sampler passes input_ids gathered at logits_indices, so local position 0 +# holds the last committed token (11) and draft tokens start at position 1. +PROMPT_LEN = 1 +COMMITTED = [5, 10, 11] +INPUT_IDS = [11, 12, 13] +LOCAL_POS = [0, 1, 2] + + +def _make_state(bad_words_token_ids: list[list[int]]) -> tuple[BadWordsState, int]: + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + req_states.add_request( + req_id="req", + prompt_len=PROMPT_LEN, + all_token_ids=COMMITTED, + num_computed_tokens=len(COMMITTED), + max_tokens=32, + ) + req_states.apply_staged_writes() + + req_idx = req_states.req_id_to_index["req"] + state = BadWordsState(req_states) + state.add_request(req_idx, SamplingParams(_bad_words_token_ids=bad_words_token_ids)) + state.apply_staged_writes() + return state, req_idx + + +def _apply(bad_words_token_ids: list[list[int]]) -> torch.Tensor: + state, req_idx = _make_state(bad_words_token_ids) + num_logits = len(INPUT_IDS) + logits = torch.zeros((num_logits, VOCAB_SIZE), device=DEVICE) + idx_mapping_np = np.array([req_idx], dtype=np.intp) + expanded_idx_mapping = torch.tensor( + [req_idx] * num_logits, dtype=torch.int32, device=DEVICE + ) + state.apply_bad_words( + logits, + expanded_idx_mapping, + idx_mapping_np, + torch.tensor(INPUT_IDS, dtype=torch.int32, device=DEVICE), + torch.tensor(LOCAL_POS, dtype=torch.int32, device=DEVICE), + ) + return logits.cpu() + + +def test_v2_bad_words_prefix_inside_draft_tokens(): + """A prefix matching entirely within the draft tokens must mask the bad + word's last token at the draft position that completes the prefix.""" + out = _apply([[12, 13, 40]]) + expected = torch.zeros_like(out) + expected[2, 40] = -float("inf") + torch.testing.assert_close(out, expected) + + +def test_v2_bad_words_prefix_spanning_committed_and_draft_tokens(): + """A prefix spanning the committed/draft boundary must mask at the row + where the prefix completes, not one draft position later.""" + out = _apply([[11, 12, 30]]) + expected = torch.zeros_like(out) + expected[1, 30] = -float("inf") + torch.testing.assert_close(out, expected) + + +def test_v2_bad_words_no_spurious_match_from_last_committed_token(): + """The last committed token must not be double-counted as the first draft + token; [11, 11] never occurs in output [10, 11] + drafts [12, 13].""" + out = _apply([[11, 11, 50]]) + expected = torch.zeros_like(out) + torch.testing.assert_close(out, expected) + + +def test_v2_bad_words_committed_prefix(): + """Baseline: a fully committed prefix masks at the first row.""" + out = _apply([[10, 11, 60]]) + expected = torch.zeros_like(out) + expected[0, 60] = -float("inf") + torch.testing.assert_close(out, expected) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 2db024dc6b96..8bbe15bb2a4e 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -35,6 +35,10 @@ from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backend import MultipleOf +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend +from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import ( + ROCMAiterMLASparseBackend, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput @@ -291,6 +295,16 @@ def test_select_common_block_size_uses_largest_shared_int(): assert selected_size == 64 +def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch): + monkeypatch.setattr(current_platform, "is_rocm", lambda: True) + + selected_size = select_common_block_size( + 16, + [DeepseekV32IndexerBackend, ROCMAiterMLASparseBackend], + ) + assert selected_size == 16 + + def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids( dist_init, ): @@ -380,9 +394,13 @@ def test_select_common_block_size_no_valid_option(): def test_set_active_mm_loras_builds_tower_and_connector_mappings(): model = Mock() - model.get_num_mm_encoder_tokens.side_effect = lambda num_embeds: num_embeds + 1 + model.get_mm_lora_token_counts.side_effect = ( + lambda *, modality, mm_kwargs, num_mm_embeds: ( + num_mm_embeds + 1, + num_mm_embeds + 11, + ) + ) model.get_mm_mapping.return_value = SimpleNamespace(connector=True) - model.get_num_mm_connector_tokens.side_effect = lambda num_tokens: num_tokens + 10 lora_manager = Mock() lora_manager.supports_tower_connector_lora.return_value = True diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index fbeaa198d0d6..ebb4beb2a5a9 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -180,6 +180,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch): hidden_states=None, aux_hidden_states=None, finished_req_ids=set(), + ec_connector_output=None, routed_experts=None, num_tokens_across_dp=None, ) diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py new file mode 100644 index 000000000000..3fd1cb3b3c68 --- /dev/null +++ b/tests/v1/worker/test_gpu_sampler_flags.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for sampler flag tests", allow_module_level=True) + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +DEVICE = torch.device("cuda") +VOCAB_SIZE = 128 + + +class MockReasoningConfig: + reasoning_start_token_ids = [90] + reasoning_end_token_ids = [91] + natural_reasoning_end_token_ids = [91] + + +def _make_sampler() -> Sampler: + req_states = RequestState( + max_num_reqs=4, + max_model_len=64, + max_num_batched_tokens=16, + num_speculative_steps=1, + vocab_size=VOCAB_SIZE, + device=DEVICE, + ) + return Sampler( + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + reasoning_config=MockReasoningConfig(), + ) + + +@pytest.mark.parametrize( + ("sampling_params", "expected"), + [ + pytest.param(SamplingParams(), False, id="defaults"), + pytest.param(SamplingParams(temperature=0.0), False, id="greedy"), + pytest.param( + SamplingParams(thinking_token_budget=3), True, id="thinking-budget" + ), + pytest.param(SamplingParams(logit_bias={1: 1.0}), True, id="logit-bias"), + pytest.param(SamplingParams(frequency_penalty=0.1), True, id="penalty"), + pytest.param(SamplingParams(_bad_words_token_ids=[[1]]), True, id="bad-words"), + pytest.param(SamplingParams(temperature=0.7), True, id="temperature"), + pytest.param(SamplingParams(min_p=0.1), True, id="min-p"), + pytest.param(SamplingParams(top_k=10), True, id="top-k"), + pytest.param(SamplingParams(top_p=0.9), True, id="top-p"), + pytest.param( + SamplingParams.for_sampler_warmup(), True, id="all-logits-processors" + ), + ], +) +def test_logits_processing_cache_matches_request_features( + sampling_params: SamplingParams, expected: bool +): + sampler = _make_sampler() + sampler.add_request(3, prompt_len=1, sampling_params=sampling_params) + + assert sampler.needs_logits_processing[3] == expected + + +def test_logits_processing_cache_is_overwritten_when_slot_is_reused(): + sampler = _make_sampler() + sampler.add_request(3, 1, SamplingParams.for_sampler_warmup()) + sampler.add_request(3, 1, SamplingParams()) + + assert not sampler.needs_logits_processing[3] + + +def test_logits_processing_cache_only_checks_active_requests(): + sampler = _make_sampler() + sampler.add_request(0, 1, SamplingParams(temperature=0.0)) + sampler.add_request(2, 1, SamplingParams.for_sampler_warmup()) + + sampling_only = np.array([0], dtype=np.int32) + with_processing = np.array([0, 2], dtype=np.int32) + + assert not np.any(sampler.needs_logits_processing[sampling_only]) + assert np.any(sampler.needs_logits_processing[with_processing]) diff --git a/tests/v1/worker/test_gpu_thinking_budget.py b/tests/v1/worker/test_gpu_thinking_budget.py index 0b0f19f7dbf0..2bc697b08438 100644 --- a/tests/v1/worker/test_gpu_thinking_budget.py +++ b/tests/v1/worker/test_gpu_thinking_budget.py @@ -12,6 +12,7 @@ ) from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState from vllm.v1.worker.gpu.states import RequestState @@ -185,6 +186,44 @@ def test_v2_thinking_budget_ignores_plain_request(): assert torch.all(out == 0) +def test_v2_greedy_sampling_applies_thinking_budget(): + """Greedy-only requests must not bypass thinking-budget processing.""" + req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1) + sampler = Sampler( + max_num_reqs=4, + vocab_size=VOCAB_SIZE, + device=DEVICE, + req_states=req_states, + reasoning_config=MockReasoningConfig(), + ) + sampler.add_request( + req_idx=3, + prompt_len=1, + sampling_params=SamplingParams( + temperature=0.0, + thinking_token_budget=3, + ), + ) + sampler.apply_staged_writes() + + idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE) + idx_mapping_np = idx_mapping.cpu().numpy() + expanded_idx_mapping = idx_mapping.clone() + input_ids = torch.tensor([12], dtype=torch.int32, device=DEVICE) + logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE) + out = sampler.apply_sampling_params( + logits, + expanded_idx_mapping, + idx_mapping, + idx_mapping_np, + torch.tensor([4], dtype=torch.int32, device=DEVICE), + input_ids, + torch.tensor([0], dtype=torch.int32, device=DEVICE), + ) + + assert out[0, END].item() == pytest.approx(1.0e9) + + def test_v2_thinking_budget_latest_prefill_end_disables_forcing(): req_states = _make_req_states( [1, START, 10, 11, 12, END, 13], diff --git a/tests/v1/worker/test_jit_warmup_migration.py b/tests/v1/worker/test_jit_warmup_migration.py new file mode 100644 index 000000000000..ef887c5351f6 --- /dev/null +++ b/tests/v1/worker/test_jit_warmup_migration.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate the registry reference kernel against runtime dispatch.""" + +import pytest + +from vllm.platforms import current_platform + +if not current_platform.is_cuda_alike(): + pytest.skip("NVIDIA dispatch tests require CUDA", allow_module_level=True) + +from vllm.v1.worker.block_table import ComputeSlotMappingKernel + + +@pytest.mark.parametrize( + ("kv_cache_block_size", "blocks_per_kv_block", "block_size", "block_size_rep"), + [ + (256, 1, 256, 16), + (256, 4, 64, 16), + (64, 1, 64, 16), + (8, 1, 8, 2), + (4, 1, 4, 2), + ], +) +def test_compute_slot_mapping_warmup_matches_runtime_specializations( + kv_cache_block_size: int, + blocks_per_kv_block: int, + block_size: int, + block_size_rep: int, +) -> None: + kernel = ComputeSlotMappingKernel() + kwargs = dict( + kv_cache_block_size=kv_cache_block_size, + blocks_per_kv_block=blocks_per_kv_block, + total_cp_world_size=1, + total_cp_rank=0, + cp_kv_cache_interleave_size=1, + block_table_stride=32768, + block_size=block_size, + ) + expected = kernel.CompileKey( + kv_cache_block_size=kv_cache_block_size, + blocks_per_kv_block=blocks_per_kv_block, + total_cp_world_size=1, + total_cp_rank=0, + cp_kv_cache_interleave_size=1, + block_table_stride=16, + block_size=block_size_rep, + ) + + assert kernel.dispatch(**kwargs) == expected + assert kernel.get_warmup_keys(**kwargs) == [expected] diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 545d23f90912..749821274318 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + import pytest import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, +) from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @@ -18,6 +26,7 @@ def test_postprocess_state_scalar_with_int32_mapping( (4,), 9, dtype=torch.int32, device="cuda" ) state._align_mode = False + state.recoverssm = None state._mamba_ctx = None idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda") @@ -27,3 +36,61 @@ def test_postprocess_state_scalar_with_int32_mapping( [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda" ) torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) + + +def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: + state = RecoverSSMState() + metadata = Mock(spec=RecoverSSMMetadata) + metadata.commit_recoverssm_state.return_value = None + num_sampled = torch.tensor([3, 1], dtype=torch.int32) + idx_mapping = torch.tensor([0, 1], dtype=torch.int32) + num_accepted_tokens = torch.ones(2, dtype=torch.int32) + group = SimpleNamespace(layer_names=["layer"]) + + state.record_step({"layer": metadata}, [[group]], for_capture=False) + state.commit_step( + num_sampled, + idx_mapping, + state_indices=None, + num_accepted_tokens=num_accepted_tokens, + ) + state.commit_step( + num_sampled, + idx_mapping, + state_indices=None, + num_accepted_tokens=num_accepted_tokens, + ) + + metadata.commit_recoverssm_state.assert_called_once_with(num_sampled) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None: + state = object.__new__(MambaHybridModelState) + state._align_mode = True + state._mamba_ctx = None + state._mamba_state_idx_gpu = torch.full((5,), -1, dtype=torch.int32, device="cuda") + state.recoverssm = RecoverSSMState() + state.num_accepted_tokens_gpu = torch.full( + (5,), 9, dtype=torch.int32, device="cuda" + ) + metadata = Mock(spec=RecoverSSMMetadata) + metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata( + num_spec_decodes=1, + request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"), + block_size=8, + block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"), + ) + num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda") + group = SimpleNamespace(layer_names=["layer"]) + + state.recoverssm.record_step({"layer": metadata}, [[group]], for_capture=False) + + state.postprocess_state(idx_mapping, num_sampled) + + expected_state_indices = [-1, 1, -1, -1, -1] + assert state._mamba_state_idx_gpu.tolist() == expected_state_indices + expected_accepted = [9, 1, 9, 2, 9] + assert state.num_accepted_tokens_gpu.tolist() == expected_accepted diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a338b934b54f..0534795f9e84 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -18,6 +18,7 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, + batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, preprocess_mamba, @@ -489,6 +490,42 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): ) +def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: + device = torch.device("cpu") + config = _TestConfig(num_layers=1) + layer_names = ["layer_0"] + kv_cache_config = _make_kv_cache_config(config, layer_names) + conv_state = torch.empty( + config.num_blocks, + config.conv_width, + config.conv_inner_dim, + dtype=config.dtype, + ) + temporal_state = torch.empty( + config.num_blocks, config.temporal_state_dim, dtype=config.dtype + ) + attention = MagicMock() + attention.kv_cache = [ + conv_state, + temporal_state, + *(torch.empty(config.num_blocks, 1) for _ in range(4)), + ] + context = _make_gpu_ctx(config, kv_cache_config, device) + + context.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, + [torch.zeros(1, 1, dtype=torch.int32)], + ) + + assert context.is_initialized + assert context.state_base_addrs.tolist() == [ + conv_state.data_ptr(), + temporal_state.data_ptr(), + ] + + def _run_gpu_postprocess( gpu_ctx: MambaSpecDecodeGPUContext, *, @@ -536,6 +573,34 @@ def device(self): def test_config(self): return _TestConfig() + def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device): + batch = 128 + row_bytes = 32 * 1024 + shift = 16 + copy_size = row_bytes - shift + + pattern = (torch.arange(row_bytes, dtype=torch.int32, device=device) % 251).to( + torch.uint8 + ) + state = pattern.expand(batch, -1).clone() + snapshot = state.clone() + + row_stride_bytes = state.stride(0) * state.element_size() + row_offsets = ( + torch.arange(batch, dtype=torch.int64, device=device) * row_stride_bytes + ) + dst_ptrs = (row_offsets + state.data_ptr()).to(torch.uint64) + src_ptrs = (row_offsets + state.data_ptr() + shift).to(torch.uint64) + sizes = torch.full((batch,), copy_size, dtype=torch.int32, device=device) + + expected = snapshot.clone() + expected[:, :copy_size].copy_(snapshot[:, shift:]) + for _ in range(10): + state.copy_(snapshot) + batch_memcpy(src_ptrs, dst_ptrs, sizes) + torch.accelerator.synchronize() + torch.testing.assert_close(state, expected, rtol=0, atol=0) + def test_matches_python_postprocess_mamba(self, device, test_config): """ Golden test: GPU kernel produces identical results to Python impl. @@ -1190,12 +1255,27 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( # --- Verify Python behavior (ground truth) --- dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1 - # Conv state should be modified (shifted copy within block) - conv_changed = not torch.allclose( - conv_state_py[dest_block_id], conv_state_orig[dest_block_id] + # This is an overlapping in-place left shift, so comparing only the + # Python and fused paths can hide the same memcpy race in both. Build + # the memmove result from the untouched snapshot and check each path + # independently. + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-1].copy_( + conv_state_orig[dest_block_id, 1:] ) - assert conv_changed, ( - "Python: Conv state should be modified when accept_token_bias > 0" + torch.testing.assert_close( + conv_state_py, + expected_conv_state, + rtol=0, + atol=0, + msg="Python: overlapping conv copy should have memmove semantics", + ) + torch.testing.assert_close( + conv_state_gpu, + expected_conv_state, + rtol=0, + atol=0, + msg="GPU: overlapping conv copy should have memmove semantics", ) # Temporal state should be modified (copy from different block) @@ -2122,13 +2202,10 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): actual_src_block_idx = src_block_idx + accept_token_bias actual_src_block_id = block_table[req, actual_src_block_idx] - All prior regression tests exercise only ``bias == 1``, i.e. they - only ever read one slot ahead of ``src_block_idx`` in the block - table. An off-by-one (or missing scale) in the address computation - on line 143 of ``mamba_utils.py`` would be invisible to every - existing test but would silently read the wrong physical block on - any speculative-decode cycle that accepts multiple tokens across a - block boundary, feeding a stale hidden state forward one step. + A ``bias == 1`` case only reads one slot ahead of ``src_block_idx`` + in the block table. This test isolates the larger-stride case, where + an off-by-one would read the wrong physical block after multiple + tokens are accepted across a block boundary. Setup (block_size=16): - running = 28 + 2 - 0 = 30 @@ -2141,8 +2218,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): With identity block_ids = [0,1,2,3,...], an off-by-one that used bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3, - producing a clear state-value mismatch against the Python - reference. + producing a clear mismatch against the untouched snapshot. """ cfg = test_config torch.manual_seed(7002) @@ -2166,6 +2242,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): fwd_py, fwd_gpu, ) = _make_dual_layer_state(cfg, device) + conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Python reference --- @@ -2212,12 +2289,22 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): device=device, ) - # --- Ground truth: Python must have sourced temporal from block 3 --- + # --- Ground truth from untouched snapshots --- actual_src_block_id = block_ids_per_req[0][3] # == 3 dest_block_id = block_ids_per_req[0][1] # == 1 + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-2].copy_( + conv_state_orig[dest_block_id, 2:] + ) + torch.testing.assert_close(conv_state_py, expected_conv_state, rtol=0, atol=0) + torch.testing.assert_close(conv_state_gpu, expected_conv_state, rtol=0, atol=0) + + # Python must have sourced temporal from block 3. torch.testing.assert_close( temporal_state_py[dest_block_id], temporal_state_orig[actual_src_block_id], + rtol=0, + atol=0, msg=( "Python reference did not copy from block_ids[src+bias]=3; " "test preconditions are wrong" @@ -2251,21 +2338,44 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): msg="num_accepted_tokens mismatch at accept_token_bias=2", ) - def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( - self, device, test_config, monkeypatch + @pytest.mark.parametrize( + "same_physical_block", [True, False], ids=["same", "distinct"] + ) + @pytest.mark.parametrize("accept_token_bias", [1, 2, 3]) + @pytest.mark.parametrize( + "dtype", + [torch.float16, torch.float32, torch.float64], + ids=["fp16", "fp32", "fp64"], + ) + def test_sd_and_ds_conv_layouts_match_snapshot( + self, + device, + test_config, + monkeypatch, + accept_token_bias, + same_physical_block, + dtype, ): - """DS conv postprocess should match SD when accept_token_bias > 0.""" + """SD and DS copies should independently match memmove semantics.""" from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils cfg = test_config + cfg.dtype = dtype torch.manual_seed(38898) req_ids = ["req_0"] - num_computed_tokens = [30] - num_scheduled_tokens = {"req_0": 1} + # Keep new_num_computed on an aligned boundary while varying how far + # below it the running state starts. This makes the copy bias exactly + # ``accept_token_bias`` for each case. The 32 boundary keeps source and + # destination in logical block 1; the 64 boundary copies block 2 -> 3. + aligned_boundary = 32 if same_physical_block else 64 + num_computed_tokens = [aligned_boundary - 2 * accept_token_bias] + num_scheduled_tokens = {"req_0": accept_token_bias} num_draft_tokens: dict[str, int] = {} - num_accepted_tokens = [2] # Results in accept_token_bias = 1 - mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx + num_accepted_tokens = [accept_token_bias + 1] + dest_block_idx = aligned_boundary // cfg.block_size - 1 + src_block_idx = dest_block_idx if same_physical_block else dest_block_idx - 1 + mamba_state_idx = [src_block_idx] block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] @@ -2288,7 +2398,8 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device ) - # SD GPU path. Default layout is SD. + # SD GPU path. + monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False) model_mamba_utils.get_conv_state_layout.cache_clear() sd_conv = sd_source_conv.clone() sd_temporal = sd_source_temporal.clone() @@ -2312,9 +2423,32 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( ) torch.accelerator.synchronize() - # Sanity: SD path actually modified the state (copy was performed). - assert not torch.equal(sd_conv, sd_source_conv), ( - "SD baseline did not modify conv state; test setup is wrong" + src_block_id = block_ids_per_req[0][src_block_idx] + dest_block_id = block_ids_per_req[0][dest_block_idx] + expected_conv = sd_source_conv.clone() + expected_conv[dest_block_id, :-accept_token_bias].copy_( + sd_source_conv[src_block_id, accept_token_bias:] + ) + torch.testing.assert_close( + sd_conv, + expected_conv, + rtol=0, + atol=0, + msg="SD conv copy did not match the untouched source snapshot", + ) + + actual_temporal_src_idx = src_block_idx + accept_token_bias + actual_temporal_src_id = block_ids_per_req[0][actual_temporal_src_idx] + expected_temporal = sd_source_temporal.clone() + expected_temporal[dest_block_id].copy_( + sd_source_temporal[actual_temporal_src_id] + ) + torch.testing.assert_close( + sd_temporal, + expected_temporal, + rtol=0, + atol=0, + msg="SD temporal copy did not match the untouched source snapshot", ) # DS GPU path on the DS twin. @@ -2346,22 +2480,39 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( # Reset the lru cache so other tests see the default layout again. model_mamba_utils.get_conv_state_layout.cache_clear() - # DS bytes, un-permuted, should match the SD result. + # Validate DS independently against the snapshot; otherwise a shared + # SD/DS bug would remain invisible. + ds_conv_sd_layout = ds_conv.permute(0, 2, 1).contiguous() torch.testing.assert_close( - ds_conv.permute(0, 2, 1).contiguous(), - sd_conv, - msg=( - "DS conv post-kernel does not match SD baseline; the DS " - "row-loop in postprocess_mamba_fused_kernel is wrong." - ), + ds_conv_sd_layout, + expected_conv, + rtol=0, + atol=0, + msg="DS conv copy did not match the untouched source snapshot", ) torch.testing.assert_close( ds_temporal, - sd_temporal, - msg="DS temporal state diverged from SD", + expected_temporal, + rtol=0, + atol=0, + msg="DS temporal copy did not match the untouched source snapshot", + ) + + expected_accepted = 1 if same_physical_block else accept_token_bias + 1 + expected_accepted_tensor = torch.tensor( + [expected_accepted], dtype=torch.int32, device=device ) torch.testing.assert_close( - gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], gpu_ctx_sd.num_accepted_tokens_out[:num_reqs], - msg="DS num_accepted_tokens diverged from SD", + expected_accepted_tensor, + rtol=0, + atol=0, + msg="SD num_accepted_tokens result is wrong", + ) + torch.testing.assert_close( + gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], + expected_accepted_tensor, + rtol=0, + atol=0, + msg="DS num_accepted_tokens result is wrong", ) diff --git a/tests/v1/worker/test_workspace.py b/tests/v1/worker/test_workspace.py new file mode 100644 index 000000000000..a909ea9af89d --- /dev/null +++ b/tests/v1/worker/test_workspace.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import cast + +import pytest +import torch + +import vllm.v1.worker.workspace as workspace +from vllm.config import VllmConfig +from vllm.v1.worker.gpu_worker import _num_workspace_lanes + + +class _SpecConfig: + def __init__(self, dspark: bool) -> None: + self._dspark = dspark + + def use_dspark(self) -> bool: + return self._dspark + + +class _VllmConfig: + def __init__(self, spec_config: _SpecConfig | None) -> None: + self.speculative_config = spec_config + + +@pytest.mark.parametrize( + ("use_v2_model_runner", "spec_config", "expected"), + [ + (True, _SpecConfig(True), 2), + (False, _SpecConfig(True), 1), + (True, _SpecConfig(False), 1), + (True, None, 1), + ], +) +def test_workspace_lane_count_is_dspark_only( + use_v2_model_runner: bool, + spec_config: _SpecConfig | None, + expected: int, +) -> None: + config = cast(VllmConfig, _VllmConfig(spec_config)) + assert _num_workspace_lanes(config, use_v2_model_runner) == expected + + +def test_workspace_lanes_do_not_alias_and_restore_context(monkeypatch) -> None: + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + + assert manager._current_workspaces == [None, None, None, None] + + (target,) = manager.get_simultaneous(((512,), torch.uint8)) + with workspace.use_workspace_lane(1): + (draft,) = manager.get_simultaneous(((256,), torch.uint8)) + (draft_reused,) = manager.get_simultaneous(((8,), torch.uint8)) + (target_reused,) = manager.get_simultaneous(((8,), torch.uint8)) + + assert manager._current_workspaces[0].numel() == 512 # type: ignore[union-attr] + assert manager._current_workspaces[1].numel() == 256 # type: ignore[union-attr] + assert manager._current_workspaces[2:] == [None, None] + assert target.data_ptr() != draft.data_ptr() + assert draft.data_ptr() == draft_reused.data_ptr() + assert target.data_ptr() == target_reused.data_ptr() + + +def test_workspace_lanes_compose_with_ubatches(monkeypatch) -> None: + active_ubatch = [0] + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch[0]) + manager = workspace.WorkspaceManager( + torch.device("cpu"), num_ubatches=2, num_lanes=2 + ) + + pointers = set() + for ubatch_id in range(2): + active_ubatch[0] = ubatch_id + for lane in range(2): + with workspace.use_workspace_lane(lane): + (buffer,) = manager.get_simultaneous(((16,), torch.uint8)) + pointers.add(buffer.data_ptr()) + + assert len(pointers) == 4 + + +def test_workspace_lane_validation(monkeypatch) -> None: + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0) + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=1) + + with ( + pytest.raises(ValueError, match="non-negative"), + workspace.use_workspace_lane(-1), + ): + pass + + with ( + workspace.use_workspace_lane(1), + pytest.raises(RuntimeError, match="is not configured"), + ): + manager.get_simultaneous(((1,), torch.uint8)) + + with pytest.raises(ValueError, match="at least one"): + workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 52a95ce1d8d9..59163a23e09b 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -92,6 +92,15 @@ class ForbiddenImport: ), allowed_files={"vllm/triton_utils/importing.py"}, ), + "tilelang": ForbiddenImport( + pattern=r"^(from|import)\s+tilelang(\s|\.|$)", + tip="Use 'from vllm.tilelang_utils import tilelang, T' instead.", + allowed_pattern=re.compile( + r"from\s+vllm\.tilelang_utils\s+import\s+" + r"(tilelang|T|T, tilelang|tilelang, T)\b" + ), + allowed_files={"vllm/tilelang_utils/__init__.py"}, + ), "huggingface_hub repo API": ForbiddenImport( # Catch `from huggingface_hub import `, including parenthesized, # multi-line imports. @@ -182,6 +191,29 @@ def matches(rule: str, content: str) -> bool: f"(expected {should_match}, got {result})" ) + tilelang_cases = [ + # Should match + ("import tilelang", True), + ("import tilelang.language as T", True), + ("from tilelang.jit import JITImpl", True), + ("from tilelang.jit.kernel import JITKernel", True), + # Should not match: indented (local) imports are allowed, mirroring + # the "triton" rule, so mocked-module test imports are not flagged. + (" import tilelang", False), + (" from tilelang.jit import JITImpl", False), + ("from vllm.tilelang_utils import tilelang", False), + ("from vllm.tilelang_utils import T", False), + ("from vllm.tilelang_utils import T, tilelang", False), + ("from vllm.tilelang_utils import tilelang, T", False), + ("import tilelang_kernels", False), + ] + for i, (content, should_match) in enumerate(tilelang_cases): + result = matches("tilelang", content) + assert result == should_match, ( + f"tilelang case {i} failed: {content!r} " + f"(expected {should_match}, got {result})" + ) + hf_cases = [ # Should match ("from huggingface_hub import snapshot_download", True), diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index c4ff05c75dfa..5d63ffa0de37 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2777,6 +2777,43 @@ def fused_kda_decode( return out +def fused_gdn_decode_post_conv_mtp( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + num_accepted_tokens: torch.Tensor, + state: torch.Tensor, + output_gate: torch.Tensor, + norm_weight: torch.Tensor, + out: torch.Tensor | None = None, + scale: float = 128**-0.5, + norm_eps: float = 1e-5, +) -> torch.Tensor: + if out is None: + out = torch.empty_like(output_gate) + torch.ops._C.fused_gdn_decode_post_conv_mtp( + mixed_qkv, + a, + b, + A_log, + dt_bias, + state_indices, + cu_seqlens, + num_accepted_tokens, + state, + output_gate, + norm_weight, + out, + scale, + norm_eps, + ) + return out + + def concat_and_cache_mla( kv_c: torch.Tensor, k_pe: torch.Tensor, diff --git a/vllm/benchmarks/datasets/datasets.py b/vllm/benchmarks/datasets/datasets.py index 2db092580a05..d928123dde87 100644 --- a/vllm/benchmarks/datasets/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -79,7 +79,7 @@ class SampleRequest: Represents a single inference request for benchmarking. """ - prompt: str | list[str] | list[dict] + prompt: str | list[str] | list[int] | list[dict] prompt_len: int expected_output_len: int = 0 multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None @@ -1572,7 +1572,6 @@ def sample( prompt_ids = self._expand_prompt( entry.get(self.label_hash_ids, []), input_length, tokenizer ) - prompt = tokenizer.decode(prompt_ids) # Get timestamp with proper error handling ts_value = entry.get(self.label_ts) @@ -1588,7 +1587,7 @@ def sample( samples.append( SampleRequest( - prompt=prompt, + prompt=prompt_ids, prompt_len=prompt_len, expected_output_len=new_output_len, lora_request=None, diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index 59cbc0e2e6c4..1023700da4b6 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -66,7 +66,7 @@ def add_chunk(self, chunk_bytes: bytes) -> list[str]: class RequestFuncInput: """The input for the request function.""" - prompt: str | list[str] | list[dict[str, Any]] + prompt: str | list[str] | list[int] | list[dict[str, Any]] api_url: str prompt_len: int output_len: int diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index e4ab5a583a04..9b993689235f 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -2097,6 +2097,11 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]: args.ignore_eos = True if args.dataset_name == "timed_trace": + if args.backend not in ("vllm", "openai"): + raise ValueError( + "timed_trace dataset passes pre-tokenized prompts (list[int])" + " and requires a completions backend ('vllm' or 'openai')." + ) # timed_trace carries per-request timestamps; # ignore EOS so generation runs to the trace's specified output length, # and default to using those timestamps for scheduling unless the user diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index feb15953234e..fdfc5b8be463 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -93,8 +93,8 @@ def _norm_input_weight_dtype_match(match: pm.Match) -> bool: _flashinfer_comm, "create_allreduce_fusion_workspace" ): flashinfer_comm = _flashinfer_comm - except ImportError: - pass + except Exception as e: + logger.debug_once("flashinfer.comm import failed: %s", e) if hasattr(torch.ops._C, "scaled_fp4_quant"): STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.out diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 994be05f54ea..907917d091cb 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -7,10 +7,13 @@ from pydantic import field_validator from vllm.config.utils import config +from vllm.logger import init_logger from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum -IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"] +logger = init_logger(__name__) + +IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"] MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"] @@ -65,12 +68,15 @@ class AttentionConfig: use_prefill_query_quantization: bool = False """If set, quantize query for attention in prefill.""" - use_fp4_indexer_cache: bool = False - """If set, use fp4 indexer cache for dsv32 family model (not support yet)""" + use_fp4_indexer_cache: bool | None = None + """Deprecated alias for `indexer_kv_dtype`; use that instead. True maps to + `mxfp4`, False is a no-op (it selected the model default already).""" - indexer_kv_dtype: IndexerKVDType = "bf16" - """Data type for the sparse-attention indexer K cache. Quantized formats - (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + indexer_kv_dtype: IndexerKVDType = "auto" + """Data type for the sparse-attention indexer K cache. "auto" picks the + model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse + indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel + support in the backend.""" use_non_causal: bool = False """Whether to use non-causal (bidirectional) attention.""" @@ -114,6 +120,26 @@ def __post_init__(self) -> None: # layers still use the platform's normal automatic backend. self.backend = None + if self.use_fp4_indexer_cache is not None: + logger.warning( + "use_fp4_indexer_cache is deprecated and will be removed in " + "v0.19. Use indexer_kv_dtype instead (True -> 'mxfp4')." + ) + if self.use_fp4_indexer_cache: + if self.indexer_kv_dtype not in ("auto", "mxfp4"): + raise ValueError( + "use_fp4_indexer_cache=True conflicts with " + f"indexer_kv_dtype={self.indexer_kv_dtype!r}. Set only " + "indexer_kv_dtype." + ) + self.indexer_kv_dtype = "mxfp4" + + def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType: + """Resolve `indexer_kv_dtype`, substituting `default` for "auto".""" + if self.indexer_kv_dtype == "auto": + return default + return self.indexer_kv_dtype + def compute_hash(self) -> str: """ Provide a hash that uniquely identifies all the configs @@ -124,7 +150,8 @@ def compute_hash(self) -> str: """ from vllm.config.utils import get_hash_factors, hash_factors - ignored_factors: set[str] = set() + # Folded into indexer_kv_dtype by __post_init__. + ignored_factors: set[str] = {"use_fp4_indexer_cache"} factors = get_hash_factors(self, ignored_factors) return hash_factors(factors) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 654e9a53589b..27e46534dcd3 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -7,7 +7,7 @@ from pydantic import Field, field_validator, model_validator -from vllm.config.utils import config +from vllm.config.utils import config, get_from_deprecated_env_if_set from vllm.logger import init_logger from vllm.utils.torch_utils import ( is_quantized_kv_cache, @@ -35,6 +35,17 @@ "nvfp4", "nvfp4_4over6", ] + + +def _get_prefix_cache_retention_interval() -> int | None: + env_value = get_from_deprecated_env_if_set( + "VLLM_PREFIX_CACHE_RETENTION_INTERVAL", + "v0.29", + "prefix_cache_retention_interval", + ) + return 0 if env_value is None else int(env_value) + + MambaDType = Literal["auto", "float32", "float16", "bfloat16"] MambaCacheMode = Literal["all", "align", "none"] PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"] @@ -111,6 +122,15 @@ class CacheConfig: security risk tolerance against the performance benefits before turning this on. - "xxhash_cbor" combines canonical CBOR serialization with xxHash for reproducible hashing. Requires the optional ``xxhash`` package.""" + prefix_cache_retention_interval: int | None = Field( + default_factory=_get_prefix_cache_retention_interval, ge=0 + ) + """Token interval between retained sliding-window and Mamba prefix-cache + checkpoints. ``0`` retains only semantic checkpoints, including the latest + replay boundary and shared-prefix junctions. Positive values additionally + retain periodic checkpoints at the specified interval, which must be a + multiple of the scheduler block size. ``None`` retains checkpoints densely. + Applies only to sliding-window and Mamba cache groups.""" kv_cache_dtype_skip_layers: list[str] = field(default_factory=list) """Layer patterns to skip KV cache quantization. Accepts layer indices (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window').""" @@ -144,9 +164,8 @@ class CacheConfig: caching is enabled. """ replayssm_buffer_len: int = Field(default=16, gt=0) - """ReplaySSM history buffer length B: with use_replayssm, standard decode - caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint - state to HBM every B steps. Default 16.""" + """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3 + speculative decoding does not use B. Default 16.""" use_replayssm: bool = False """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. @@ -154,6 +173,8 @@ class CacheConfig: mamba backend; standard (non-speculative) decode only. In align mode flushes are most efficient when mamba_block_size is a multiple of replayssm_buffer_len, but this is not required.""" + use_kda_recoverssm: bool = field(default=False, init=False) + """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) @@ -217,6 +238,7 @@ def compute_hash(self) -> str: "num_gpu_blocks_override", "enable_prefix_caching", "prefix_caching_hash_algo", + "prefix_cache_retention_interval", # Prefix-caching implementation detail (doesn't affect compiled graph). "prefix_match_unit", "mamba_page_size_padded", diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 0776a60a4697..3cd227d72ce4 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -769,6 +769,7 @@ class CompilationConfig: "vllm::short_conv", "vllm::linear_attention", "vllm::qwen_gdn_attention_core", + "vllm::qwen_gdn_attention_core_fused_norm_packed", "vllm::gdn_attention_core_xpu", "vllm::olmo_hybrid_gdn_full_forward", "vllm::sparse_attn_indexer", diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 74e959f33ab9..7287c1ce923f 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -147,6 +147,7 @@ def with_default( "flashinfer_trtllm", "flashinfer_cudnn", "flashinfer_b12x", + "b12x", "marlin", "humming", "triton", @@ -223,6 +224,7 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+) + - "b12x": Use native B12X FP8 and FP4 linear kernels on SM12x - "marlin": Use Marlin kernels - "triton": Use Triton-based kernels - "deep_gemm": Use DeepGEMM kernels diff --git a/vllm/config/model.py b/vllm/config/model.py index b0dd48b7e503..fcdce3ef5858 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import warnings from collections.abc import Callable from dataclasses import InitVar, field from functools import cached_property @@ -349,8 +348,6 @@ class ModelConfig: - "transformers" will use the Transformers model implementation. - "terratorch" will use the TerraTorch model implementation. """ - override_attention_dtype: str | None = None - """Override dtype for attention""" logits_processors: list[str | type[LogitsProcessor]] | None = None """One or more logits processors' fully-qualified class names or class definitions""" @@ -433,7 +430,6 @@ def compute_hash(self) -> str: "config_format", "hf_token", "hf_overrides", - "override_attention_dtype", "logits_processors", "io_processor_plugin", "pooler_config", @@ -595,12 +591,6 @@ def __post_init__( self.hf_token, ) - if self.override_attention_dtype is not None and not current_platform.is_rocm(): - warnings.warn( - "override-attention-dtype is set but not using ROCm platform", - stacklevel=2, - ) - if self.enable_sleep_mode: if not current_platform.is_sleep_mode_available(): raise ValueError("Sleep mode is not supported on current platform.") diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index d8f049755ce4..c3e0866b453d 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -630,8 +630,7 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: hf_config.model_type = "inkling_mtp" hf_config.update( { - # Inkling currently exposes only the first checkpoint depth. - "n_predict": 1, + "n_predict": checkpoint_depths, "num_nextn_predict_layers": checkpoint_depths, "chain_hidden_post_norm": mtp_config.get( "chain_hidden_post_norm", False @@ -944,12 +943,20 @@ def __post_init__(self): self.method = "eagle" elif "eagle3" in self.draft_model_config.model.lower(): self.method = "eagle3" - elif "dflash" in self.draft_model_config.model.lower(): + elif ( + "dflash" in self.draft_model_config.model.lower() + or "MuseGlimmerAssistantModel" + in self.draft_model_config.architectures + ): self.method = "dflash" elif ( "dspark" in self.draft_model_config.model.lower() or "Qwen3DSparkModel" in self.draft_model_config.architectures or "Gemma4DSparkModel" in self.draft_model_config.architectures + or ( + "DSparkDraftModel" in self.draft_model_config.architectures + and self.draft_model_config.hf_config.model_type == "qwen3" + ) ): self.method = "dspark" elif self.draft_model_config.hf_config.model_type == "medusa": @@ -1007,7 +1014,16 @@ def __post_init__(self): self.draft_model_config.hf_config = eagle_config self.update_arch_() - if self.method == "dspark" and ( + if ( + self.method == "dspark" + and "DSparkDraftModel" in self.draft_model_config.architectures + and self.draft_model_config.hf_config.model_type == "qwen3" + ): + self.draft_model_config.hf_config.architectures = [ + "Qwen3DSparkModel" + ] + self.update_arch_() + elif self.method == "dspark" and ( "Qwen3DSparkModel" not in self.draft_model_config.architectures and "Gemma4DSparkModel" not in self.draft_model_config.architectures and "K3DSparkModel" not in self.draft_model_config.architectures @@ -1043,16 +1059,6 @@ def __post_init__(self): if self.method in ("dflash", "dspark"): self.parallel_drafting = True - if ( - self.method == "dspark" - and "K3DSparkModel" in self.draft_model_config.architectures - and self.target_parallel_config.decode_context_parallel_size > 1 - ): - raise ValueError( - "MLA DSpark does not currently support decode context " - "parallelism; set decode_context_parallel_size=1." - ) - if self.num_speculative_tokens is not None and hasattr( self.draft_model_config.hf_config, "num_lookahead_tokens" ): @@ -1083,43 +1089,11 @@ def __post_init__(self): "`num_speculative_tokens` was not provided" ) - if ( - self.draft_model_config.hf_config.model_type == "inkling_mtp" - and self.num_speculative_tokens != 1 - ): - raise ValueError( - "Inkling MTP currently supports exactly one speculative token" - ) - if self.dspark_draft_topk is not None and self.method != "dspark": raise ValueError("dspark_draft_topk is only supported by DSpark") dspark_draft_topk = None if self.method == "dspark": - # DSpark is a semi-autoregressive *block* drafter. A - # speculative length smaller than the checkpoint's block - # feeds the block / Markov-head machinery an unsupported - # layout and yields incorrect (garbled) output rather than - # merely lower acceptance. Require num_speculative_tokens to - # be at least the block size (e.g. 5 or 7 for DeepSeek-V4). - dspark_block_size = getattr( - self.draft_model_config.hf_config, - "dspark_block_size", - None, - ) - if ( - dspark_block_size is not None - and self.num_speculative_tokens < dspark_block_size - ): - raise ValueError( - "DSpark requires num_speculative_tokens >= " - f"dspark_block_size ({dspark_block_size}); got " - f"{self.num_speculative_tokens}. Smaller values " - "produce incorrect output. Use " - f"num_speculative_tokens={dspark_block_size} or " - "larger (e.g. 7)." - ) - hf_config = self.draft_model_config.hf_config dspark_draft_topk = self.dspark_draft_topk if dspark_draft_topk is None: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index eeb2d840ffd8..a2d6b293bb12 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,10 +66,6 @@ logger = init_logger(__name__) -MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES = frozenset( - {"DeepseekV4ForCausalLM"} -) - DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "DeepseekV2ForCausalLM", @@ -83,15 +79,6 @@ } ) -# Architectures that default to V1 on ROCm: the V2 runner faults during the -# profile run. VLLM_USE_V2_MODEL_RUNNER=1 still forces V2. -# TODO: fix V2 enablement -ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( - { - "KimiK3ForConditionalGeneration", - } -) - @lru_cache def default_v2_model_runner_architectures() -> frozenset[str]: @@ -99,10 +86,10 @@ def default_v2_model_runner_architectures() -> frozenset[str]: from vllm.platforms import current_platform if current_platform.is_rocm(): - return ( - DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - - ROCM_EXCLUDED_V2_MODEL_RUNNER_ARCHITECTURES - ) + # TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The + # attention layer picks the eager cudagraph region MRV1 needs, so + # this is a perf default only; drop it once MRV2 catches up. + return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"} return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES @@ -707,25 +694,6 @@ def _is_default_v2_model_runner_model(self) -> bool: return False return is_default_v2_architecture or not model_config.is_moe - def _validate_mrv1_piecewise_cudagraph(self) -> None: - if self.use_v2_model_runner: - return - model_config = self.model_config - if model_config is None: - return - if not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs(): - return - architectures = getattr(model_config, "architectures", []) - if any( - arch in MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES - for arch in architectures - ): - raise ValueError( - "DeepSeek V4 does not support PIECEWISE CUDA graphs with " - "Model Runner V1. Use Model Runner V2 or disable PIECEWISE " - "CUDA graphs." - ) - @property def needs_dp_coordinator(self) -> bool: """ @@ -1090,6 +1058,14 @@ def __post_init__(self): "--enable-return-routed-experts is incompatible with " "pipeline parallelism (PP > 1)." ) + if ( + self.parallel_config.decode_context_parallel_size > 1 + or self.parallel_config.prefill_context_parallel_size > 1 + ): + raise ValueError( + "--enable-return-routed-experts is incompatible with context " + "parallelism (DCP > 1 or PCP > 1)." + ) # Incompatible with any KV connector — covers both PD disaggregation # (kv_producer/kv_consumer: routing captured on P can't reach D) and @@ -1652,8 +1628,6 @@ def has_blocked_weights(): "pipeline parallelism", ) - self._validate_mrv1_piecewise_cudagraph() - # final check of cudagraph mode after all possible updates if current_platform.is_cuda_alike(): if ( @@ -2566,22 +2540,47 @@ def validate_mamba_block_size(self) -> "VllmConfig": @model_validator(mode="after") def validate_mamba_cached_kernel(self) -> "VllmConfig": if not self.cache_config.use_replayssm: + self.cache_config.use_kda_recoverssm = False return self - # ReplaySSM adds a 3-tensor ring to the mamba state; only models that - # opt in (supports_replayssm) build a consistent shape on both the layer - # and config paths. Reject others so the mamba page size cannot desync. + self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0 + if self.model_config is not None and not self.model_config.supports_replayssm: raise ValueError( - "--use-replayssm is only supported for Nemotron-H models " - f"(got architecture {self.model_config.architecture!r})" + "--use-replayssm is not supported for architecture " + f"{self.model_config.architecture!r}" ) - if self.cache_config.mamba_cache_mode == "all": + if self.cache_config.use_kda_recoverssm: + if self.model_config is not None and self.model_config.architecture not in ( + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", + ): + raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA") + if self.mamba_config.enable_stochastic_rounding: + raise ValueError( + "RecoverSSM supports bfloat16/float32 " + "SSM state caches, not --enable-mamba-cache-stochastic-" + "rounding, which requires an explicit float16 cache" + ) + if self.cache_config.mamba_cache_mode not in ("none", "align"): + raise ValueError( + "RecoverSSM supports only none and align Mamba cache modes" + ) + if ( + self.cache_config.mamba_cache_mode == "align" + and not self.use_v2_model_runner + ): + raise ValueError( + "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + ) + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError( + "RecoverSSM currently requires pipeline_parallel_size=1" + ) + elif self.cache_config.mamba_cache_mode == "all": raise ValueError( "--use-replayssm supports prefix caching only in align mode; " "pass --mamba-cache-mode align" ) - if self.num_speculative_tokens > 0: - raise ValueError("--use-replayssm does not support speculative decoding") if self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError("--use-replayssm requires --mamba-backend triton") if ( diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index d86a550eeabc..23721aac440c 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -40,6 +40,18 @@ def recast_val(x, dtype, *, loc=None, ip=None): return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip))) +@dsl_user_op +def to_cta0_smem(ptr: cute.Pointer, *, loc=None, ip=None): + return cute.make_ptr( + ptr.dtype, + ptr.toint(loc=loc, ip=ip) & 0xFEFF_FFFF, + cute.AddressSpace.smem, + assumed_align=8, + loc=loc, + ip=ip, + ) + + def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None): """A simple helper that wraps group_modes() and tma_partition() NOTE: this should be called WITHOUT cute.elect_one() diff --git a/vllm/cute_utils/mbarrier.py b/vllm/cute_utils/mbarrier.py new file mode 100644 index 000000000000..d9135258102a --- /dev/null +++ b/vllm/cute_utils/mbarrier.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from cutlass import Int32 +from cutlass._mlir.dialects import nvvm +from cutlass.cutlass_dsl import dsl_user_op + +_SPACE_MAP = { + "cta": nvvm.MBarrierSpaceKind.CTA, + "cluster": nvvm.MBarrierSpaceKind.CLUSTER, +} + +_MEMORY_ORDER_MAP = { + "weak": nvvm.MemOrderKind.WEAK, + "relaxed": nvvm.MemOrderKind.RELAXED, + "acquire": nvvm.MemOrderKind.ACQUIRE, + "release": nvvm.MemOrderKind.RELEASE, + "acq_rel": nvvm.MemOrderKind.ACQ_REL, +} + + +@dsl_user_op +def arrive(mbar, space: str = "cta", order: str = "relaxed", *, loc=None, ip=None): + nvvm.mbarrier_txn( + mbar.to_llvm_ptr(loc=loc, ip=ip), + Int32(1).ir_value(loc=loc, ip=ip), + kind=nvvm.MBarrierTxnKind.ARRIVE, + space=_SPACE_MAP[space], + order=_MEMORY_ORDER_MAP[order], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def arrive_expect_tx( + mbar, size, space: str = "cta", order: str = "relaxed", *, loc=None, ip=None +): + nvvm.mbarrier_txn( + mbar.to_llvm_ptr(loc=loc, ip=ip), + Int32(size).ir_value(loc=loc, ip=ip), + kind=nvvm.MBarrierTxnKind.ARRIVE_EXPECT_TX, + space=_SPACE_MAP[space], + order=_MEMORY_ORDER_MAP[order], + loc=loc, + ip=ip, + ) diff --git a/vllm/device_allocator/__init__.py b/vllm/device_allocator/__init__.py index 66e8b146d29b..02d51603163e 100644 --- a/vllm/device_allocator/__init__.py +++ b/vllm/device_allocator/__init__.py @@ -27,6 +27,8 @@ def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ... def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ... + def discard(self, tags: tuple[str, ...] | str) -> None: ... + def wake_up(self, tags: list[str] | None = None) -> None: ... def get_current_usage(self) -> int: ... diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 2cb9805bae39..a16f1b01486e 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -247,8 +247,15 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: total_bytes = 0 backup_bytes = 0 + has_policy_conflict = False for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + requests_offload = data.tag in offload_tags + was_offloaded = data.cpu_backup_tensor is not None + if requests_offload != was_offloaded: + has_policy_conflict = True + continue handle = data.handle total_bytes += handle[1] if data.tag in offload_tags: @@ -277,9 +284,46 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: (total_bytes - backup_bytes) / 1024**3, ) + if has_policy_conflict: + logger.warning( + "CuMemAllocator: sleep cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + gc.collect() torch.cuda.empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Discard mapped allocations with the given tags without CPU backup.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + has_policy_conflict = False + for data in self.pointer_to_data.values(): + if data.tag not in tags: + continue + if data.is_asleep: + if data.cpu_backup_tensor is not None: + has_policy_conflict = True + continue + torch.accelerator.synchronize(data.handle[0]) + unmap_and_release(data.handle) + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "CuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_bytes / 1024**3, + tags, + ) + + if has_policy_conflict: + logger.warning( + "CuMemAllocator: discard cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + def wake_up(self, tags: list[str] | None = None) -> None: """ Wake up the allocator from sleep mode. @@ -295,6 +339,8 @@ def wake_up(self, tags: list[str] | None = None) -> None: torch.accelerator.empty_cache() for ptr, data in self.pointer_to_data.items(): + if not data.is_asleep: + continue if tags is None or data.tag in tags: handle = data.handle create_and_map(handle) diff --git a/vllm/device_allocator/xpumem.py b/vllm/device_allocator/xpumem.py index e0f359b200d2..146ea965db07 100644 --- a/vllm/device_allocator/xpumem.py +++ b/vllm/device_allocator/xpumem.py @@ -174,12 +174,20 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: total_bytes = 0 backup_bytes = 0 + has_policy_conflict = False for ptr, data in self.pointer_to_data.items(): + if data.is_asleep: + requests_offload = data.tag in offload_tags + was_offloaded = data.cpu_backup_tensor is not None + if requests_offload != was_offloaded: + has_policy_conflict = True + continue size_in_bytes = data.handle[1] total_bytes += size_in_bytes if data.tag not in offload_tags: unmap_and_release(data.handle) + data.is_asleep = True continue backup_bytes += size_in_bytes @@ -201,6 +209,7 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: data.cpu_backup_tensor = cpu_backup_tensor unmap_and_release(data.handle) + data.is_asleep = True logger.info( "XpuMemAllocator: sleep freed %.2f GiB memory in total, of which " @@ -211,16 +220,56 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: (total_bytes - backup_bytes) / 1024**3, ) + if has_policy_conflict: + logger.warning( + "XpuMemAllocator: sleep cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + gc.collect() xpu_empty_cache = getattr(torch.xpu, "empty_cache", None) if callable(xpu_empty_cache): xpu_empty_cache() + def discard(self, tags: tuple[str, ...] | str) -> None: + """Discard mapped allocations with the given tags without CPU backup.""" + if isinstance(tags, str): + tags = (tags,) + + discarded_bytes = 0 + has_policy_conflict = False + for data in self.pointer_to_data.values(): + if data.tag not in tags: + continue + if data.is_asleep: + if data.cpu_backup_tensor is not None: + has_policy_conflict = True + continue + torch.accelerator.synchronize(data.handle[0]) + unmap_and_release(data.handle) + data.is_asleep = True + discarded_bytes += data.handle[1] + + logger.info( + "XpuMemAllocator: discarded %.2f GiB for tags %s.", + discarded_bytes / 1024**3, + tags, + ) + + if has_policy_conflict: + logger.warning( + "XpuMemAllocator: discard cannot change the policy of " + "already-asleep allocations; the existing policy was kept." + ) + def wake_up(self, tags: list[str] | None = None) -> None: for ptr, data in self.pointer_to_data.items(): + if not data.is_asleep: + continue if tags is not None and data.tag not in tags: continue create_and_allocate(data.handle) + data.is_asleep = False cpu_backup_tensor = data.cpu_backup_tensor if cpu_backup_tensor is None: diff --git a/vllm/distributed/ec_transfer/ec_connector/base.py b/vllm/distributed/ec_transfer/ec_connector/base.py index 3c20a4a1f749..9138203bf4dd 100644 --- a/vllm/distributed/ec_transfer/ec_connector/base.py +++ b/vllm/distributed/ec_transfer/ec_connector/base.py @@ -20,6 +20,8 @@ get_finished() - called with ids of finished requests, returns ids of requests that have completed async sending/recving. + build_connector_worker_meta() - builds metadata to be sent + back to the scheduler-side connector """ import enum @@ -56,6 +58,27 @@ class ECConnectorMetadata(ABC): # noqa: B024 pass +class ECConnectorWorkerMetadata(ABC): + """ + Abstract Metadata used to communicate back + Worker ECConnector -> Scheduler ECConnector. + + Each worker can output its own metadata. + For a single engine step, all metadata objects returned by workers + will be aggregated using the `aggregate` method below, before + being passed to the Scheduler ECConnector. + """ + + @abstractmethod + def aggregate( + self, other: "ECConnectorWorkerMetadata" + ) -> "ECConnectorWorkerMetadata": + """ + Aggregate metadata with another `ECConnectorWorkerMetadata` object. + """ + pass + + class ECConnectorBase(ABC): def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole): self._connector_metadata: ECConnectorMetadata | None = None @@ -190,6 +213,16 @@ def get_finished( """ return None, None + def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None: + """ + Build the ECConnector worker metadata for this engine step. + + Returns: + ECConnectorWorkerMetadata: the worker metadata. + None if no worker metadata is available. + """ + return None + # ============================== # Scheduler-side methods # ============================== diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py new file mode 100644 index 000000000000..f5f78e6c3c33 --- /dev/null +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EC connector helper utilities.""" + +from vllm.v1.outputs import ECConnectorOutput, ModelRunnerOutput + + +class ECOutputAggregator: + """Merge every worker's EC connector output onto the single + ModelRunnerOutput that reaches the scheduler. + + Mirrors KVOutputAggregator: only `output_rank`'s output is returned to the + scheduler, but the EC connector may have run on any rank. + """ + + def aggregate( + self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0 + ) -> ModelRunnerOutput | None: + output = outputs[output_rank] + if not output: + return None + + finished_sending = set[str]() + finished_recving = set[str]() + worker_meta = None + for model_runner_output in outputs: + assert model_runner_output is not None + ec_output = model_runner_output.ec_connector_output + if not ec_output: + continue + + finished_sending |= ec_output.finished_sending or set() + finished_recving |= ec_output.finished_recving or set() + + if meta := ec_output.ec_connector_worker_meta: + worker_meta = ( + meta if worker_meta is None else worker_meta.aggregate(meta) + ) + + aggregated = ECConnectorOutput( + finished_sending=finished_sending or None, + finished_recving=finished_recving or None, + ec_connector_worker_meta=worker_meta, + ) + if aggregated.is_empty(): + output.ec_connector_output = None + return output + + # `output` is the shared empty output whenever `output_rank` had no work, + # so attach through the copy-on-write helper. + return ModelRunnerOutput.with_ec_conn_output(output, aggregated) diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index f9a9a8a90a81..3f04b7bc91ed 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -33,6 +33,7 @@ from vllm.distributed.utils import is_weak_contiguous from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) @@ -222,10 +223,11 @@ def build_ops() -> None: # Wait for all D2H copies to finish # before issuing gloo batch_isend_irecv operations. - if self._cuda_stream is not None: - self._cuda_stream.synchronize() - else: - torch.cuda.current_stream().synchronize() + with gpu_sync_allowed(): + if self._cuda_stream is not None: + self._cuda_stream.synchronize() + else: + torch.cuda.current_stream().synchronize() reqs = batch_isend_irecv(p2p_ops) for req in reqs: diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 267beda9f188..5b2bb05ed0aa 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -47,6 +47,7 @@ from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MixtureOfExperts from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from .async_worker import start_async_worker from .eplb_communicator import EplbCommunicator, create_eplb_communicator @@ -596,10 +597,11 @@ def step( num_tokens_per_rank ) - # Just to make type checker happy - tokens_tensors: list[float] = torch.stack( - [avg_tokens_tensor, max_tokens_tensor] - ).tolist() + # This gpu/cpu sync only happens with expert stat logging enabled. + with gpu_sync_allowed(): + tokens_tensors: list[float] = torch.stack( + [avg_tokens_tensor, max_tokens_tensor] + ).tolist() avg_tokens, max_tokens = tokens_tensors balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 @@ -825,15 +827,17 @@ def rearrange( self.model_states.values(), global_expert_load_windows ): if not self.is_async or is_profile: - # Get new expert mappings for the model - new_physical_to_logical_map = self.policy.rebalance_experts( - global_expert_load_window.cpu(), - num_replicas, - num_groups, - num_nodes, - num_gpus, - eplb_model_state.physical_to_logical_map.cpu(), - ) + # Get new expert mappings for the model. The policy runs on the + # host, so the load window and current map have to come back. + with gpu_sync_allowed(): + new_physical_to_logical_map = self.policy.rebalance_experts( + global_expert_load_window.cpu(), + num_replicas, + num_groups, + num_nodes, + num_gpus, + eplb_model_state.physical_to_logical_map.cpu(), + ) skip_rearrange = False if ( @@ -1248,7 +1252,9 @@ def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: src_padding = dst.shape[-1] - src.shape[-1] assert src_padding >= 0 new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1) - dst.copy_(new_src) + # The map is committed from the host once per layer per rearrangement. + with gpu_sync_allowed(): + dst.copy_(new_src) def _commit_eplb_maps_for_layer( @@ -1276,10 +1282,7 @@ def _commit_eplb_maps_for_layer( num_logical_experts = model_state.logical_to_physical_map.shape[1] new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts) # Commit logical_to_physical_map - _pad_out_tensor( - src=new_logical, - dst=model_state.logical_to_physical_map[layer], - ) + _pad_out_tensor(src=new_logical, dst=model_state.logical_to_physical_map[layer]) # Commit logical_replica_count src = new_replica_count diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 53b0356dcd80..cbd513294565 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -16,6 +16,7 @@ from vllm.distributed.eplb.eplb_communicator import EplbCommunicator from vllm.distributed.eplb.eplb_utils import CpuGpuEvent from vllm.logger import init_logger +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) @@ -590,8 +591,11 @@ def rearrange_expert_weights_inplace( weights_buffer = list(expert_buffer) - old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy() - new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy() + # The per-layer transfer plan is built in Python, so both maps have to + # come back to the host. Once per rearrangement. + with gpu_sync_allowed(): + old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy() + new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy() for layer_idx in range(num_moe_layers): transfer_metadata = move_to_buffer( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index 93b7f7c3c1b6..782972b03a8e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py @@ -15,6 +15,7 @@ ) from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.hashing import safe_hash from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput @@ -243,7 +244,8 @@ def extract_kv_from_layer( layer_name, request.token_ids, request.mm_hashes ) kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping) - tensors = {"kv_cache": kv_cache.detach().cpu()} + with gpu_sync_allowed(): + tensors = {"kv_cache": kv_cache.detach().cpu()} safetensors.torch.save_file(tensors, filename) def wait_for_save(self): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index ca32c3345c38..297fc22ad8a0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py @@ -26,6 +26,7 @@ KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + KVConnectorWorkerMetadata, SupportsHMA, ) from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( @@ -37,6 +38,7 @@ from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig @@ -201,6 +203,14 @@ def update_state_after_alloc( request, blocks, num_external_tokens ) + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.bind_gpu_block_pool(gpu_block_pool) + + def has_pending_push_work(self) -> bool: + assert self.connector_scheduler is not None + return self.connector_scheduler.has_pending_push_work() + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -208,6 +218,10 @@ def build_connector_meta( assert self.connector_scheduler is not None return self.connector_scheduler.build_connector_meta(scheduler_output) + def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: + assert self.connector_worker is not None + return self.connector_worker.build_connector_worker_meta() + def request_finished( self, request: Request, @@ -220,8 +234,9 @@ def request_finished_all_groups( request: Request, block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - assert self.connector_scheduler is not None - return self.connector_scheduler.request_finished(request, block_ids) + # An in-flight store job holds its own reference on the blocks it reads, + # so a finishing request never has to defer freeing them. + return False, None def reset_cache(self) -> bool | None: """Reset the external Mooncake store on prefix-cache reset. @@ -241,6 +256,9 @@ def reset_cache(self) -> bool | None: return None def update_connector_output(self, connector_output: KVConnectorOutput): + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + kv_cache_events = connector_output.kv_cache_events if not kv_cache_events or not isinstance( kv_cache_events, MooncakeStoreKVEvents diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py index 6daa1e82ea6d..55ffc040cfd1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py @@ -6,7 +6,7 @@ """Data classes for MooncakeStoreConnector.""" from collections.abc import Iterable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import cast import numpy as np @@ -14,6 +14,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, + KVConnectorWorkerMetadata, ) from vllm.logger import init_logger from vllm.utils.math_utils import cdiv @@ -366,6 +367,10 @@ class ReqMeta: token_ids: list[int] | None = None num_prompt_tokens: int | None = None + # Identifies this store job for the engine's lifetime. A request id cannot + # serve that purpose: it is reused once a preempted request resumes, so it + # would release the wrong job's blocks. + store_job_id: int | None = None # Core-provided per-mamba-group # (group_id, cow_block_id, boundary_tokens) for this request's partial tail. # Present only on the producer's CoW step; drives the connector's offload @@ -434,6 +439,23 @@ def from_request_tracker( ) +@dataclass +class MooncakeStoreWorkerMetadata(KVConnectorWorkerMetadata): + """Maps ``ReqMeta.store_job_id`` to the number of ranks done with that job.""" + + completed_saves: dict[int, int] = field(default_factory=dict) + + def aggregate( + self, other: "KVConnectorWorkerMetadata" + ) -> "MooncakeStoreWorkerMetadata": + assert isinstance(other, MooncakeStoreWorkerMetadata) + for store_job_id, count in other.completed_saves.items(): + self.completed_saves[store_job_id] = ( + self.completed_saves.get(store_job_id, 0) + count + ) + return self + + class MooncakeStoreConnectorMetadata(KVConnectorMetadata): """Metadata passed from scheduler to worker.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 42c6f3fa99af..f583b3fcfdf1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -5,8 +5,6 @@ # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/). """Scheduler-side logic for MooncakeStoreConnector.""" -from typing import Any - from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorMetadata, @@ -17,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501 LoadSpec, MooncakeStoreConnectorMetadata, + MooncakeStoreWorkerMetadata, ReqMeta, RequestTracker, ) @@ -24,10 +23,12 @@ LookupKeyClient, ) from vllm.logger import init_logger +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import Request logger = init_logger(__name__) @@ -78,6 +79,15 @@ def __init__( self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {} self._unfinished_request_ids: set[str] = set() + self._gpu_block_pool: BlockPool | None = None + self._num_workers = vllm_config.parallel_config.world_size + self._next_store_job_id = 0 + # store_job_id -> (referenced block ids, ranks yet to report completion) + self._pinned_saves: dict[int, tuple[list[int], int]] = {} + + def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None: + self._gpu_block_pool = gpu_block_pool + def get_num_new_matched_tokens( self, request: Request, @@ -392,33 +402,75 @@ def build_connector_meta( ) ) + self._reference_save_blocks(meta) return meta - def request_finished( - self, - request: Request, - block_ids: tuple[list[int], ...], - ) -> tuple[bool, dict[str, Any] | None]: - """Determine whether to delay freeing blocks for async save.""" - if self.kv_role == "kv_consumer": - return False, None - tracker = self._request_trackers.get(request.request_id) - # Missing tracker can happen when the request is aborted before the - # connector observes the normal finished lifecycle or is preempted - # before finishing. - if tracker is None or ( - tracker.num_saved_tokens <= 0 and not tracker.has_pending_offload - ): - return False, None - total_blocks = sum(len(g) for g in block_ids) - delay_free_blocks = total_blocks > 0 - if delay_free_blocks: - logger.debug( - "Delaying free of %d blocks for request %s", - total_blocks, - request.request_id, + def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None: + """Take a GPU block reference for every store job this step emits. + + The worker DMAs out of these blocks after the step that scheduled them, + so a reference keeps them out of the free queue even once the request + itself is freed, until every rank reports the job done. + """ + pool = self._gpu_block_pool + for req_meta in meta.requests: + if not req_meta.can_save: + continue + assert pool is not None, ( + "GPU block pool must be bound before any store job is emitted" + ) + req_meta.store_job_id = store_job_id = self._next_store_job_id + self._next_store_job_id += 1 + block_ids: list[int] = [] + if req_meta.partial_tail_offloads: + # A partial-tail CoW block is deliberately kept out of the + # request's block table, so it is absent from `block_ids` even + # though the worker DMAs out of it just as asynchronously. + # It leads the list, as in `pop_blocks_for_free`. + block_ids += [bid for _, bid, _ in req_meta.partial_tail_offloads] + # Every allocated block is referenced, not just the ones covering + # this job's token range: a rank resumes from its own last + # successful offset, which lags the scheduler's whenever a save was + # skipped or failed, so it may read anywhere below the range. + block_ids += [bid for group in req_meta.block_ids for bid in group] + if not block_ids: + continue + self._pinned_saves[store_job_id] = (block_ids, self._num_workers) + pool.touch([pool.blocks[bid] for bid in block_ids]) + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + """Drop the block references of store jobs every rank has finished.""" + meta = connector_output.kv_connector_worker_meta + if not isinstance(meta, MooncakeStoreWorkerMetadata): + return + pool = self._gpu_block_pool + assert pool is not None + for store_job_id, count in meta.completed_saves.items(): + pinned = self._pinned_saves.get(store_job_id) + if pinned is None: + # The job referenced no blocks, so nothing was recorded for it. + continue + block_ids, remaining = pinned + remaining -= count + if remaining > 0: + self._pinned_saves[store_job_id] = (block_ids, remaining) + continue + assert remaining == 0, ( + f"store job {store_job_id} reported by too many ranks" ) - return delay_free_blocks, None + del self._pinned_saves[store_job_id] + # Tail-first, as elsewhere, so the shared prefix is evicted last. + pool.free_blocks(pool.blocks[bid] for bid in reversed(block_ids)) + + def has_pending_push_work(self) -> bool: + """Keep the engine stepping while any store job still holds block refs. + + Completions only reach the scheduler as worker metadata on a step, so an + engine that quiesced with jobs in flight would leave those references + held indefinitely. Nothing else keeps it alive now that a finishing + request no longer defers its own free. + """ + return bool(self._pinned_saves) def reset_store(self) -> bool: """Trigger a global ``remove_all(force=True)`` on the Mooncake master. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 762e5eba9263..83c6d16e1b9e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -17,7 +17,6 @@ import socket import threading import time -from collections import defaultdict from collections.abc import Callable, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass @@ -46,6 +45,7 @@ ChunkedTokenDatabase, KeyMetadata, MooncakeStoreConnectorMetadata, + MooncakeStoreWorkerMetadata, PoolKey, ReqMeta, ) @@ -506,7 +506,16 @@ def __init__( self.group_put_steps = group_put_steps self.coord = coord self.kv_role = kv_role - self.stored_requests: defaultdict[str, int] = defaultdict(int) + # req_id -> ids of its store jobs that are still queued or running. + # Keying by store_job_id, which never repeats for the engine's lifetime, + # rather than counting jobs per request id makes the ledger immune to id + # reuse across preemption: a job left over from a retired generation is + # missing from the set its resumed generation builds, so it can no longer + # retire that generation, rewind its resume offset, or mark it skipped. + self.stored_requests: dict[str, set[int]] = {} + # store_job_id -> times this rank finished with it, drained every step + # so the scheduler can release the blocks it referenced for those jobs. + self._completed_saves: dict[int, int] = {} self.enable_kv_event = enable_kv_event # Caller always passes a non-None ReplicateConfig — see # MooncakeStoreWorker.__init__ where store_replicate_config is built. @@ -522,14 +531,20 @@ def __init__( # batch resumes here, so pressure-skipped or failed ranges are retried. self._saved_offset: dict[str, int] = {} - def add_stored_request(self, req_id: str): + def add_request(self, request: ReqMeta) -> None: + # Register before enqueueing so a job is never picked up unledgered. + assert request.store_job_id is not None with self.done_task_lock: - self.stored_requests[req_id] += 1 + self.stored_requests.setdefault(request.req_id, set()).add( + request.store_job_id + ) + super().add_request(request) - def dec_stored_request(self, req_id: str): + def is_live_store_job(self, req_meta: ReqMeta) -> bool: with self.done_task_lock: - if req_id in self.stored_requests: - self.stored_requests[req_id] -= 1 + return req_meta.store_job_id in self.stored_requests.get( + req_meta.req_id, () + ) def delete_finished_stored_request(self, req_id: str): with self.done_task_lock: @@ -538,21 +553,51 @@ def delete_finished_stored_request(self, req_id: str): self._skip_store_requests.discard(req_id) self._saved_offset.pop(req_id, None) - def _record_saved(self, req_id: str, token_len: int) -> None: - # Guard on liveness so a concurrent finish/preempt pop isn't recreated. + def finish_store_job(self, req_meta: ReqMeta) -> None: + """Retire a job from the ledger and report its blocks as no longer read. + + Every path out of a job must reach this, skips and failures included: a + job that never reports leaves its blocks referenced for the rest of the + run. The discard is a no-op for a job whose generation already retired. + """ + store_job_id = req_meta.store_job_id + assert store_job_id is not None, ( + "a queued store job always carries a store_job_id" + ) with self.done_task_lock: - if req_id in self.stored_requests: - self._saved_offset[req_id] = token_len + live = self.stored_requests.get(req_meta.req_id) + if live is not None: + live.discard(store_job_id) + self._completed_saves[store_job_id] = ( + self._completed_saves.get(store_job_id, 0) + 1 + ) + + def take_completed_saves(self) -> dict[int, int]: + with self.done_task_lock: + completed = self._completed_saves + self._completed_saves = {} + return completed + + def _record_saved(self, req_meta: ReqMeta, token_len: int) -> None: + # Guard on job liveness so neither a concurrent finish/preempt pop nor a + # stale job's offset is written back over the live generation's. + with self.done_task_lock: + if req_meta.store_job_id in self.stored_requests.get(req_meta.req_id, ()): + self._saved_offset[req_meta.req_id] = token_len def _should_skip_request(self, req_id: str) -> bool: with self.done_task_lock: return self._store_pressure_active and req_id in self._skip_store_requests - def _mark_request_skipped_for_pressure(self, req_id: str) -> bool: + def _mark_request_skipped_for_pressure(self, req_meta: ReqMeta) -> bool: + req_id = req_meta.req_id with self.done_task_lock: already_skipped = req_id in self._skip_store_requests self._store_pressure_active = True - self._skip_store_requests.add(req_id) + # The pressure itself is global, but only a live job may sentence its + # own request to being skipped. + if req_meta.store_job_id in self.stored_requests.get(req_id, ()): + self._skip_store_requests.add(req_id) return already_skipped def _clear_store_pressure(self) -> bool: @@ -727,7 +772,7 @@ def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: failed_codes, ) if MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes: - self._mark_request_skipped_for_pressure(req_meta.req_id) + self._mark_request_skipped_for_pressure(req_meta) return False if self._clear_store_pressure(): @@ -738,22 +783,20 @@ def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool: return True def _handle_request(self, req_meta: ReqMeta): - # Cache hits are always a multiple of ``lcm_block_size`` tokens, which - # is also ``store_mask``'s precondition. - lcm_block_size = self.coord.lcm_block_size - token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size - block_ids_per_group = req_meta.block_ids - req_id = req_meta.req_id - current_event = req_meta.current_event - - if req_id not in self.stored_requests: - self.request_queue.task_done() - return - - # Decrement the in-flight counter and signal task_done() in `finally` - # so the scheduler can release the GPU blocks it pinned for this - # request (via `delay_free_blocks`) even when the store path raises. + # The single `finally` is the only way out, so the scheduler releases + # this job's GPU block references however the job ends. try: + # Cache hits are always a multiple of ``lcm_block_size`` tokens, + # which is also ``store_mask``'s precondition. + lcm_block_size = self.coord.lcm_block_size + token_len = req_meta.token_len_chunk // lcm_block_size * lcm_block_size + block_ids_per_group = req_meta.block_ids + req_id = req_meta.req_id + current_event = req_meta.current_event + + if not self.is_live_store_job(req_meta): + return + if self._should_skip_request(req_id): logger.debug( "Skipping Mooncake store for request %s while CPU/disk " @@ -808,7 +851,7 @@ def _handle_request(self, req_meta: ReqMeta): group_indices.append(g_idx) if not keys: - self._record_saved(req_id, token_len) + self._record_saved(req_meta, token_len) return # Check which blocks already exist (dedup) @@ -834,7 +877,7 @@ def _handle_request(self, req_meta: ReqMeta): ] if not missing_indices: - self._record_saved(req_id, token_len) + self._record_saved(req_meta, token_len) return if len(missing_indices) != len(keys): @@ -954,7 +997,7 @@ def _handle_request(self, req_meta: ReqMeta): ) if ( MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes - and not self._mark_request_skipped_for_pressure(req_id) + and not self._mark_request_skipped_for_pressure(req_meta) ): logger.warning( "Detected Mooncake CPU/disk offloading pressure " @@ -964,7 +1007,7 @@ def _handle_request(self, req_meta: ReqMeta): req_id, ) else: - self._record_saved(req_id, token_len) + self._record_saved(req_meta, token_len) if self._clear_store_pressure(): logger.info( "Mooncake CPU/disk offloading pressure cleared " @@ -984,7 +1027,7 @@ def _handle_request(self, req_meta: ReqMeta): if self.enable_kv_event and stored_events: self.update_kv_event(stored_events) finally: - self.dec_stored_request(req_id) + self.finish_store_job(req_meta) self.request_queue.task_done() @@ -1390,7 +1433,7 @@ def __init__( scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, use_eagle=use_eagle, - retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL, + retention_interval=kv_cache_config.prefix_cache_retention_interval, ) # One ChunkedTokenDatabase per group; addresses populated in # register_kv_caches once the kv-cache layout is known. Each group's @@ -1670,16 +1713,15 @@ def get_finished( continue request.current_event = current_event assert self.kv_send_thread is not None - self.kv_send_thread.add_stored_request(request.req_id) self.kv_send_thread.add_request(request) - # Check completion of previously queued transfers - done_sending = ( - self._get_and_clear_finished_sending(finished_req_ids, meta) - if self.kv_role in ["kv_producer", "kv_both"] - else set() - ) + if self.kv_role in ["kv_producer", "kv_both"]: + self._close_ended_store_requests(finished_req_ids, meta) + # Blocks read by a store job are released by the scheduler when the job + # reports back (see build_connector_worker_meta), so no request ever waits + # on a `finished_sending` signal to get its blocks back. + done_sending: set[str] = set() done_recving: set[str] = set() if self.load_async: for recv_thread in self.kv_recv_threads: @@ -1727,35 +1769,37 @@ def get_kv_connector_stats(self) -> MooncakeStoreConnectorStats | None: self.kv_connector_stats = MooncakeStoreConnectorStats() return kv_connector_stats - def _get_and_clear_finished_sending( + def _close_ended_store_requests( self, finished_req_ids: set[str], meta: MooncakeStoreConnectorMetadata, - ) -> set[str]: + ) -> None: + """Retire the ledger entries of requests that finished or were preempted. + + An entry may only go once its jobs have drained, because they still read + the resume offset it owns; a request that comes back after preemption + then saves from the start rather than from where the last attempt got to. + """ assert self.kv_send_thread is not None - finished_sending: set[str] = set() for req_id in meta.preempted_req_ids: self.kv_send_thread.delete_finished_stored_request(req_id) - for req_id in self.kv_send_thread.stored_requests.copy(): - if ( - self.kv_send_thread.stored_requests[req_id] == 0 - and req_id in self.finished_store_req - ): - self.finished_store_req.remove(req_id) - finished_sending.add(req_id) - self.kv_send_thread.delete_finished_stored_request(req_id) - - for req_id in finished_req_ids: - req_remain_jobs = self.kv_send_thread.stored_requests.get(req_id) - if req_remain_jobs == 0: - finished_sending.add(req_id) - self.kv_send_thread.delete_finished_stored_request(req_id) - elif req_remain_jobs is not None: + for req_id in finished_req_ids | self.finished_store_req: + if self.kv_send_thread.stored_requests.get(req_id): + # Queued jobs still need the resume offset; retire on a later step. self.finished_store_req.add(req_id) + else: + self.finished_store_req.discard(req_id) + self.kv_send_thread.delete_finished_stored_request(req_id) - return finished_sending + def build_connector_worker_meta(self) -> MooncakeStoreWorkerMetadata | None: + if self.kv_send_thread is None: + return None + completed_saves = self.kv_send_thread.take_completed_saves() + if not completed_saves: + return None + return MooncakeStoreWorkerMetadata(completed_saves=completed_saves) def lookup(self, num_tokens: int, block_hashes: Sequence[BlockHash]) -> int: """Check how many prefix tokens exist in the store. 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 afa16279ff1c..bd93c2323379 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 @@ -51,6 +51,11 @@ class NixlBaseConnectorScheduler: """Base implementation of Scheduler side methods shared by pull and push.""" + # Emitted in kv_transfer_params so an external router can distinguish a + # pull (READ) producer from a push (WRITE) one. Overridden by the push + # scheduler. + _TRANSFER_MODE: str = "pull" + def __init__( self, vllm_config: "VllmConfig", 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 815eaba7decd..d2c073676de9 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 @@ -67,6 +67,7 @@ ) from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import get_kv_cache_layout @@ -90,6 +91,11 @@ class NixlBaseConnectorWorker: """Base implementation of Worker side methods shared by pull and push.""" + # Transfer mode included in the NIXL compatibility hash so that a push + # (WRITE) connector and a pull (READ) connector never handshake together. + # Overridden by NixlPushConnectorWorker. + _TRANSFER_MODE: str = "pull" + def _compute_desc_ids( self, block_ids: BlockIds, @@ -976,6 +982,7 @@ def _register_packed_kv_cache( self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks, + transfer_mode=self._TRANSFER_MODE, ) total_size = storage.nbytes() @@ -1066,7 +1073,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): is_mamba=self._has_mamba, ) self.compat_hash = compute_nixl_compatibility_hash( - self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks + self.vllm_config, + self.backend_name, + self.transfer_topo.cross_layers_blocks, + transfer_mode=self._TRANSFER_MODE, ) if self.use_host_buffer: @@ -1873,14 +1883,16 @@ def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta): local_block_ids = meta.local_physical_block_ids # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups - for group_block_ids in local_block_ids: - self.copy_blocks( - self.host_xfer_buffers, - self.device_kv_caches, - group_block_ids, - group_block_ids, - "h2d", - ) + # The h2d block copies below are intentionally synchronous. + with gpu_sync_allowed(): + for group_block_ids in local_block_ids: + self.copy_blocks( + self.host_xfer_buffers, + self.device_kv_caches, + group_block_ids, + group_block_ids, + "h2d", + ) if logger.isEnabledFor(logging.DEBUG): logger.debug( "synced recved kv of request[%s] to device kv buffer," @@ -1894,26 +1906,28 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): assert self.use_host_buffer assert self.copy_blocks is not None - for req_id, meta in metadata.reqs_to_save.items(): - meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids, self._physical_blocks_per_logical_kv_block - ) - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "save_load_kv for request[%s] to host xfer buffer." - "local_block_ids: %s. ", - req_id, - ",".join(map(str, meta.local_physical_block_ids)), - ) - # blocking - for group_block_ids in meta.local_physical_block_ids: - self.copy_blocks( - self.device_kv_caches, - self.host_xfer_buffers, - group_block_ids, - group_block_ids, - "d2h", + # The d2h block copies below are intentionally synchronous. + with gpu_sync_allowed(): + for req_id, meta in metadata.reqs_to_save.items(): + meta.local_physical_block_ids = self._logical_to_kernel_block_ids( + meta.local_block_ids, self._physical_blocks_per_logical_kv_block ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "save_load_kv for request[%s] to host xfer buffer." + "local_block_ids: %s. ", + req_id, + ",".join(map(str, meta.local_physical_block_ids)), + ) + # blocking + for group_block_ids in meta.local_physical_block_ids: + self.copy_blocks( + self.device_kv_caches, + self.host_xfer_buffers, + group_block_ids, + group_block_ids, + "d2h", + ) @cached_property def _attention_kv_caches(self) -> list[torch.Tensor]: 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 4758bab5d73c..ed2e030c7341 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -42,8 +42,9 @@ # 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake # clock-sync timestamp # 6: Validate EAGLE/MTP speculative configuration compatibility +# 7: Include NIXL transfer mode (push vs pull) in the compatibility hash # -NIXL_CONNECTOR_VERSION: int = 6 +NIXL_CONNECTOR_VERSION: int = 7 @dataclass @@ -123,7 +124,10 @@ def _get_speculative_compatibility_factors( def compute_nixl_compatibility_hash( - vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool + vllm_config: VllmConfig, + attn_backend_name: str, + cross_layers_blocks: bool, + transfer_mode: str = "pull", ) -> str: """ Compute compatibility hash for NIXL KV transfer. @@ -137,6 +141,11 @@ def compute_nixl_compatibility_hash( - KV cache format (dtype, sliding window) - Attention backend - EAGLE/MTP configuration that affects transferred state + - Transfer mode (push vs pull) + + The transfer mode is included because the push (WRITE) and pull (READ) + connectors use incompatible transfer protocols; a push connector and a + pull connector must never complete a handshake with each other. Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout are validated at runtime in _validate_remote_agent_handshake and are not @@ -171,6 +180,8 @@ def compute_nixl_compatibility_hash( "cross_layers_blocks": cross_layers_blocks, "is_hma_enabled": is_hma_enabled, "speculative_config": _get_speculative_compatibility_factors(vllm_config), + # push (WRITE) and pull (READ) connectors are protocol-incompatible + "transfer_mode": transfer_mode, } compat_hash = hash_factors(factors) 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 9beceeabada6..d0a5383c5c5f 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 @@ -277,4 +277,5 @@ def request_finished( tp_size=self.vllm_config.parallel_config.tensor_parallel_size, remote_num_tokens=remote_num_tokens, remote_blocks_expiry_time=blocks_expiry_time, + transfer_mode=self._TRANSFER_MODE, ) 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 a02491e4872e..6efe491be863 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 @@ -60,6 +60,8 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler): hooks. """ + _TRANSFER_MODE: str = "push" + def __init__( self, vllm_config: VllmConfig, @@ -290,6 +292,7 @@ def request_finished( tp_size=self.vllm_config.parallel_config.tensor_parallel_size, pp_size=self.vllm_config.parallel_config.pipeline_parallel_size, remote_num_tokens=remote_num_tokens, + transfer_mode=self._TRANSFER_MODE, ) def build_connector_meta( 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 7fdde27dcc27..0721bf9d1dd9 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 @@ -77,6 +77,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): """Push-specific (WRITE) worker logic. See module docstring.""" + # Distinguishes push from pull in the NIXL compatibility hash. + _TRANSFER_MODE: str = "push" + def __init__( self, vllm_config: "VllmConfig", diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 4e4448a09c60..1fb1388f100c 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -522,6 +522,9 @@ class EngineArgs: prefix_caching_hash_algo: PrefixCachingHashAlgo = ( CacheConfig.prefix_caching_hash_algo ) + prefix_cache_retention_interval: int | None = get_field( + CacheConfig, "prefix_cache_retention_interval" + ) disable_sliding_window: bool = ModelConfig.disable_sliding_window disable_cascade_attn: bool = ModelConfig.disable_cascade_attn offload_backend: str = OffloadConfig.offload_backend @@ -700,7 +703,6 @@ class EngineArgs: ModelConfig, "override_generation_config" ) model_impl: str = ModelConfig.model_impl - override_attention_dtype: str | None = ModelConfig.override_attention_dtype attention_backend: AttentionBackendEnum | None = AttentionConfig.backend kv_cache_dtype_skip_layers: list[str] = get_field( @@ -917,9 +919,6 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-cumem-allocator", **model_kwargs["enable_cumem_allocator"] ) model_group.add_argument("--model-impl", **model_kwargs["model_impl"]) - model_group.add_argument( - "--override-attention-dtype", **model_kwargs["override_attention_dtype"] - ) model_group.add_argument( "--logits-processors", **model_kwargs["logits_processors"] ) @@ -1228,6 +1227,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: cache_group.add_argument( "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"] ) + cache_group.add_argument( + "--prefix-cache-retention-interval", + **cache_kwargs["prefix_cache_retention_interval"], + ) cache_group.add_argument( "--kv-cache-dtype-skip-layers", **cache_kwargs["kv_cache_dtype_skip_layers"] ) @@ -1787,7 +1790,6 @@ def create_model_config(self) -> ModelConfig: enable_sleep_mode=self.enable_sleep_mode, enable_cumem_allocator=self.enable_cumem_allocator, model_impl=self.model_impl, - override_attention_dtype=self.override_attention_dtype, logits_processors=self.logits_processors, video_pruning_rate=self.video_pruning_rate, video_pruning_method=self.video_pruning_method, @@ -2006,6 +2008,7 @@ def create_engine_config( sliding_window=sliding_window, enable_prefix_caching=self.enable_prefix_caching, prefix_caching_hash_algo=self.prefix_caching_hash_algo, + prefix_cache_retention_interval=self.prefix_cache_retention_interval, kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers, kv_sharing_fast_prefill=self.kv_sharing_fast_prefill, mamba_cache_dtype=self.mamba_cache_dtype, @@ -2158,7 +2161,7 @@ def create_engine_config( elif self.data_parallel_size_local is not None: data_parallel_size_local = self.data_parallel_size_local - if self.data_parallel_start_rank and not headless: + if self.data_parallel_start_rank is not None and not headless: # Infer hybrid LB mode. self.data_parallel_hybrid_lb = True @@ -2177,7 +2180,9 @@ def create_engine_config( self.data_parallel_hybrid_lb = False self.data_parallel_rank = ( - self.data_parallel_start_rank or inferred_data_parallel_rank + self.data_parallel_start_rank + if self.data_parallel_start_rank is not None + else inferred_data_parallel_rank ) if self.nnodes > 1: logger.info( diff --git a/vllm/entrypoints/anthropic/api_router.py b/vllm/entrypoints/anthropic/api_router.py index 414f87f308ae..584a2958018e 100644 --- a/vllm/entrypoints/anthropic/api_router.py +++ b/vllm/entrypoints/anthropic/api_router.py @@ -17,9 +17,11 @@ ) from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, - sanitize_message, validate_json_request, with_cancellation, ) @@ -71,15 +73,7 @@ async def create_messages(request: AnthropicMessagesRequest, raw_request: Reques generator = await handler.create_messages(request, raw_request) except Exception as e: logger.exception("Error in create_messages: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=AnthropicErrorResponse( - error=AnthropicError( - type="internal_error", - message=sanitize_message(str(e)), - ) - ).model_dump(), - ) + return translate_error_response(create_error_response(e)) if isinstance(generator, ErrorResponse): return translate_error_response(generator) @@ -117,15 +111,7 @@ async def count_tokens(request: AnthropicCountTokensRequest, raw_request: Reques response = await handler.count_tokens(request, raw_request) except Exception as e: logger.exception("Error in count_tokens: %s", e) - return JSONResponse( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - content=AnthropicErrorResponse( - error=AnthropicError( - type="internal_error", - message=sanitize_message(str(e)), - ) - ).model_dump(), - ) + return translate_error_response(create_error_response(e)) if isinstance(response, ErrorResponse): return translate_error_response(response) diff --git a/vllm/entrypoints/anthropic/protocol.py b/vllm/entrypoints/anthropic/protocol.py index 80c354dacbbc..02b9356113ca 100644 --- a/vllm/entrypoints/anthropic/protocol.py +++ b/vllm/entrypoints/anthropic/protocol.py @@ -95,6 +95,7 @@ class AnthropicToolChoice(BaseModel): type: Literal["auto", "any", "tool", "none"] name: str | None = None + disable_parallel_tool_use: bool | None = None @model_validator(mode="after") def validate_name_required_for_tool(self) -> "AnthropicToolChoice": diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 7c3b0597b5c4..4ec5c14e019b 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -46,7 +46,7 @@ UsageInfo, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.utils.api_utils import sanitize_message +from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.renderers.online_renderer import OnlineRenderer @@ -539,6 +539,9 @@ def _convert_tool_choice( req.tool_choice = None return + req.parallel_tool_calls = ( + not anthropic_request.tool_choice.disable_parallel_tool_use + ) tool_choice_type = anthropic_request.tool_choice.type if tool_choice_type == "auto": req.tool_choice = "auto" diff --git a/vllm/entrypoints/cohere/api_router.py b/vllm/entrypoints/cohere/api_router.py index b532e60920a8..85eee0fd5548 100644 --- a/vllm/entrypoints/cohere/api_router.py +++ b/vllm/entrypoints/cohere/api_router.py @@ -33,9 +33,9 @@ import vllm.envs as envs from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, - sanitize_message, validate_json_request, with_cancellation, ) diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index 3873f7043316..b285319e7e20 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -86,7 +86,7 @@ StreamOptions, ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.utils.api_utils import sanitize_message +from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.parser.abstract_parser import Parser from vllm.renderers.cohere import MESSAGES_CITATIONS_KEY, POSITION_TO_SOURCE_KEY diff --git a/vllm/entrypoints/launchers/__init__.py b/vllm/entrypoints/launchers/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/launchers/api_server/__init__.py b/vllm/entrypoints/launchers/api_server/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py new file mode 100644 index 000000000000..39d670c019b7 --- /dev/null +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from argparse import Namespace + +from fastapi import FastAPI + +from vllm import envs +from vllm.config import ModelConfig +from vllm.tasks import POOLING_TASKS, SupportedTask + + +def register_api_routers( + args: Namespace, + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], + model_config: ModelConfig | None = None, +): + from vllm.entrypoints.serve import register_vllm_serve_api_routers + + register_vllm_serve_api_routers(app) + + from vllm.entrypoints.openai.models.api_router import ( + attach_router as register_models_api_router, + ) + + register_models_api_router(app) + + from vllm.entrypoints.serve.sagemaker.api_router import ( + attach_router as register_sagemaker_api_router, + ) + + register_sagemaker_api_router(app, supported_tasks, model_config) + + if envs.VLLM_SERVER_DEV_MODE: + from vllm.entrypoints.serve import register_vllm_dev_api_routers + + register_vllm_dev_api_routers(app) + + if "generate" in supported_tasks: + from vllm.entrypoints.generate.api_router import ( + register_generate_api_routers, + ) + + register_generate_api_routers(app) + + from vllm.entrypoints.serve.elastic_ep.api_router import ( + attach_router as elastic_ep_attach_router, + ) + + elastic_ep_attach_router(app) + + if "generate" in supported_tasks or "render" in supported_tasks: + from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers + + register_scale_out_api_routers(app, supported_tasks) + + if "transcription" in supported_tasks or "realtime" in supported_tasks: + from vllm.entrypoints.speech_to_text.factories import ( + register_speech_to_text_api_routers, + ) + + register_speech_to_text_api_routers(app, supported_tasks) + + if any(task in POOLING_TASKS for task in supported_tasks): + from vllm.entrypoints.pooling.factories import register_pooling_api_routers + + register_pooling_api_routers(app, supported_tasks, model_config) + + if getattr(args, "enable_fault_tolerance", False): + from vllm.entrypoints.serve.fault_tolerance.api_router import ( + register_fault_tolerance_api_router, + ) + + register_fault_tolerance_api_router(app) diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 5a6dd8c83dd0..c541f03f445e 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import importlib -import inspect import multiprocessing import multiprocessing.forkserver as forkserver import os @@ -16,9 +14,7 @@ from typing import Any, cast import uvloop -from fastapi import FastAPI, HTTPException -from fastapi.exceptions import RequestValidationError -from fastapi.middleware.cors import CORSMiddleware +from fastapi import FastAPI from starlette.datastructures import State import vllm.envs as envs @@ -27,10 +23,12 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http +from vllm.entrypoints.launchers.api_server.routers import register_api_routers from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware +from vllm.entrypoints.serve.exception_handling.register import init_exception_handler +from vllm.entrypoints.serve.middleware.register import init_entrypoints_middleware from vllm.entrypoints.serve.sagemaker.api_router import sagemaker_standards_bootstrap from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( @@ -41,15 +39,9 @@ ) from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.entrypoints.serve.utils.server_utils import ( - exception_handler, get_uvicorn_log_config, - http_exception_handler, lifespan, - log_response, - validation_exception_handler, - vllm_error_handler, ) -from vllm.exceptions import VLLMError from vllm.logger import init_logger from vllm.reasoning import ReasoningParserManager from vllm.renderers.online_derenderer import OnlineDerenderer @@ -210,64 +202,9 @@ def build_app( else: app = FastAPI(lifespan=lifespan) app.state.args = args + app.root_path = args.root_path - from vllm.entrypoints.serve import register_vllm_serve_api_routers - - register_vllm_serve_api_routers(app) - - from vllm.entrypoints.openai.models.api_router import ( - attach_router as register_models_api_router, - ) - - register_models_api_router(app) - - from vllm.entrypoints.serve.sagemaker.api_router import ( - attach_router as register_sagemaker_api_router, - ) - - register_sagemaker_api_router(app, supported_tasks, model_config) - - if envs.VLLM_SERVER_DEV_MODE: - from vllm.entrypoints.serve import register_vllm_dev_api_routers - - register_vllm_dev_api_routers(app) - - if "generate" in supported_tasks: - from vllm.entrypoints.generate.api_router import ( - register_generate_api_routers, - ) - - register_generate_api_routers(app) - - from vllm.entrypoints.serve.elastic_ep.api_router import ( - attach_router as elastic_ep_attach_router, - ) - - elastic_ep_attach_router(app) - - if "generate" in supported_tasks or "render" in supported_tasks: - from vllm.entrypoints.scale_out.factories import register_scale_out_api_routers - - register_scale_out_api_routers(app, supported_tasks) - - if "transcription" in supported_tasks or "realtime" in supported_tasks: - from vllm.entrypoints.speech_to_text.factories import ( - register_speech_to_text_api_routers, - ) - - register_speech_to_text_api_routers(app, supported_tasks) - - if any(task in POOLING_TASKS for task in supported_tasks): - from vllm.entrypoints.pooling.factories import register_pooling_api_routers - - register_pooling_api_routers(app, supported_tasks, model_config) - - if args.enable_fault_tolerance: - from vllm.entrypoints.serve.fault_tolerance.api_router import ( - register_fault_tolerance_api_router, - ) - - register_fault_tolerance_api_router(app) + register_api_routers(args, app, supported_tasks, model_config) # Endpoint plugins are attached last so their routes are registered after all core # routers. This runs even for the CPU only render server. A plugin eligible for @@ -275,79 +212,8 @@ def build_app( # `engine_client=None` at Phase B (see `_init_endpoint_plugins_state`). _attach_endpoint_plugins(app, supported_tasks) - app.root_path = args.root_path - app.add_middleware( - CORSMiddleware, - allow_origins=args.allowed_origins, - allow_credentials=args.allow_credentials, - allow_methods=args.allowed_methods, - allow_headers=args.allowed_headers, - ) - - # Exception handlers are registered in four layers: - # 1. framework errors raised by FastAPI/Starlette - # 2. vLLM-specific errors dispatched via a single ``VLLMError`` handler - # 3. fallback handlers for raw exceptions not yet migrated to ``VLLMError`` - # 4. the raw ``Exception`` handler as a safety net - # Registering specific exception types (rather than only ``Exception``) - # ensures they are handled by ``ExceptionMiddleware`` (inside the Prometheus - # middleware) rather than ``ServerErrorMiddleware`` (outside it), so their - # status codes are recorded correctly. - app.exception_handler(HTTPException)(http_exception_handler) - app.exception_handler(RequestValidationError)(validation_exception_handler) - - app.exception_handler(VLLMError)(vllm_error_handler) - - # TODO(zqzten): remove these fallback handlers after migration to VLLMError - app.exception_handler(ValueError)(exception_handler) - app.exception_handler(TypeError)(exception_handler) - app.exception_handler(OverflowError)(exception_handler) - app.exception_handler(NotImplementedError)(exception_handler) - - app.exception_handler(Exception)(exception_handler) - - # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY - if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]: - from vllm.entrypoints.serve.utils.server_utils import AuthenticationMiddleware - - app.add_middleware(AuthenticationMiddleware, tokens=tokens) - - if args.enable_request_id_headers: - from vllm.entrypoints.serve.utils.server_utils import XRequestIdMiddleware - - app.add_middleware(XRequestIdMiddleware) - - # Add scaling middleware to check for scaling state - app.add_middleware(ScalingMiddleware) - - if "realtime" in supported_tasks: - # Add WebSocket metrics middleware - from vllm.entrypoints.speech_to_text.factories import ( - add_websocket_metrics_middleware, - ) - - add_websocket_metrics_middleware(app) - - if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE: - logger.warning( - "CAUTION: Enabling log response in the API Server. " - "This can include sensitive information and should be " - "avoided in production." - ) - app.middleware("http")(log_response) - - for middleware in args.middleware: - module_path, object_name = middleware.rsplit(".", 1) - imported = getattr(importlib.import_module(module_path), object_name) - if inspect.isclass(imported): - app.add_middleware(imported) # type: ignore[arg-type] - elif inspect.iscoroutinefunction(imported): - app.middleware("http")(imported) - else: - raise ValueError( - f"Invalid middleware {middleware}. Must be a function or a class." - ) - + init_exception_handler(app) + init_entrypoints_middleware(args, app, supported_tasks) app = sagemaker_standards_bootstrap(app) return app diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 13b81331728d..e8bb4e567826 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -114,7 +114,7 @@ class ChatCompletionResponseChoice(OpenAIBaseModel): token_ids: list[int] | None = None # Per-token expert routing decisions, base64-encoded ``.npy`` bytes # (numpy serialization). Shape after decode: - # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32 # ``num_tokens - 1`` because the last sampled token has not been # forwarded yet and therefore has no routing data. # Decode: @@ -422,6 +422,11 @@ class ChatCompletionRequest(OpenAIBaseModel): "need to map generated text back to input tokens." ), ) + routed_experts_prompt_start: int = Field( + default=0, + ge=0, + description="Skip the first N prompt tokens from returned routed-expert data.", + ) return_token_offsets: bool | None = Field( default=False, description=( @@ -741,6 +746,7 @@ def to_sampling_params( extra_args=extra_args or None, skip_clone=True, # Created fresh per request, safe to skip clone repetition_detection=self.repetition_detection, + routed_experts_prompt_start=self.routed_experts_prompt_start, ) @model_validator(mode="before") diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index f2d4936f1e67..28ec1e8a1a59 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -2,15 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import io import time from collections.abc import AsyncGenerator, AsyncIterator from collections.abc import Sequence as GenericSequence from http import HTTPStatus from typing import Any, Final, cast -import numpy as np -import pybase64 as base64 from fastapi import Request from vllm.engine.protocol import EngineClient @@ -39,6 +36,7 @@ ChatMessage, ) from vllm.entrypoints.openai.engine.protocol import ( + CompletionTokenUsageInfo, DeltaMessage, ErrorResponse, FunctionCall, @@ -64,6 +62,7 @@ from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.utils.collection_utils import as_list +from vllm.utils.serial_utils import numpy2base64 logger = init_logger(__name__) @@ -107,6 +106,12 @@ def _make_prompt_tokens_details( ) +def _make_completion_tokens_details( + reasoning_tokens: int, +) -> CompletionTokenUsageInfo: + return CompletionTokenUsageInfo(reasoning_tokens=reasoning_tokens) + + class OpenAIServingChat(GenerateBaseServing): def __init__( self, @@ -148,6 +153,7 @@ def __init__( self.enable_log_deltas = enable_log_deltas self.enable_auto_tools: bool = enable_auto_tools + self._include_reasoning_tokens_details = bool(reasoning_parser) self.parser_cls = ParserManager.get_parser( tool_parser_name=tool_parser, reasoning_parser_name=reasoning_parser, @@ -441,6 +447,9 @@ async def chat_completion_stream_generator( # Send response for each token for each request.n (index) num_choices = 1 if request.n is None else request.n previous_num_tokens = [0] * num_choices + # TODO: Remove once all reasoning parsers use the Parser Engine. + generated_token_ids: list[list[int]] = [[] for _ in range(num_choices)] + previous_reasoning_tokens = [0] * num_choices finish_reason_sent = [False] * num_choices num_prompt_tokens = 0 num_cached_tokens = None @@ -539,6 +548,11 @@ async def chat_completion_stream_generator( prompt_tokens=num_prompt_tokens, completion_tokens=0, total_tokens=num_prompt_tokens, + completion_tokens_details=( + _make_completion_tokens_details(0) + if self._include_reasoning_tokens_details + else None + ), ) data = chunk.model_dump_json(exclude_unset=True) @@ -575,6 +589,11 @@ async def chat_completion_stream_generator( prompt_tokens=num_prompt_tokens, completion_tokens=0, total_tokens=num_prompt_tokens, + completion_tokens_details=( + _make_completion_tokens_details(0) + if self._include_reasoning_tokens_details + else None + ), ) data = chunk.model_dump_json(exclude_unset=True) @@ -633,6 +652,11 @@ async def chat_completion_stream_generator( # set the previous values for the next iteration previous_num_tokens[i] += len(output.token_ids) + if parser is not None: + generated_token_ids[i].extend(output.token_ids) + previous_reasoning_tokens[i] = parser.count_reasoning_tokens( + tuple(generated_token_ids[i]) + ) # if the message delta is None (e.g. because it was a # "control token" for tool calls or the parser otherwise @@ -758,6 +782,13 @@ async def chat_completion_stream_generator( prompt_tokens=num_prompt_tokens, completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, + completion_tokens_details=( + _make_completion_tokens_details( + previous_reasoning_tokens[i] + ) + if self._include_reasoning_tokens_details + else None + ), ) data = chunk.model_dump_json(exclude_unset=True) @@ -771,6 +802,11 @@ async def chat_completion_stream_generator( prompt_tokens=num_prompt_tokens, completion_tokens=completion_tokens, total_tokens=num_prompt_tokens + completion_tokens, + completion_tokens_details=_make_completion_tokens_details( + sum(previous_reasoning_tokens) + ) + if self._include_reasoning_tokens_details + else None, ) final_usage.prompt_tokens_details = _make_prompt_tokens_details( self.enable_prompt_tokens_details, @@ -815,6 +851,11 @@ async def chat_completion_stream_generator( prompt_tokens=num_prompt_tokens, completion_tokens=num_completion_tokens, total_tokens=num_prompt_tokens + num_completion_tokens, + completion_tokens_details=_make_completion_tokens_details( + sum(previous_reasoning_tokens) + ) + if self._include_reasoning_tokens_details + else None, ) # Log complete streaming response if output logging is enabled @@ -873,6 +914,7 @@ async def chat_completion_full_generator( ) choices: list[ChatCompletionResponseChoice] = [] + total_reasoning_tokens = 0 role = self.get_chat_request_role(request) tool_parser_cls = ( @@ -910,6 +952,7 @@ async def chat_completion_full_generator( suppress_metadata = not request.include_reasoning and parser is not None if not request.include_reasoning: reasoning = None + total_reasoning_tokens += parser.count_reasoning_tokens(token_ids) if suppress_metadata: logprobs = None else: @@ -1007,15 +1050,11 @@ async def chat_completion_full_generator( and output.finish_reason == "stop" ) - # Encode routed_experts for transport. JSON can't carry raw - # bytes, so we write the ndarray as a ``.npy`` byte stream - # and base64-encode it. ``pybase64`` is ~3x faster than the - # stdlib ``base64`` on large payloads thanks to SIMD. - routed_experts_b64 = None - if output.routed_experts is not None: - buf = io.BytesIO() - np.save(buf, output.routed_experts) - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + routed_experts_b64 = ( + numpy2base64(output.routed_experts) + if output.routed_experts is not None + else None + ) choice_data = ChatCompletionResponseChoice( index=output.index, @@ -1064,6 +1103,11 @@ async def chat_completion_full_generator( prompt_tokens=num_prompt_tokens, completion_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, + completion_tokens_details=_make_completion_tokens_details( + total_reasoning_tokens + ) + if self._include_reasoning_tokens_details + else None, ) usage.prompt_tokens_details = _make_prompt_tokens_details( self.enable_prompt_tokens_details, diff --git a/vllm/entrypoints/openai/cli_args.py b/vllm/entrypoints/openai/cli_args.py index 64dab7416d3b..b387fb63573a 100644 --- a/vllm/entrypoints/openai/cli_args.py +++ b/vllm/entrypoints/openai/cli_args.py @@ -141,9 +141,9 @@ class BaseFrontendArgs: """Enable the `/tokenizer_info` endpoint. May expose chat templates and other tokenizer configuration.""" enable_log_outputs: bool = False - """If set to True, log model outputs (generations). - Requires `--enable-log-requests`. As with `--enable-log-requests`, - information is only logged at INFO level at maximum.""" + """If set to True, log model outputs (generations). Requires + `--enable-log-requests`. Output text and finish reasons are logged at INFO, + while output token IDs are logged at DEBUG.""" enable_log_deltas: bool = True """If set to False, output deltas will not be logged. Relevant only if --enable-log-outputs is set. diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 54de327c6481..0cb830294687 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -175,6 +175,11 @@ class CompletionRequest(OpenAIBaseModel): "need to map generated text back to input tokens." ), ) + routed_experts_prompt_start: int = Field( + default=0, + ge=0, + description="Skip the first N prompt tokens from returned routed-expert data.", + ) return_token_offsets: bool | None = Field( default=False, description=( @@ -394,6 +399,7 @@ def to_sampling_params( skip_clone=True, # Created fresh per request, safe to skip clone repetition_detection=self.repetition_detection, thinking_token_budget=self.thinking_token_budget, + routed_experts_prompt_start=self.routed_experts_prompt_start, ) @model_validator(mode="before") @@ -407,6 +413,8 @@ def normalize_null_max_tokens(cls, data): @model_validator(mode="before") @classmethod def validate_response_format(cls, data): + if not isinstance(data, dict): + return data response_format = data.get("response_format") if response_format is None: return data @@ -438,6 +446,8 @@ def validate_response_format(cls, data): @model_validator(mode="before") @classmethod def check_structured_outputs_count(cls, data): + if not isinstance(data, dict): + return data if data.get("structured_outputs", None) is None: return data @@ -466,6 +476,8 @@ def check_structured_outputs_count(cls, data): @model_validator(mode="before") @classmethod def check_logprobs(cls, data): + if not isinstance(data, dict): + return data if data.get("logprob_token_ids") and data.get("use_beam_search"): raise VLLMValidationError( "`logprob_token_ids` is not supported with beam search.", @@ -526,6 +538,8 @@ def check_logprobs(cls, data): @model_validator(mode="before") @classmethod def validate_stream_options(cls, data): + if not isinstance(data, dict): + return data if data.get("stream_options") and not data.get("stream"): raise VLLMValidationError( "Stream options can only be defined when `stream=True`.", @@ -537,6 +551,8 @@ def validate_stream_options(cls, data): @model_validator(mode="before") @classmethod def validate_prompt_and_prompt_embeds(cls, data): + if not isinstance(data, dict): + return data prompt = data.get("prompt") prompt_embeds = data.get("prompt_embeds") @@ -556,6 +572,8 @@ def validate_prompt_and_prompt_embeds(cls, data): @model_validator(mode="before") @classmethod def validate_prompt_list_length(cls, data): + if not isinstance(data, dict): + return data max_prompts = envs.VLLM_MAX_COMPLETION_PROMPTS prompt = data.get("prompt") @@ -611,7 +629,7 @@ class CompletionResponseChoice(OpenAIBaseModel): prompt_token_ids: list[int] | None = None # For prompt # Per-token expert routing decisions, base64-encoded ``.npy`` bytes # (numpy serialization). Shape after decode: - # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32 # ``num_tokens - 1`` because the last sampled token has not been # forwarded yet and therefore has no routing data. # Decode: diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 66b6016d9f15..37f55f4116c2 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -2,14 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import io import time from collections.abc import AsyncGenerator, AsyncIterator from collections.abc import Sequence as GenericSequence from typing import cast -import numpy as np -import pybase64 as base64 from fastapi import Request from vllm.engine.protocol import EngineClient @@ -48,6 +45,7 @@ from vllm.tokenizers import TokenizerLike from vllm.utils.async_utils import merge_async_iterators from vllm.utils.collection_utils import as_list +from vllm.utils.serial_utils import numpy2base64 logger = init_logger(__name__) @@ -570,17 +568,11 @@ def request_output_to_completion_response( else: logprobs = None - # Encode routed_experts for transport. JSON can't carry raw - # bytes, so we write the ndarray as a ``.npy`` byte stream - # and base64-encode it. ``pybase64`` is ~3x faster than the - # stdlib ``base64`` on large payloads thanks to SIMD. - routed_experts_b64 = None - if output.routed_experts is not None: - buf = io.BytesIO() - np.save(buf, output.routed_experts) - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode( - "ascii" - ) + routed_experts_b64 = ( + numpy2base64(output.routed_experts) + if output.routed_experts is not None + else None + ) choice_data = CompletionResponseChoice( index=len(choices), diff --git a/vllm/entrypoints/openai/engine/protocol.py b/vllm/entrypoints/openai/engine/protocol.py index 4d07c2940832..c32be4459c90 100644 --- a/vllm/entrypoints/openai/engine/protocol.py +++ b/vllm/entrypoints/openai/engine/protocol.py @@ -117,11 +117,16 @@ class PromptTokenUsageInfo(OpenAIBaseModel): request has no multimodal input.""" +class CompletionTokenUsageInfo(OpenAIBaseModel): + reasoning_tokens: int = 0 + + class UsageInfo(OpenAIBaseModel): prompt_tokens: int = 0 total_tokens: int = 0 completion_tokens: int | None = 0 prompt_tokens_details: PromptTokenUsageInfo | None = None + completion_tokens_details: CompletionTokenUsageInfo | None = None class PerRequestTimingMetrics(OpenAIBaseModel): @@ -261,7 +266,7 @@ def validate_structural_tag_payload(payload: Any, *, parameter: str) -> None: structured_outputs=StructuredOutputsParams(structural_tag=payload) ) ) - except (TypeError, ValueError) as exc: + except (TypeError, ValueError, VLLMValidationError) as exc: raise VLLMValidationError( f"Invalid {parameter} structural_tag specification.", parameter=parameter, diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py index b886d92641e7..9ceb8c7dfbb8 100644 --- a/vllm/entrypoints/openai/models/serving.py +++ b/vllm/entrypoints/openai/models/serving.py @@ -14,11 +14,11 @@ ModelPermission, ) from vllm.entrypoints.openai.models.protocol import BaseModelPath, LoRAModulePath +from vllm.entrypoints.serve import create_error_response from vllm.entrypoints.serve.lora.protocol import ( LoadLoRAAdapterRequest, UnloadLoRAAdapterRequest, ) -from vllm.entrypoints.serve.utils.error_response import create_error_response from vllm.exceptions import LoRAAdapterNotFoundError from vllm.logger import init_logger from vllm.lora.request import LoRARequest diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 7ad9f5711623..ff3331612f9f 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -467,6 +467,8 @@ def is_include_output_logprobs(self) -> bool: @model_validator(mode="before") @classmethod def validate_background(cls, data): + if not isinstance(data, dict): + return data if not data.get("background"): return data if not data.get("store", True): @@ -479,6 +481,8 @@ def validate_background(cls, data): @model_validator(mode="before") @classmethod def validate_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("prompt") is not None: raise VLLMValidationError( "prompt template is not supported", parameter="prompt" @@ -499,6 +503,8 @@ def input_item_parsing(cls, data): Invalid structures are left for Pydantic to reject. """ + if not isinstance(data, dict): + return data input_data = data.get("input") # Early return for None, strings, or bytes diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 9d8c5dafb9f1..20a331f42aa0 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -487,6 +487,7 @@ async def _create_responses( response_parser=response_parser, ) + reasoning_parser_kwargs = None if ( context.response_parser is not None and context.response_parser.reasoning_parser is not None @@ -518,9 +519,7 @@ async def _create_responses( priority=self._get_priority(request, raw_request), trace_headers=trace_headers, session_id=session_id, - reasoning_parser_kwargs=reasoning_parser_kwargs - if self.parser and self.parser.reasoning_parser_cls is not None - else None, + reasoning_parser_kwargs=reasoning_parser_kwargs, ) generators.append(generator) @@ -873,6 +872,8 @@ async def responses_full_generator( if final_output.finish_reason == "length": status = "incomplete" + # TODO: Build final response items from the accumulated streaming + # parser results instead of reparsing the complete output. output = self._make_response_output_items( request, final_output, @@ -901,13 +902,10 @@ async def responses_full_generator( num_reasoning_tokens == 0 and isinstance(context, (SimpleContext, ParsableContext)) and context.response_parser is not None - and context.response_parser.reasoning_parser is not None ): accumulated = getattr(context, "_accumulated_token_ids", []) or [] - num_reasoning_tokens = ( - context.response_parser.reasoning_parser.count_reasoning_tokens( - accumulated - ) + num_reasoning_tokens = context.response_parser.count_reasoning_tokens( + accumulated ) usage = ResponseUsage( diff --git a/vllm/entrypoints/openai/run_batch.py b/vllm/entrypoints/openai/run_batch.py index 6ae608da0adb..e9b7c582d6de 100644 --- a/vllm/entrypoints/openai/run_batch.py +++ b/vllm/entrypoints/openai/run_batch.py @@ -52,7 +52,7 @@ ScoreRequest, ScoreResponse, ) -from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve import create_error_response from vllm.entrypoints.speech_to_text.transcription.protocol import ( TranscriptionRequest, TranscriptionResponse, diff --git a/vllm/entrypoints/pooling/base/protocol.py b/vllm/entrypoints/pooling/base/protocol.py index fcc7c24b28c5..6e592e30f88c 100644 --- a/vllm/entrypoints/pooling/base/protocol.py +++ b/vllm/entrypoints/pooling/base/protocol.py @@ -239,6 +239,8 @@ class ChatRequestOptionsMixin(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_generation_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("continue_final_message") and data.get("add_generation_prompt"): raise VLLMValidationError( "Cannot set both `continue_final_message` and " diff --git a/vllm/entrypoints/pooling/embed/protocol.py b/vllm/entrypoints/pooling/embed/protocol.py index 96a87fea2018..d319131a1d13 100644 --- a/vllm/entrypoints/pooling/embed/protocol.py +++ b/vllm/entrypoints/pooling/embed/protocol.py @@ -12,6 +12,7 @@ from collections.abc import Sequence from typing import Annotated, Any, Literal, TypeAlias +import numpy as np import pybase64 as base64 from pydantic import BaseModel, BeforeValidator, Field, model_validator @@ -298,31 +299,28 @@ def _pack_binary_embeddings( ) -> list[list[int]]: """Bit-pack float embeddings: positive -> 1, negative -> 0. - Each bit is shifted left by ``7 - idx%8``, and every 8 bits are packed - into one byte. + Bits are packed MSB-first, eight per byte. """ - result: list[list[int]] = [] - for embedding in float_embeddings: - dim = len(embedding) - if dim % 8 != 0: - raise ValueError( - "Embedding dimension must be a multiple of 8 for binary " - f"embedding types, but got {dim}." - ) - packed_len = dim // 8 - packed: list[int] = [] - byte_val = 0 - for idx, value in enumerate(embedding): - bit = 1 if value >= 0 else 0 - byte_val += bit << (7 - idx % 8) - if (idx + 1) % 8 == 0: - if signed: - byte_val -= _UNSIGNED_TO_SIGNED_DIFF - packed.append(byte_val) - byte_val = 0 - assert len(packed) == packed_len - result.append(packed) - return result + if not float_embeddings: + return [] + + array = np.asarray(float_embeddings, dtype=np.float64) + if array.ndim != 2: + raise ValueError( + f"Expected a 2D batch of embeddings, but got {array.ndim}D input." + ) + + dim = array.shape[1] + if dim % 8 != 0: + raise ValueError( + "Embedding dimension must be a multiple of 8 for binary " + f"embedding types, but got {dim}." + ) + + packed = np.packbits(array >= 0, axis=-1) + if signed: + packed = packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF + return packed.tolist() def _encode_base64_embeddings( diff --git a/vllm/entrypoints/pooling/factories.py b/vllm/entrypoints/pooling/factories.py index dd3d873b3116..c7615054dd26 100644 --- a/vllm/entrypoints/pooling/factories.py +++ b/vllm/entrypoints/pooling/factories.py @@ -65,10 +65,16 @@ def init_pooling_io_processors( processors["token_embed"] = TokenEmbedIOProcessor - if has_io_processor( + if pooling_task == "embed&token_classify": + from .pooling.io_processor import UnsupportedCombinedTaskIOProcessor + + processors[pooling_task] = UnsupportedCombinedTaskIOProcessor + + has_plugin = has_io_processor( vllm_config, model_config.io_processor_plugin, - ): + ) + if has_plugin: from .pooling.io_processor import PluginWithIOProcessorPlugins processors["plugin"] = PluginWithIOProcessorPlugins diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index ecc356c477a8..4b18cfc9e184 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -24,6 +24,18 @@ logger = init_logger(__name__) +class UnsupportedCombinedTaskIOProcessor(PoolingIOProcessor): + name = "embed&token_classify" + + def create_pooling_params(self, request): + raise ValueError( + "The 'embed&token_classify' pooling task is only available " + "through an IO processor plugin. Send a plugin request with " + "a 'data' field, " + "or select a concrete task with --pooler-config.task." + ) + + class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): # Some models, such as Terratorch (tests/models/test_terratorch.py), # use plugin tasks in the pooler but do not use IO Processor plugins. diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 966a878528a0..f304bf677bab 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -202,7 +202,7 @@ class GenerateResponseChoice(BaseModel): token_ids: list[int] | None = None # Per-token expert routing decisions, base64-encoded ``.npy`` bytes # (numpy serialization). Shape after decode: - # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32 # ``num_tokens - 1`` because the last sampled token has not been # forwarded yet and therefore has no routing data. # Decode: diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 775187302522..9e9ace877a7c 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -3,14 +3,11 @@ import asyncio -import io import time from collections.abc import AsyncGenerator from collections.abc import Sequence as GenericSequence import msgspec -import numpy as np -import pybase64 as base64 from fastapi import Request from vllm.engine.protocol import EngineClient @@ -46,6 +43,7 @@ from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.utils.collection_utils import as_list +from vllm.utils.serial_utils import numpy2base64 from .mm_serde import decode_mm_kwargs_item from .protocol import ( @@ -303,17 +301,11 @@ async def serve_tokens_full_generator( else: logprobs = None - # Encode routed_experts for transport. JSON can't carry raw - # bytes, so we write the ndarray as a ``.npy`` byte stream - # and base64-encode it. ``pybase64`` is ~3x faster than the - # stdlib ``base64`` on large payloads thanks to SIMD. - # This is the only base64 hop in the pipeline -- the - # engine<->API-server link is binary msgpack + zmq. - routed_experts_b64 = None - if output.routed_experts is not None: - buf = io.BytesIO() - np.save(buf, output.routed_experts) - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + routed_experts_b64 = ( + numpy2base64(output.routed_experts) + if output.routed_experts is not None + else None + ) sampling_mask = None if output.sampling_mask is not None: @@ -435,13 +427,11 @@ async def serve_tokens_stream_generator( else: logprobs = None - routed_experts_b64 = None - if output.routed_experts is not None: - buf = io.BytesIO() - np.save(buf, output.routed_experts) - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode( - "ascii" - ) + routed_experts_b64 = ( + numpy2base64(output.routed_experts) + if output.routed_experts is not None + else None + ) chunk = GenerateStreamResponse( request_id=request_id, diff --git a/vllm/entrypoints/serve/__init__.py b/vllm/entrypoints/serve/__init__.py index 57491d45f631..c358b9a6db26 100644 --- a/vllm/entrypoints/serve/__init__.py +++ b/vllm/entrypoints/serve/__init__.py @@ -5,8 +5,16 @@ from vllm.logger import init_logger +from .exception_handling.error_response import create_error_response + logger = init_logger(__name__) +__all__ = [ + "create_error_response", + "register_vllm_serve_api_routers", + "register_vllm_dev_api_routers", +] + def register_vllm_serve_api_routers(app: FastAPI): from .instrumentator import register_instrumentator_api_routers diff --git a/vllm/entrypoints/serve/engine/serving.py b/vllm/entrypoints/serve/engine/serving.py index 5908e20f71c1..a31e47d9c55a 100644 --- a/vllm/entrypoints/serve/engine/serving.py +++ b/vllm/entrypoints/serve/engine/serving.py @@ -12,8 +12,8 @@ OpenAIServingModels, ) from vllm.entrypoints.pooling.typing import AnyPoolingRequest +from vllm.entrypoints.serve import create_error_response from vllm.entrypoints.serve.engine.typing import AnyRequest -from vllm.entrypoints.serve.utils.error_response import create_error_response from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.exceptions import VLLMNotFoundError from vllm.inputs import EngineInput diff --git a/vllm/entrypoints/serve/exception_handling/__init__.py b/vllm/entrypoints/serve/exception_handling/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/utils/error_response.py b/vllm/entrypoints/serve/exception_handling/error_response.py similarity index 97% rename from vllm/entrypoints/serve/utils/error_response.py rename to vllm/entrypoints/serve/exception_handling/error_response.py index 2aa785c53bb1..bb31c1d335cc 100644 --- a/vllm/entrypoints/serve/utils/error_response.py +++ b/vllm/entrypoints/serve/exception_handling/error_response.py @@ -7,9 +7,10 @@ ErrorResponse, GenerationError, ) -from vllm.entrypoints.serve.utils.api_utils import sanitize_message from vllm.logger import init_logger +from .utils import sanitize_message + logger = init_logger(__name__) diff --git a/vllm/entrypoints/serve/exception_handling/handlers/__init__.py b/vllm/entrypoints/serve/exception_handling/handlers/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/exception_handling/handlers/exception.py b/vllm/entrypoints/serve/exception_handling/handlers/exception.py new file mode 100644 index 000000000000..5afef771abea --- /dev/null +++ b/vllm/entrypoints/serve/exception_handling/handlers/exception.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from fastapi import Request +from starlette.responses import JSONResponse + +from vllm.logger import init_logger + +from ..error_response import create_error_response + +logger = init_logger(__name__) + + +async def exception_handler(req: Request, exc: Exception): + if req.app.state.args.log_error_stack: + logger.error( + "Exception caught. Request id: %s", + req.state.request_metadata.request_id + if hasattr(req.state, "request_metadata") + else None, + ) + + err = create_error_response(exc) + return JSONResponse(err.model_dump(), status_code=err.error.code) diff --git a/vllm/entrypoints/serve/exception_handling/handlers/http.py b/vllm/entrypoints/serve/exception_handling/handlers/http.py new file mode 100644 index 000000000000..d838c0ae1052 --- /dev/null +++ b/vllm/entrypoints/serve/exception_handling/handlers/http.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from http import HTTPStatus + +from fastapi import HTTPException, Request +from starlette.responses import JSONResponse + +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm.logger import init_logger + +from ..utils import sanitize_message + +logger = init_logger(__name__) + + +async def http_exception_handler(req: Request, exc: HTTPException): + if req.app.state.args.log_error_stack: + logger.exception( + "HTTPException caught. Request id: %s", + req.state.request_metadata.request_id + if hasattr(req.state, "request_metadata") + else None, + ) + err = ErrorResponse( + error=ErrorInfo( + message=sanitize_message(exc.detail), + type=HTTPStatus(exc.status_code).phrase, + code=exc.status_code, + ) + ) + return JSONResponse(err.model_dump(), status_code=exc.status_code) diff --git a/vllm/entrypoints/serve/exception_handling/handlers/validation.py b/vllm/entrypoints/serve/exception_handling/handlers/validation.py new file mode 100644 index 000000000000..fc02389ef60c --- /dev/null +++ b/vllm/entrypoints/serve/exception_handling/handlers/validation.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from http import HTTPStatus + +import regex as re +from fastapi import Request +from fastapi.exceptions import RequestValidationError +from starlette.responses import JSONResponse + +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm.exceptions import VLLMValidationError +from vllm.logger import init_logger + +from ..utils import sanitize_message + +logger = init_logger(__name__) + + +_BRACKETED_INTERNAL_RE = re.compile(r"[\[\]{}()]") + +# NOTE: this list is pydantic-core's internal schema-kind vocabulary, +# not a stable public API -- it can grow when pydantic-core adds new +# wrapper/validator kinds. To refresh it after a pydantic upgrade: +# 1. Fuzz the validation-error-prone endpoints (e.g. /tokenize, +# /v1/completions, /v1/chat/completions) with deliberately +# malformed values for union-typed and wrapped fields (e.g. `stop`, +# `prompt`), and inspect the raw `loc` tuples in the response. +# 2. Any *unbracketed* segment that isn't a real field name or list +# index is a new internal marker -- add it here. Bracketed/ +# parenthesized markers (e.g. "list[...]", "function-wrap[...]") +# are already caught structurally by _BRACKETED_INTERNAL_RE and +# don't need a list entry. +# 3. pydantic-core's source (the `error.rs`/schema-kind definitions +# in the pydantic-core Rust crate) is the canonical reference if +# you want to check before it shows up in a live fuzz run. +_INTERNAL_LOC_MARKERS = frozenset( + { + "function-wrap", + "function-after", + "function-before", + "function-plain", + "json-or-python", + "lax-or-strict", + "chain", + "default", + "nullable", + "tagged-union", + "union", + "call", + "arguments", + "is-instance", + "is-subclass", + "callable", + "str", + "int", + "float", + "bool", + "bytes", + "bytearray", + "list", + "tuple", + "dict", + "set", + "frozenset", + "complex", + "none", + "nonetype", + } +) + + +def _is_internal_loc_segment(segment: str) -> bool: + """True if `segment` is a Pydantic-internal wrapper/union-branch + marker rather than a user-meaningful field name or list index.""" + if _BRACKETED_INTERNAL_RE.search(segment): + return True + return segment.lower() in _INTERNAL_LOC_MARKERS + + +def clean_loc_for_param(loc: tuple) -> str: + """Join a Pydantic error `loc` tuple into a clean dotted `param` + path, dropping internal wrapper/union-branch markers that don't + correspond to a real field name an API consumer would recognize. + + E.g. ('body', 'function-wrap[__log_extra_fields__()]', 'prompt') + -> "body.prompt", not "body.function-wrap[__log_extra_fields__()].prompt". + """ + parts = [str(p) for p in loc if not _is_internal_loc_segment(str(p))] + if not parts: + return ".".join(str(p) for p in loc) + return ".".join(parts) + + +async def validation_exception_handler(req: Request, exc: RequestValidationError): + if req.app.state.args.log_error_stack: + logger.exception( + "RequestValidationError caught. Request id: %s", + req.state.request_metadata.request_id + if hasattr(req.state, "request_metadata") + else None, + ) + + param = None + errors = exc.errors() + for error in errors: + if "ctx" in error and "error" in error["ctx"]: + ctx_error = error["ctx"]["error"] + if isinstance(ctx_error, VLLMValidationError): + param = ctx_error.parameter + break + + if param is None and errors: + first_error = errors[0] + loc = first_error.get("loc") if isinstance(first_error, dict) else None + if loc: + param = clean_loc_for_param(loc) + + # Build the message from exc.errors() instead of str(exc) - str(exc) + # leaks the server's file path via FastAPI's endpoint context. + if errors: + count = len(errors) + label = "error" if count == 1 else "errors" + message = f"{count} validation {label}:\n" + message += "".join(f" {err}\n" for err in errors) + message = message.rstrip() + else: + message = "Validation error" + + err = ErrorResponse( + error=ErrorInfo( + message=sanitize_message(message), + type=HTTPStatus.BAD_REQUEST.phrase, + code=HTTPStatus.BAD_REQUEST, + param=param, + ) + ) + return JSONResponse(err.model_dump(), status_code=HTTPStatus.BAD_REQUEST) diff --git a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py new file mode 100644 index 000000000000..8ac9756f3d7f --- /dev/null +++ b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from fastapi import Request +from starlette.responses import JSONResponse + +from vllm.entrypoints.launcher import terminate_if_errored +from vllm.entrypoints.openai.engine.protocol import GenerationError +from vllm.exceptions import VLLMError +from vllm.logger import init_logger +from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError + +from ..error_response import create_error_response +from .exception import exception_handler + +logger = init_logger(__name__) + + +async def vllm_error_handler(req: Request, exc: VLLMError): + """Dispatch a vLLM-specific error to the appropriate handler.""" + if isinstance(exc, (EngineGenerateError, EngineDeadError)): + return await engine_error_handler(req, exc) + elif isinstance(exc, GenerationError): + return await generation_error_handler(req, exc) + else: + return await exception_handler(req, exc) + + +async def engine_error_handler( + req: Request, exc: EngineDeadError | EngineGenerateError +): + """ + VLLM V1 AsyncLLM catches exceptions and returns + only two types: EngineGenerateError and EngineDeadError. + + EngineGenerateError is raised by the per request generate() + method. This error could be request specific (and therefore + recoverable - e.g. if there is an error in input processing). + + EngineDeadError is raised by the background output_handler + method. This error is global and therefore not recoverable. + + We register these @app.exception_handlers to return nice + responses to the end user if they occur and shut down if needed. + See https://fastapi.tiangolo.com/tutorial/handling-errors/ + for more details on how exception handlers work. + + If an exception is encountered in a StreamingResponse + generator, the exception is not raised, since we already sent + a 200 status. Rather, we send an error message as the next chunk. + Since the exception is not raised, this means that the server + will not automatically shut down. Instead, we use the watchdog + background task for check for errored state. + """ + + if req.app.state.args.log_error_stack: + logger.exception( + "Engine Exception caught. Request id: %s", + req.state.request_metadata.request_id + if hasattr(req.state, "request_metadata") + else None, + ) + + terminate_if_errored( + server=req.app.state.server, + engine=req.app.state.engine_client, + ) + err = create_error_response(exc) + return JSONResponse(err.model_dump(), status_code=err.error.code) + + +async def generation_error_handler(req: Request, exc: GenerationError): + """Handle GenerationError without logging stack traces. + + GenerationError is a known, expected error (e.g. KV cache load failure) + that should be returned to the client as a 500 response without polluting + server logs with stack traces. + """ + err = create_error_response(exc) + return JSONResponse(err.model_dump(), status_code=err.error.code) diff --git a/vllm/entrypoints/serve/exception_handling/register.py b/vllm/entrypoints/serve/exception_handling/register.py new file mode 100644 index 000000000000..9f94873939a0 --- /dev/null +++ b/vllm/entrypoints/serve/exception_handling/register.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +vLLM Exception handlers are registered in four layers: +1. framework errors raised by FastAPI/Starlette +2. vLLM-specific errors dispatched via a single ``VLLMError`` handler +3. fallback handlers for raw exceptions not yet migrated to ``VLLMError`` +4. the raw ``Exception`` handler as a safety net +Registering specific exception types (rather than only ``Exception``) +ensures they are handled by ``ExceptionMiddleware`` (inside the Prometheus +middleware) rather than ``ServerErrorMiddleware`` (outside it), so their +status codes are recorded correctly. +""" + +from fastapi import FastAPI, HTTPException +from fastapi.exceptions import RequestValidationError + +from vllm.exceptions import VLLMError + +from .handlers.exception import exception_handler +from .handlers.http import http_exception_handler +from .handlers.validation import validation_exception_handler +from .handlers.vllm_error import vllm_error_handler + + +def init_exception_handler(app: FastAPI): + # 1. framework errors raised by FastAPI/Starlette + app.exception_handler(HTTPException)(http_exception_handler) + app.exception_handler(RequestValidationError)(validation_exception_handler) + + # 2. vLLM-specific errors dispatched via a single ``VLLMError`` handler + app.exception_handler(VLLMError)(vllm_error_handler) + + # 3. fallback handlers for raw exceptions not yet migrated to ``VLLMError`` + # TODO(zqzten): remove these fallback handlers after migration to VLLMError + app.exception_handler(ValueError)(exception_handler) + app.exception_handler(TypeError)(exception_handler) + app.exception_handler(OverflowError)(exception_handler) + app.exception_handler(NotImplementedError)(exception_handler) + + # 4. the raw ``Exception`` handler as a safety net + app.exception_handler(Exception)(exception_handler) diff --git a/vllm/entrypoints/serve/exception_handling/utils.py b/vllm/entrypoints/serve/exception_handling/utils.py new file mode 100644 index 000000000000..73d9992237d7 --- /dev/null +++ b/vllm/entrypoints/serve/exception_handling/utils.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import regex as re + + +def sanitize_message(message: str) -> str: + """Strip memory addresses, tracebacks, and file paths from error messages.""" + message = re.sub(r" at 0x[0-9a-f]+>", ">", message) + message = re.sub(r'\n?\s*File "[^"]+", line \d+, in \S+(\n\s+.*)?', "", message) + message = re.sub( + r"/(?:home|usr|opt|var|tmp|root|lib|mnt|srv)(?:/[\w.\-]+)+", "", message + ) + message = re.sub(r"(?:/[\w\-]+)+/[\w\-]+\.\w+", "", message) + return message.strip() diff --git a/vllm/entrypoints/serve/middleware/__init__.py b/vllm/entrypoints/serve/middleware/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm/entrypoints/serve/middleware/authenticate.py b/vllm/entrypoints/serve/middleware/authenticate.py new file mode 100644 index 000000000000..e8a61b21ee46 --- /dev/null +++ b/vllm/entrypoints/serve/middleware/authenticate.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import hashlib +import secrets +from collections.abc import Awaitable + +from starlette.datastructures import Headers +from starlette.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send + +GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere") + + +class AuthenticationMiddleware: + """ + Pure ASGI middleware that authenticates each request by checking + if the Authorization Bearer token exists and equals anyof "{api_key}". + + Notes + ----- + There are two cases in which authentication is skipped: + 1. The HTTP method is OPTIONS. + 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health). + """ + + def __init__(self, app: ASGIApp, tokens: list[str]) -> None: + self.app = app + self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens] + + def verify_token(self, headers: Headers) -> bool: + authorization_header_value = headers.get("Authorization") + if not authorization_header_value: + return False + + scheme, _, param = authorization_header_value.partition(" ") + if scheme.lower() != "bearer": + return False + + param_hash = hashlib.sha256(param.encode("utf-8")).digest() + + token_match = False + for token_hash in self.api_tokens: + token_match |= secrets.compare_digest(param_hash, token_hash) + + return token_match + + def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: + if ( + scope["type"] not in ("http", "websocket") + or scope.get("method") == "OPTIONS" + ): + # scope["type"] can be "lifespan" or "startup" for example, + # in which case we don't need to do anything + return self.app(scope, receive, send) + root_path = scope.get("root_path", "") + url_path = scope["path"].removeprefix(root_path) + headers = Headers(scope=scope) + # Type narrow to satisfy mypy. + if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers): + response = JSONResponse(content={"error": "Unauthorized"}, status_code=401) + return response(scope, receive, send) + return self.app(scope, receive, send) diff --git a/vllm/entrypoints/serve/middleware/log_response.py b/vllm/entrypoints/serve/middleware/log_response.py new file mode 100644 index 000000000000..e2f5e5bde51e --- /dev/null +++ b/vllm/entrypoints/serve/middleware/log_response.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pydantic +from fastapi import Request +from starlette.concurrency import iterate_in_threadpool + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class SSEDecoder: + """Robust Server-Sent Events decoder for streaming responses.""" + + def __init__(self): + self.buffer = "" + self.content_buffer = [] + + def decode_chunk(self, chunk: bytes) -> list[dict]: + """Decode a chunk of SSE data and return parsed events.""" + import json + + try: + chunk_str = chunk.decode("utf-8") + except UnicodeDecodeError: + # Skip malformed chunks + return [] + + self.buffer += chunk_str + events = [] + + # Process complete lines + while "\n" in self.buffer: + line, self.buffer = self.buffer.split("\n", 1) + line = line.rstrip("\r") # Handle CRLF + + if line.startswith("data: "): + data_str = line[6:].strip() + if data_str == "[DONE]": + events.append({"type": "done"}) + elif data_str: + try: + event_data = json.loads(data_str) + events.append({"type": "data", "data": event_data}) + except json.JSONDecodeError: + # Skip malformed JSON + continue + + return events + + def extract_content(self, event_data: dict) -> str: + """Extract content from event data.""" + return _extract_content_from_chunk(event_data) + + def add_content(self, content: str) -> None: + """Add content to the buffer.""" + if content: + self.content_buffer.append(content) + + def get_complete_content(self) -> str: + """Get the complete buffered content.""" + return "".join(self.content_buffer) + + +def _extract_content_from_chunk(chunk_data: dict) -> str: + """Extract content from a streaming response chunk.""" + try: + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionStreamResponse, + ) + from vllm.entrypoints.openai.completion.protocol import ( + CompletionStreamResponse, + ) + + # Try using Completion types for type-safe parsing + if chunk_data.get("object") == "chat.completion.chunk": + chat_response = ChatCompletionStreamResponse.model_validate(chunk_data) + if chat_response.choices and chat_response.choices[0].delta.content: + return chat_response.choices[0].delta.content + elif chunk_data.get("object") == "text_completion": + completion_response = CompletionStreamResponse.model_validate(chunk_data) + if completion_response.choices and completion_response.choices[0].text: + return completion_response.choices[0].text + except pydantic.ValidationError: + # Fallback to manual parsing + if "choices" in chunk_data and chunk_data["choices"]: + choice = chunk_data["choices"][0] + if "delta" in choice and choice["delta"].get("content"): + return choice["delta"]["content"] + elif choice.get("text"): + return choice["text"] + return "" + + +def _log_streaming_response(response, response_body: list) -> None: + """Log streaming response with robust SSE parsing.""" + + sse_decoder = SSEDecoder() + chunk_count = 0 + + def buffered_iterator(): + nonlocal chunk_count + + for chunk in response_body: + chunk_count += 1 + yield chunk + + # Parse SSE events from chunk + events = sse_decoder.decode_chunk(chunk) + + for event in events: + if event["type"] == "data": + content = sse_decoder.extract_content(event["data"]) + sse_decoder.add_content(content) + elif event["type"] == "done": + # Log complete content when done + full_content = sse_decoder.get_complete_content() + if full_content: + # Truncate if too long + if len(full_content) > 2048: + full_content = full_content[:2048] + "...[truncated]" + logger.info( + "response_body={streaming_complete: content=%r, chunks=%d}", + full_content, + chunk_count, + ) + else: + logger.info( + "response_body={streaming_complete: no_content, chunks=%d}", + chunk_count, + ) + return + + response.body_iterator = iterate_in_threadpool(buffered_iterator()) + logger.info("response_body={streaming_started: chunks=%d}", len(response_body)) + + +def _log_non_streaming_response(response_body: list) -> None: + """Log non-streaming response.""" + try: + decoded_body = response_body[0].decode() + logger.info("response_body={%s}", decoded_body) + except UnicodeDecodeError: + logger.info("response_body={}") + + +async def log_response(request: Request, call_next): + response = await call_next(request) + response_body = [section async for section in response.body_iterator] + response.body_iterator = iterate_in_threadpool(iter(response_body)) + # Check if this is a streaming response by looking at content-type + content_type = response.headers.get("content-type", "") + is_streaming = content_type == "text/event-stream; charset=utf-8" + + # Log response body based on type + if not response_body: + logger.info("response_body={}") + elif is_streaming: + _log_streaming_response(response, response_body) + else: + _log_non_streaming_response(response_body) + return response diff --git a/vllm/entrypoints/serve/middleware/register.py b/vllm/entrypoints/serve/middleware/register.py new file mode 100644 index 000000000000..3b132ce0886d --- /dev/null +++ b/vllm/entrypoints/serve/middleware/register.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib +import inspect +from argparse import Namespace + +from fastapi import FastAPI +from starlette.middleware.cors import CORSMiddleware + +from vllm import envs +from vllm.logger import init_logger +from vllm.tasks import SupportedTask + +from .log_response import log_response + +logger = init_logger(__name__) + + +def init_entrypoints_middleware( + args: Namespace, + app: FastAPI, + supported_tasks: tuple["SupportedTask", ...], +): + app.add_middleware( + CORSMiddleware, + allow_origins=args.allowed_origins, + allow_credentials=args.allow_credentials, + allow_methods=args.allowed_methods, + allow_headers=args.allowed_headers, + ) + + # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY + if tokens := [key for key in (args.api_key or [envs.VLLM_API_KEY]) if key]: + from .authenticate import AuthenticationMiddleware + + app.add_middleware(AuthenticationMiddleware, tokens=tokens) + + if args.enable_request_id_headers: + from .x_request_id import XRequestIdMiddleware + + app.add_middleware(XRequestIdMiddleware) + + if "generate" in supported_tasks: + # Add scaling middleware to check for scaling state + from vllm.entrypoints.serve.elastic_ep.middleware import ScalingMiddleware + + app.add_middleware(ScalingMiddleware) + + if "realtime" in supported_tasks: + # Add WebSocket metrics middleware + from vllm.entrypoints.speech_to_text.realtime.metrics import ( + WebSocketMetricsMiddleware, + ) + + app.add_middleware(WebSocketMetricsMiddleware) + + if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE: + logger.warning( + "CAUTION: Enabling log response in the API Server. " + "This can include sensitive information and should be " + "avoided in production." + ) + app.middleware("http")(log_response) + + for middleware in args.middleware: + module_path, object_name = middleware.rsplit(".", 1) + imported = getattr(importlib.import_module(module_path), object_name) + if inspect.isclass(imported): + app.add_middleware(imported) # type: ignore[arg-type] + elif inspect.iscoroutinefunction(imported): + app.middleware("http")(imported) + else: + raise ValueError( + f"Invalid middleware {middleware}. Must be a function or a class." + ) diff --git a/vllm/entrypoints/serve/middleware/x_request_id.py b/vllm/entrypoints/serve/middleware/x_request_id.py new file mode 100644 index 000000000000..02e85306cd12 --- /dev/null +++ b/vllm/entrypoints/serve/middleware/x_request_id.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import uuid +from collections.abc import Awaitable + +from starlette.datastructures import Headers, MutableHeaders +from starlette.types import ASGIApp, Message, Receive, Scope, Send + + +class XRequestIdMiddleware: + """ + Middleware the set's the X-Request-Id header for each response + to a random uuid4 (hex) value if the header isn't already + present in the request, otherwise use the provided request id. + """ + + def __init__(self, app: ASGIApp) -> None: + self.app = app + + def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: + if scope["type"] not in ("http", "websocket"): + return self.app(scope, receive, send) + + # Extract the request headers. + request_headers = Headers(scope=scope) + + async def send_with_request_id(message: Message) -> None: + """ + Custom send function to mutate the response headers + and append X-Request-Id to it. + """ + if message["type"] == "http.response.start": + response_headers = MutableHeaders(raw=message["headers"]) + request_id = request_headers.get("X-Request-Id", uuid.uuid4().hex) + response_headers.append("X-Request-Id", request_id) + await send(message) + + return self.app(scope, receive, send_with_request_id) diff --git a/vllm/entrypoints/serve/tokenize/protocol.py b/vllm/entrypoints/serve/tokenize/protocol.py index 66c122da87de..e68302d1f95b 100644 --- a/vllm/entrypoints/serve/tokenize/protocol.py +++ b/vllm/entrypoints/serve/tokenize/protocol.py @@ -120,6 +120,8 @@ class TokenizeChatRequest(OpenAIBaseModel): @model_validator(mode="before") @classmethod def check_generation_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("continue_final_message") and data.get("add_generation_prompt"): raise VLLMValidationError( "Cannot set both `continue_final_message` and " diff --git a/vllm/entrypoints/serve/utils/api_utils.py b/vllm/entrypoints/serve/utils/api_utils.py index 5fc855d9b5a9..c75cd16e54f4 100644 --- a/vllm/entrypoints/serve/utils/api_utils.py +++ b/vllm/entrypoints/serve/utils/api_utils.py @@ -10,7 +10,6 @@ from string import Template from typing import Any -import regex as re from fastapi import Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse, StreamingResponse @@ -309,17 +308,6 @@ def process_lora_modules( return lora_modules -def sanitize_message(message: str) -> str: - """Strip memory addresses, tracebacks, and file paths from error messages.""" - message = re.sub(r" at 0x[0-9a-f]+>", ">", message) - message = re.sub(r'\n?\s*File "[^"]+", line \d+, in \S+(\n\s+.*)?', "", message) - message = re.sub( - r"/(?:home|usr|opt|var|tmp|root|lib|mnt|srv)(?:/[\w.\-]+)+", "", message - ) - message = re.sub(r"(?:/[\w\-]+)+/[\w\-]+\.\w+", "", message) - return message.strip() - - def log_version_and_model(lgr: Logger, version: str, model_name: str) -> None: if envs.VLLM_DISABLE_LOG_LOGO or (formatter := current_formatter_type(lgr)) is None: message = "vLLM server version %s, serving model %s" diff --git a/vllm/entrypoints/serve/utils/request_logger.py b/vllm/entrypoints/serve/utils/request_logger.py index c2a77fbb4e56..ac10feab1011 100644 --- a/vllm/entrypoints/serve/utils/request_logger.py +++ b/vllm/entrypoints/serve/utils/request_logger.py @@ -77,24 +77,28 @@ def log_outputs( delta: bool = False, ) -> None: max_log_len = self.max_log_len - if max_log_len is not None: - if outputs is not None: - outputs = outputs[:max_log_len] - - if output_token_ids is not None: - # Convert to list and apply truncation - output_token_ids = list(output_token_ids)[:max_log_len] + if max_log_len is not None and outputs is not None: + outputs = outputs[:max_log_len] stream_info = "" if is_streaming: stream_info = " (streaming delta)" if delta else " (streaming complete)" + if logger.isEnabledFor(logging.DEBUG): + if max_log_len is not None and output_token_ids is not None: + output_token_ids = list(output_token_ids)[:max_log_len] + + logger.debug( + "Generated response %s%s details: output_token_ids: %s", + request_id, + stream_info, + output_token_ids, + ) + logger.info( - "Generated response %s%s: output: %r, " - "output_token_ids: %s, finish_reason: %s", + "Generated response %s%s: output: %r, finish_reason: %s", request_id, stream_info, outputs, - output_token_ids, finish_reason, ) diff --git a/vllm/entrypoints/serve/utils/server_utils.py b/vllm/entrypoints/serve/utils/server_utils.py index 03e02e34f743..97910d019ec7 100644 --- a/vllm/entrypoints/serve/utils/server_utils.py +++ b/vllm/entrypoints/serve/utils/server_utils.py @@ -1,129 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio -import hashlib import json -import secrets -import uuid from argparse import Namespace -from collections.abc import Awaitable from contextlib import asynccontextmanager -from http import HTTPStatus -import pydantic -import regex as re -from fastapi import FastAPI, HTTPException, Request -from fastapi.exceptions import RequestValidationError -from fastapi.responses import JSONResponse -from starlette.concurrency import iterate_in_threadpool -from starlette.datastructures import Headers, MutableHeaders -from starlette.types import ASGIApp, Message, Receive, Scope, Send +from fastapi import FastAPI from vllm import envs from vllm.engine.protocol import EngineClient -from vllm.entrypoints.launcher import terminate_if_errored -from vllm.entrypoints.openai.engine.protocol import ( - ErrorInfo, - ErrorResponse, - GenerationError, -) -from vllm.entrypoints.serve.utils.error_response import ( - create_error_response, - sanitize_message, -) -from vllm.exceptions import VLLMError, VLLMValidationError from vllm.logger import init_logger from vllm.utils.gc_utils import freeze_gc_heap -from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError logger = init_logger("vllm.entrypoints.openai.server_utils") -GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere") - - -class AuthenticationMiddleware: - """ - Pure ASGI middleware that authenticates each request by checking - if the Authorization Bearer token exists and equals anyof "{api_key}". - - Notes - ----- - There are two cases in which authentication is skipped: - 1. The HTTP method is OPTIONS. - 2. The request path doesn't start with GUARDED_PREFIX (e.g. /health). - """ - - def __init__(self, app: ASGIApp, tokens: list[str]) -> None: - self.app = app - self.api_tokens = [hashlib.sha256(t.encode("utf-8")).digest() for t in tokens] - - def verify_token(self, headers: Headers) -> bool: - authorization_header_value = headers.get("Authorization") - if not authorization_header_value: - return False - - scheme, _, param = authorization_header_value.partition(" ") - if scheme.lower() != "bearer": - return False - - param_hash = hashlib.sha256(param.encode("utf-8")).digest() - - token_match = False - for token_hash in self.api_tokens: - token_match |= secrets.compare_digest(param_hash, token_hash) - - return token_match - - def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: - if ( - scope["type"] not in ("http", "websocket") - or scope.get("method") == "OPTIONS" - ): - # scope["type"] can be "lifespan" or "startup" for example, - # in which case we don't need to do anything - return self.app(scope, receive, send) - root_path = scope.get("root_path", "") - url_path = scope["path"].removeprefix(root_path) - headers = Headers(scope=scope) - # Type narrow to satisfy mypy. - if url_path.startswith(GUARDED_PREFIX) and not self.verify_token(headers): - response = JSONResponse(content={"error": "Unauthorized"}, status_code=401) - return response(scope, receive, send) - return self.app(scope, receive, send) - - -class XRequestIdMiddleware: - """ - Middleware the set's the X-Request-Id header for each response - to a random uuid4 (hex) value if the header isn't already - present in the request, otherwise use the provided request id. - """ - - def __init__(self, app: ASGIApp) -> None: - self.app = app - - def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: - if scope["type"] not in ("http", "websocket"): - return self.app(scope, receive, send) - - # Extract the request headers. - request_headers = Headers(scope=scope) - - async def send_with_request_id(message: Message) -> None: - """ - Custom send function to mutate the response headers - and append X-Request-Id to it. - """ - if message["type"] == "http.response.start": - response_headers = MutableHeaders(raw=message["headers"]) - request_id = request_headers.get("X-Request-Id", uuid.uuid4().hex) - response_headers.append("X-Request-Id", request_id) - await send(message) - - return self.app(scope, receive, send_with_request_id) - - def load_log_config(log_config_file: str | None) -> dict | None: if not log_config_file: return None @@ -170,377 +61,6 @@ def get_uvicorn_log_config(args: Namespace) -> dict | None: return None -def _extract_content_from_chunk(chunk_data: dict) -> str: - """Extract content from a streaming response chunk.""" - try: - from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionStreamResponse, - ) - from vllm.entrypoints.openai.completion.protocol import ( - CompletionStreamResponse, - ) - - # Try using Completion types for type-safe parsing - if chunk_data.get("object") == "chat.completion.chunk": - chat_response = ChatCompletionStreamResponse.model_validate(chunk_data) - if chat_response.choices and chat_response.choices[0].delta.content: - return chat_response.choices[0].delta.content - elif chunk_data.get("object") == "text_completion": - completion_response = CompletionStreamResponse.model_validate(chunk_data) - if completion_response.choices and completion_response.choices[0].text: - return completion_response.choices[0].text - except pydantic.ValidationError: - # Fallback to manual parsing - if "choices" in chunk_data and chunk_data["choices"]: - choice = chunk_data["choices"][0] - if "delta" in choice and choice["delta"].get("content"): - return choice["delta"]["content"] - elif choice.get("text"): - return choice["text"] - return "" - - -class SSEDecoder: - """Robust Server-Sent Events decoder for streaming responses.""" - - def __init__(self): - self.buffer = "" - self.content_buffer = [] - - def decode_chunk(self, chunk: bytes) -> list[dict]: - """Decode a chunk of SSE data and return parsed events.""" - import json - - try: - chunk_str = chunk.decode("utf-8") - except UnicodeDecodeError: - # Skip malformed chunks - return [] - - self.buffer += chunk_str - events = [] - - # Process complete lines - while "\n" in self.buffer: - line, self.buffer = self.buffer.split("\n", 1) - line = line.rstrip("\r") # Handle CRLF - - if line.startswith("data: "): - data_str = line[6:].strip() - if data_str == "[DONE]": - events.append({"type": "done"}) - elif data_str: - try: - event_data = json.loads(data_str) - events.append({"type": "data", "data": event_data}) - except json.JSONDecodeError: - # Skip malformed JSON - continue - - return events - - def extract_content(self, event_data: dict) -> str: - """Extract content from event data.""" - return _extract_content_from_chunk(event_data) - - def add_content(self, content: str) -> None: - """Add content to the buffer.""" - if content: - self.content_buffer.append(content) - - def get_complete_content(self) -> str: - """Get the complete buffered content.""" - return "".join(self.content_buffer) - - -def _log_streaming_response(response, response_body: list) -> None: - """Log streaming response with robust SSE parsing.""" - from starlette.concurrency import iterate_in_threadpool - - sse_decoder = SSEDecoder() - chunk_count = 0 - - def buffered_iterator(): - nonlocal chunk_count - - for chunk in response_body: - chunk_count += 1 - yield chunk - - # Parse SSE events from chunk - events = sse_decoder.decode_chunk(chunk) - - for event in events: - if event["type"] == "data": - content = sse_decoder.extract_content(event["data"]) - sse_decoder.add_content(content) - elif event["type"] == "done": - # Log complete content when done - full_content = sse_decoder.get_complete_content() - if full_content: - # Truncate if too long - if len(full_content) > 2048: - full_content = full_content[:2048] + "" - "...[truncated]" - logger.info( - "response_body={streaming_complete: content=%r, chunks=%d}", - full_content, - chunk_count, - ) - else: - logger.info( - "response_body={streaming_complete: no_content, chunks=%d}", - chunk_count, - ) - return - - response.body_iterator = iterate_in_threadpool(buffered_iterator()) - logger.info("response_body={streaming_started: chunks=%d}", len(response_body)) - - -def _log_non_streaming_response(response_body: list) -> None: - """Log non-streaming response.""" - try: - decoded_body = response_body[0].decode() - logger.info("response_body={%s}", decoded_body) - except UnicodeDecodeError: - logger.info("response_body={}") - - -async def log_response(request: Request, call_next): - response = await call_next(request) - response_body = [section async for section in response.body_iterator] - response.body_iterator = iterate_in_threadpool(iter(response_body)) - # Check if this is a streaming response by looking at content-type - content_type = response.headers.get("content-type", "") - is_streaming = content_type == "text/event-stream; charset=utf-8" - - # Log response body based on type - if not response_body: - logger.info("response_body={}") - elif is_streaming: - _log_streaming_response(response, response_body) - else: - _log_non_streaming_response(response_body) - return response - - -async def vllm_error_handler(req: Request, exc: VLLMError): - """Dispatch a vLLM-specific error to the appropriate handler.""" - if isinstance(exc, (EngineGenerateError, EngineDeadError)): - return await engine_error_handler(req, exc) - elif isinstance(exc, GenerationError): - return await generation_error_handler(req, exc) - else: - return await exception_handler(req, exc) - - -async def engine_error_handler( - req: Request, exc: EngineDeadError | EngineGenerateError -): - """ - VLLM V1 AsyncLLM catches exceptions and returns - only two types: EngineGenerateError and EngineDeadError. - - EngineGenerateError is raised by the per request generate() - method. This error could be request specific (and therefore - recoverable - e.g. if there is an error in input processing). - - EngineDeadError is raised by the background output_handler - method. This error is global and therefore not recoverable. - - We register these @app.exception_handlers to return nice - responses to the end user if they occur and shut down if needed. - See https://fastapi.tiangolo.com/tutorial/handling-errors/ - for more details on how exception handlers work. - - If an exception is encountered in a StreamingResponse - generator, the exception is not raised, since we already sent - a 200 status. Rather, we send an error message as the next chunk. - Since the exception is not raised, this means that the server - will not automatically shut down. Instead, we use the watchdog - background task for check for errored state. - """ - - if req.app.state.args.log_error_stack: - logger.exception( - "Engine Exception caught. Request id: %s", - req.state.request_metadata.request_id - if hasattr(req.state, "request_metadata") - else None, - ) - - terminate_if_errored( - server=req.app.state.server, - engine=req.app.state.engine_client, - ) - err = create_error_response(exc) - return JSONResponse(err.model_dump(), status_code=err.error.code) - - -async def generation_error_handler(req: Request, exc: GenerationError): - """Handle GenerationError without logging stack traces. - - GenerationError is a known, expected error (e.g. KV cache load failure) - that should be returned to the client as a 500 response without polluting - server logs with stack traces. - """ - err = create_error_response(exc) - return JSONResponse(err.model_dump(), status_code=err.error.code) - - -async def exception_handler(req: Request, exc: Exception): - if req.app.state.args.log_error_stack: - logger.error( - "Exception caught. Request id: %s", - req.state.request_metadata.request_id - if hasattr(req.state, "request_metadata") - else None, - ) - - err = create_error_response(exc) - return JSONResponse(err.model_dump(), status_code=err.error.code) - - -async def http_exception_handler(req: Request, exc: HTTPException): - if req.app.state.args.log_error_stack: - logger.exception( - "HTTPException caught. Request id: %s", - req.state.request_metadata.request_id - if hasattr(req.state, "request_metadata") - else None, - ) - err = ErrorResponse( - error=ErrorInfo( - message=sanitize_message(exc.detail), - type=HTTPStatus(exc.status_code).phrase, - code=exc.status_code, - ) - ) - return JSONResponse(err.model_dump(), status_code=exc.status_code) - - -_BRACKETED_INTERNAL_RE = re.compile(r"[\[\]{}()]") - -# NOTE: this list is pydantic-core's internal schema-kind vocabulary, -# not a stable public API -- it can grow when pydantic-core adds new -# wrapper/validator kinds. To refresh it after a pydantic upgrade: -# 1. Fuzz the validation-error-prone endpoints (e.g. /tokenize, -# /v1/completions, /v1/chat/completions) with deliberately -# malformed values for union-typed and wrapped fields (e.g. `stop`, -# `prompt`), and inspect the raw `loc` tuples in the response. -# 2. Any *unbracketed* segment that isn't a real field name or list -# index is a new internal marker -- add it here. Bracketed/ -# parenthesized markers (e.g. "list[...]", "function-wrap[...]") -# are already caught structurally by _BRACKETED_INTERNAL_RE and -# don't need a list entry. -# 3. pydantic-core's source (the `error.rs`/schema-kind definitions -# in the pydantic-core Rust crate) is the canonical reference if -# you want to check before it shows up in a live fuzz run. -_INTERNAL_LOC_MARKERS = frozenset( - { - "function-wrap", - "function-after", - "function-before", - "function-plain", - "json-or-python", - "lax-or-strict", - "chain", - "default", - "nullable", - "tagged-union", - "union", - "call", - "arguments", - "is-instance", - "is-subclass", - "callable", - "str", - "int", - "float", - "bool", - "bytes", - "bytearray", - "list", - "tuple", - "dict", - "set", - "frozenset", - "complex", - "none", - "nonetype", - } -) - - -def _is_internal_loc_segment(segment: str) -> bool: - """True if `segment` is a Pydantic-internal wrapper/union-branch - marker rather than a user-meaningful field name or list index.""" - if _BRACKETED_INTERNAL_RE.search(segment): - return True - return segment.lower() in _INTERNAL_LOC_MARKERS - - -def clean_loc_for_param(loc: tuple) -> str: - """Join a Pydantic error `loc` tuple into a clean dotted `param` - path, dropping internal wrapper/union-branch markers that don't - correspond to a real field name an API consumer would recognize. - - E.g. ('body', 'function-wrap[__log_extra_fields__()]', 'prompt') - -> "body.prompt", not "body.function-wrap[__log_extra_fields__()].prompt". - """ - parts = [str(p) for p in loc if not _is_internal_loc_segment(str(p))] - if not parts: - return ".".join(str(p) for p in loc) - return ".".join(parts) - - -async def validation_exception_handler(req: Request, exc: RequestValidationError): - if req.app.state.args.log_error_stack: - logger.exception( - "RequestValidationError caught. Request id: %s", - req.state.request_metadata.request_id - if hasattr(req.state, "request_metadata") - else None, - ) - - param = None - errors = exc.errors() - for error in errors: - if "ctx" in error and "error" in error["ctx"]: - ctx_error = error["ctx"]["error"] - if isinstance(ctx_error, VLLMValidationError): - param = ctx_error.parameter - break - - if param is None and errors: - first_error = errors[0] - loc = first_error.get("loc") if isinstance(first_error, dict) else None - if loc: - param = clean_loc_for_param(loc) - - # Build the message from exc.errors() instead of str(exc) - str(exc) - # leaks the server's file path via FastAPI's endpoint context. - if errors: - count = len(errors) - label = "error" if count == 1 else "errors" - message = f"{count} validation {label}:\n" - message += "".join(f" {err}\n" for err in errors) - message = message.rstrip() - else: - message = "Validation error" - - err = ErrorResponse( - error=ErrorInfo( - message=sanitize_message(message), - type=HTTPStatus.BAD_REQUEST.phrase, - code=HTTPStatus.BAD_REQUEST, - param=param, - ) - ) - return JSONResponse(err.model_dump(), status_code=HTTPStatus.BAD_REQUEST) - - _running_tasks: set[asyncio.Task] = set() diff --git a/vllm/entrypoints/speech_to_text/factories.py b/vllm/entrypoints/speech_to_text/factories.py index 1971e32b989e..21633b22f323 100644 --- a/vllm/entrypoints/speech_to_text/factories.py +++ b/vllm/entrypoints/speech_to_text/factories.py @@ -37,12 +37,6 @@ def register_speech_to_text_api_routers( app.include_router(translation_router) -def add_websocket_metrics_middleware(app: FastAPI): - from .realtime.metrics import WebSocketMetricsMiddleware - - app.add_middleware(WebSocketMetricsMiddleware) - - def init_speech_to_text_state( engine_client: "EngineClient", state: "State", diff --git a/vllm/entrypoints/speech_to_text/realtime/connection.py b/vllm/entrypoints/speech_to_text/realtime/connection.py index 32f501f10421..15d38509c52d 100644 --- a/vllm/entrypoints/speech_to_text/realtime/connection.py +++ b/vllm/entrypoints/speech_to_text/realtime/connection.py @@ -14,7 +14,7 @@ from vllm import envs from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo -from vllm.entrypoints.serve.utils.api_utils import sanitize_message +from vllm.entrypoints.serve.exception_handling.utils import sanitize_message from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger diff --git a/vllm/entrypoints/speech_to_text/transcription/protocol.py b/vllm/entrypoints/speech_to_text/transcription/protocol.py index 3220e1505099..bde69a49bd34 100644 --- a/vllm/entrypoints/speech_to_text/transcription/protocol.py +++ b/vllm/entrypoints/speech_to_text/transcription/protocol.py @@ -285,6 +285,8 @@ def to_sampling_params( @model_validator(mode="before") @classmethod def validate_transcription_request(cls, data): + if not isinstance(data, dict): + return data if isinstance(data.get("file"), str): raise HTTPException( status_code=HTTPStatus.UNPROCESSABLE_ENTITY, diff --git a/vllm/entrypoints/speech_to_text/translation/protocol.py b/vllm/entrypoints/speech_to_text/translation/protocol.py index d8836554c565..1a7319e2f105 100644 --- a/vllm/entrypoints/speech_to_text/translation/protocol.py +++ b/vllm/entrypoints/speech_to_text/translation/protocol.py @@ -271,6 +271,8 @@ def to_sampling_params( @model_validator(mode="before") @classmethod def validate_stream_options(cls, data): + if not isinstance(data, dict): + return data stream_opts = ["stream_include_usage", "stream_continuous_usage_stats"] stream = data.get("stream", False) if any(bool(data.get(so, False)) for so in stream_opts) and not stream: diff --git a/vllm/envs.py b/vllm/envs.py index 3382fa1d67c1..15705b7c1779 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -127,6 +127,7 @@ VLLM_DISABLED_KERNELS: list[str] = [] VLLM_USE_HW_AGNOSTIC: bool = False VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True + VLLM_GDN_DECODE_KERNEL: Literal["cuda", "triton"] = "cuda" VLLM_DISABLE_PYNCCL: bool = False VLLM_USE_OINK_OPS: bool = False VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True @@ -206,6 +207,8 @@ VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True VLLM_MOE_SKIP_PADDING: bool = True VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False + VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False + VLLM_KIMI_K3_GEMM_RS: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -713,9 +716,8 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_RPC_BASE_PATH", tempfile.gettempdir() ), # If true, will load models from ModelScope instead of Hugging Face Hub. - # note that the value is true or false, not numbers "VLLM_USE_MODELSCOPE": lambda: ( - os.environ.get("VLLM_USE_MODELSCOPE", "False").lower() == "true" + os.environ.get("VLLM_USE_MODELSCOPE", "False").strip().lower() in ("1", "true") ), # If true, replace the Rust BPE backend that powers HF fast tokenizers # with the `fastokens` (https://github.com/crusoecloud/fastokens) shim. @@ -1146,12 +1148,6 @@ def _resolve_rust_cli_path() -> str | None: if "VLLM_PLUGINS" not in os.environ else os.environ["VLLM_PLUGINS"].split(",") ), - # Retain local sliding-window KV checkpoints for prefix caching. - # Unset (default) preserves the dense local checkpointing behavior. `0` - # retains only the latest completed prompt boundary. Positive values retain - # checkpoints at the specified interval boundaries (rounded up to the - # prefix-cache alignment). - # Applies to sliding-window attention for now but not yet Mamba/linear attention. "VLLM_PREFIX_CACHE_RETENTION_INTERVAL": lambda: ( int(os.environ["VLLM_PREFIX_CACHE_RETENTION_INTERVAL"]) if "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ @@ -1208,6 +1204,15 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool( int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1")) ), + # Select the GDN MTP decode implementation. "cuda" uses the fused decode + # kernel where supported and falls back to "triton" otherwise; setting it + # explicitly to "cuda" raises when unsupported. + "VLLM_GDN_DECODE_KERNEL": env_with_choices( + "VLLM_GDN_DECODE_KERNEL", + "cuda", + ["cuda", "triton"], + case_sensitive=False, + ), # Disable pynccl (using torch.distributed instead) "VLLM_DISABLE_PYNCCL": lambda: ( os.getenv("VLLM_DISABLE_PYNCCL", "False").lower() in ("true", "1") @@ -1594,6 +1599,16 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool( int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0")) ), + # Kimi K3 only, and unrelated to the MoE flags above. Tap the pre-norm + # AttnRes mixture, rather than the post-mixture sum, as the auxiliary + # hidden state handed to a DFlash drafter. This changes the numerics the + # speculator sees, so it is off by default while the effect is measured. + "VLLM_KIMI_K3_AUX_ATTN_RES_STREAM": lambda: bool( + int(os.getenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "0")) + ), + # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel + # row-parallel projections. All TP ranks must belong to one NVLink domain. + "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index bd5e91cfa51f..5db3cd4faab9 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -247,8 +247,17 @@ def _maybe_init_mm( mm_budget = MultiModalBudget(vllm_config, mm_registry) limit_per_prompt = max(mm_budget.mm_max_items_per_prompt.values()) - num_encoder_tokens = self.model.get_num_mm_encoder_tokens( - mm_budget.get_encoder_budget() + max_lora_tokens = mm_budget.get_encoder_budget() + lora_token_counts_by_modality = [ + self.model.get_mm_lora_token_counts( + modality=modality, + mm_kwargs=None, + num_mm_embeds=max_lora_tokens, + ) + for modality in mm_budget.mm_max_toks_per_item + ] + num_encoder_tokens = max( + tower_tokens for tower_tokens, _ in lora_token_counts_by_modality ) # Tower wrappers @@ -263,10 +272,15 @@ def _maybe_init_mm( # Use wrapper for connector if present. if self.mm_mapping.connector: - if hasattr(self.model, "get_num_mm_connector_tokens"): - connector_tokens = self.model.get_num_mm_connector_tokens( - num_encoder_tokens - ) + connector_tokens = max( + ( + connector_tokens + for _, connector_tokens in lora_token_counts_by_modality + if connector_tokens is not None + ), + default=None, + ) + if connector_tokens is not None: connector_punica_wrapper = get_punica_wrapper( connector_tokens, max_batches=self.max_num_seqs * limit_per_prompt, diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 18272354b470..2f746f967e37 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -14,6 +14,7 @@ from vllm.lora.layers import LoRAMapping from vllm.lora.utils import get_captured_lora_counts from vllm.triton_utils import HAS_TRITON, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import round_up if HAS_TRITON: @@ -83,9 +84,11 @@ def update_metadata( self.is_prefill = mapping.is_prefill self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size) - # Prepare cuda kernel metadata tensors - self.token_mapping_meta.prepare_tensors(self.token_lora_indices) - self.prompt_mapping_meta.prepare_tensors(self.sampler_indices) + # TODO avoid gpu<->cpu sync here + with gpu_sync_allowed(): + # Prepare cuda kernel metadata tensors + self.token_mapping_meta.prepare_tensors(self.token_lora_indices) + self.prompt_mapping_meta.prepare_tensors(self.sampler_indices) def add_shrink( self, diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 7082b7287d8e..105a99807257 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -19,6 +19,7 @@ from vllm.lora.peft_helper import PEFTHelper from vllm.lora.request import LoRARequest from vllm.lora.utils import get_adapter_absolute_path +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) @@ -223,9 +224,11 @@ def _apply_adapters(self, adapter_requests: set[Any]) -> None: def add_adapter(self, adapter_request: Any) -> bool: if adapter_request.adapter_id in self.list_adapters(): return False - loaded_adapter = self._load_adapter(adapter_request) - loaded = self._adapter_manager.add_adapter(loaded_adapter) - self._adapter_manager.activate_adapter(loaded_adapter.id) + # One-time per adapter + with gpu_sync_allowed(): + loaded_adapter = self._load_adapter(adapter_request) + loaded = self._adapter_manager.add_adapter(loaded_adapter) + self._adapter_manager.activate_adapter(loaded_adapter.id) return loaded def remove_adapter(self, adapter_id: int) -> bool: @@ -288,32 +291,34 @@ def add_adapter(self, lora_request: LoRARequest) -> bool: # This is ok because it's currently only called from # the single-threaded core engine loop. - if ( - lora_request.lora_int_id not in self.list_adapters() - or lora_request.load_inplace - ): - # Load the new adapter first to ensure it is actually valid, before - # evicting any existing adapters. - # This may cause the # of loaded lora adapters to very temporarily - # exceed `--max-cpu-loras`. - lora = self._load_adapter(lora_request) - - # Remove the existing adapter if it exists - # Use case for LoRA inplace - self._adapter_manager.remove_adapter(lora.id) - - # Loading succeeded, now check if we will exceed cache capacity and - # evict if the oldest adapter if so - if len(self._adapter_manager) + 1 > self._adapter_manager.capacity: - assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager) - self._adapter_manager.remove_oldest_adapter() - # Then add the new adapter to the cache - loaded = self._adapter_manager.add_adapter(lora) - else: - # If the lora is already loaded, just touch it to - # update its position in the caches - loaded = ( - self._adapter_manager.get_adapter(lora_request.lora_int_id) is not None - ) - self._adapter_manager.activate_adapter(lora_request.lora_int_id) + with gpu_sync_allowed(): + if ( + lora_request.lora_int_id not in self.list_adapters() + or lora_request.load_inplace + ): + # Load the new adapter first to ensure it is actually valid, before + # evicting any existing adapters. + # This may cause the # of loaded lora adapters to very temporarily + # exceed `--max-cpu-loras`. + lora = self._load_adapter(lora_request) + + # Remove the existing adapter if it exists + # Use case for LoRA inplace + self._adapter_manager.remove_adapter(lora.id) + + # Loading succeeded, now check if we will exceed cache capacity and + # evict if the oldest adapter if so + if len(self._adapter_manager) + 1 > self._adapter_manager.capacity: + assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager) + self._adapter_manager.remove_oldest_adapter() + # Then add the new adapter to the cache + loaded = self._adapter_manager.add_adapter(lora) + else: + # If the lora is already loaded, just touch it to + # update its position in the caches + loaded = ( + self._adapter_manager.get_adapter(lora_request.lora_int_id) + is not None + ) + self._adapter_manager.activate_adapter(lora_request.lora_int_id) return loaded diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 9915255a847d..3ba499155b6c 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -77,6 +77,9 @@ from vllm.model_executor.kernels.linear.mxfp4.aiter import ( AiterMxfp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.b12x import ( + B12xMxFp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.emulation import ( EmulationMxfp4LinearKernel, ) @@ -103,6 +106,9 @@ Mxfp8LinearKernel, Mxfp8LinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mxfp8.b12x import ( + B12xMxfp8LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp8.emulation import ( EmulationMxfp8LinearKernel, ) @@ -126,6 +132,9 @@ NvFp4LinearKernel, NvFp4LinearLayerConfig, ) +from vllm.model_executor.kernels.linear.nvfp4.b12x import ( + B12xNvFp4LinearKernel, +) from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( CutlassNvFp4LinearKernel, ) @@ -163,6 +172,10 @@ AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( + B12xFp8BlockScaledMMKernel, + B12xTensorFP8ScaledMMLinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( CPUFp8BlockScaledMMKernel, CPUInt8ScaledMMLinearKernel, @@ -228,6 +241,13 @@ def _get_linear_backend() -> str: # set are considered candidates. If none can implement the layer config, # an error is raised to respect the user's explicit intent. _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { + "b12x": { + B12xFp8BlockScaledMMKernel, + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, + B12xTensorFP8ScaledMMLinearKernel, + }, "cutlass": { CutlassInt8ScaledMMLinearKernel, CutlassFP8ScaledMMLinearKernel, @@ -377,6 +397,7 @@ def _resolve_backend_kernels( MarlinFP8ScaledMMLinearKernel, FlashInferFP8ScaledMMLinearKernel, CutlassFP8ScaledMMLinearKernel, + B12xTensorFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel, ChannelWiseTorchFP8ScaledMMLinearKernel, HummingFP8ScaledMMLinearKernel, @@ -411,6 +432,7 @@ def _resolve_backend_kernels( FlashInferFp8DeepGEMMDynamicBlockScaledKernel, DeepGemmFp8BlockScaledMMKernel, CutlassFp8BlockScaledMMKernel, + B12xFp8BlockScaledMMKernel, MarlinFP8ScaledMMLinearKernel, TritonFp8BlockScaledMMKernel, HummingFP8ScaledMMLinearKernel, @@ -482,6 +504,7 @@ def _resolve_backend_kernels( FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel, MarlinMxfp8LinearKernel, + B12xMxfp8LinearKernel, EmulationMxfp8LinearKernel, HummingMxfp8LinearKernel, ], @@ -507,6 +530,7 @@ def _resolve_backend_kernels( FlashInferTrtllmNvFp4LinearKernel, FlashInferCudnnNvFp4LinearKernel, FbgemmNvFp4LinearKernel, + B12xNvFp4LinearKernel, EmulationNvFp4LinearKernel, HummingNvFp4LinearKernel, ], @@ -529,6 +553,7 @@ def _resolve_backend_kernels( FlashInferMxFp4LinearKernel, MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel, + B12xMxFp4LinearKernel, EmulationMxfp4LinearKernel, ], PlatformEnum.ROCM: [ @@ -1175,6 +1200,9 @@ def register_linear_kernel( "init_mxfp8_linear_kernel", "Mxfp8LinearKernel", "Mxfp8LinearLayerConfig", + "B12xMxfp8LinearKernel", + "B12xMxFp4LinearKernel", + "B12xNvFp4LinearKernel", "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", @@ -1203,4 +1231,6 @@ def register_linear_kernel( "_KernelT", "DeepGemmFp8BlockScaledMMKernel", "FlashInferFp8DeepGEMMDynamicBlockScaledKernel", + "B12xFp8BlockScaledMMKernel", + "B12xTensorFP8ScaledMMLinearKernel", ] diff --git a/vllm/model_executor/kernels/linear/base.py b/vllm/model_executor/kernels/linear/base.py index 416b6ea1c1b6..036b30b7c46d 100644 --- a/vllm/model_executor/kernels/linear/base.py +++ b/vllm/model_executor/kernels/linear/base.py @@ -89,28 +89,6 @@ def from_layer(cls, layer: torch.nn.Module) -> "FP8Params": ) -@dataclass -class Int8Params(Params): - """Int8 layer parameters with typed fields""" - - input_zero_point: torch.Tensor | None - azp_adj: torch.Tensor | None - - INPUT_ZERO_POINT: ClassVar[str] = "input_zero_point" - AZP_ADJ: ClassVar[str] = "azp_adj" - - @classmethod - def from_layer(cls, layer: torch.nn.Module) -> "Int8Params": - """Extract parameters from layer""" - return cls( - weight=getattr(layer, cls.WEIGHT), - weight_scale=getattr(layer, cls.WEIGHT_SCALE), - input_scale=getattr(layer, cls.INPUT_SCALE, None), - input_zero_point=getattr(layer, cls.INPUT_ZERO_POINT, None), - azp_adj=getattr(layer, cls.AZP_ADJ, None), - ) - - _ParamsT = TypeVar("_ParamsT", bound=Params) _ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig) @@ -130,7 +108,7 @@ class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]): Typical Usage: 1. Define a config dataclass inheriting from MMLinearLayerConfig - 2. Define a params dataclass inheriting from Params (or FP8Params/Int8Params) + 2. Define a params dataclass inheriting from Params 3. Subclass MMLinearKernel with your config and params types 4. Implement all abstract methods 5. Register the kernel with the quantization method diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py index bc0a587b6763..8a4a493c6c55 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -400,29 +400,40 @@ def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter: if self.w_zp_name is not None: zp = getattr(layer, self.w_zp_name, None) if zp is not None: - c = self.config - K, N = c.partition_weight_shape - group_size = c.group_size if c.group_size != -1 else K - expected_shape = (K // group_size, N // 8) - transposed_shape = (N // 8, K // group_size) - - if tuple(zp.data.shape) == expected_shape: - # GPTQ/AutoGPTQ already stores qzeros in the kernel layout. - qzeros = zp.data.contiguous() - elif tuple(zp.data.shape) == transposed_shape: - # Compressed-tensors stores qzeros transposed from what the - # kernel needs. - qzeros = zp.data.t().contiguous() + # Kernel needs [K//G, N//8]: + # input(K) at dim 0, output(N) packed at dim 1. + # AutoGPTQ: + # output_dim=1 -> already [K//G, N//8], no transpose. + # compressed-tensors: + # output_dim=0 -> [N//8, K//G], needs transpose. + # None (unknown): + # infer from shape; if square (ambiguous), default to transpose. + zp_output_dim = getattr(zp, "output_dim", None) + if zp_output_dim is not None: + needs_transpose = zp_output_dim != 1 else: - raise AssertionError( - f"{self.w_zp_name} shape mismatch: {zp.data.shape}; " - f"expected {expected_shape} or {transposed_shape}" - ) - + # in case output_dim is None + c = self.config + K, N = c.partition_weight_shape + group_size = c.group_size if c.group_size != -1 else K + expected_shape = (K // group_size, N // 8) + transposed_shape = (N // 8, K // group_size) + if ( + tuple(zp.data.shape) == expected_shape + and expected_shape != transposed_shape + ): + needs_transpose = False + else: + needs_transpose = True + zp_data = ( + zp.data.t().contiguous() + if needs_transpose + else zp.data.contiguous() + ) replace_parameter( layer, self.w_zp_name, - torch.nn.Parameter(qzeros, requires_grad=False), + torch.nn.Parameter(zp_data, requires_grad=False), ) def apply_weights( @@ -437,17 +448,18 @@ def apply_weights( K = c.partition_weight_shape[0] group_size = c.group_size if c.group_size != -1 else K - # For symmetric types (uint4b8), use the scalar bias; no zeros tensor. - # Some checkpoint loaders still register qzeros parameters for GPTQ - # layers, but they are not part of the symmetric kernel contract. - zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0 - qzeros = None if c.weight_type.has_bias() else w_zp + # For symmetric types (uint4b8), use the scalar bias; no zeros tensor + if c.weight_type.has_bias(): + zp_bias = c.weight_type.bias + w_zp = None # symmetric: ignore qzeros, use scalar bias instead + else: + zp_bias = 0 output = triton_w4a16_gemm( a=x_2d, b_q=w_q, scales=w_s, - qzeros=qzeros, + qzeros=w_zp, group_size=group_size, zp_bias=zp_bias, ) diff --git a/vllm/model_executor/kernels/linear/mxfp4/b12x.py b/vllm/model_executor/kernels/linear/mxfp4/b12x.py new file mode 100644 index 000000000000..a4649d3c4f8e --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/b12x.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit +from vllm.utils.b12x import ( + get_b12x_blockscaled as _import_b12x_blockscaled, +) +from vllm.utils.b12x import get_b12x_intrinsics as _import_b12x_intrinsics + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + + +def _apply_b12x_mxfp4_linear( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale_storage: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + from vllm.utils.flashinfer import flashinfer_mxfp4_quantize + + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + + output_size = int(weight.shape[0]) + output_shape = [*x.shape[:-1], output_size] + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + x_packed, x_scale_swizzled = flashinfer_mxfp4_quantize(x_2d, backend="cute-dsl") + output = blockscaled.mm_mxfp4( + x_packed, + x_scale_swizzled, + weight, + weight_scale_storage, + out_dtype=x.dtype, + ) + if bias is not None: + output = output + bias + return output.view(*output_shape) + + +class B12xMxFp4LinearKernel(MxFp4LinearKernel): + """MXFP4 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X MXFP4 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X MXFP4 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None or _import_b12x_intrinsics() is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "b12x native MXFP4 GEMM is not supported" + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "B12X MXFP4 GEMM requires dynamic MXFP4 activations" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + intrinsics = _import_b12x_intrinsics() + assert intrinsics is not None + replace_parameter( + layer, + "weight_scale", + intrinsics.swizzle_block_scale(layer.weight_scale.data), + ) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = layer.weight_scale + n, packed_k = map(int, weight.shape) + k = packed_k * 2 + + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_mxfp4_linear(source, weight, weight_scale, None) + + return B12xWarmupUnit( + name="MXFP4", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return _apply_b12x_mxfp4_linear( + x, + layer.weight, + layer.weight_scale, + bias, + ) + + +__all__ = ["B12xMxFp4LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/mxfp8/b12x.py b/vllm/model_executor/kernels/linear/mxfp8/b12x.py new file mode 100644 index 000000000000..48343f3d694b --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/b12x.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage +from vllm.utils.b12x import ( + get_b12x_mxfp8_linear as _import_b12x_mxfp8, +) +from vllm.utils.torch_utils import current_stream + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +def _apply_b12x_mxfp8_packed_linear( + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + packed_weight = layer.b12x_mxfp8_packed_weight + + input_2d = x.reshape(-1, x.shape[-1]).contiguous() + output_shape = [*x.shape[:-1], int(packed_weight.out_features)] + + mxfp8 = _import_b12x_mxfp8() + assert mxfp8 is not None + output = mxfp8.mm( + input_2d, + packed_weight, + bias=bias, + expected_m=max(1, int(input_2d.shape[0])), + ) + return output.view(*output_shape) + + +class B12xMxfp8LinearKernel(Mxfp8LinearKernel): + """ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "b12x MXFP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x MXFP8 kernels require a Blackwell 12x device" + mxfp8 = _import_b12x_mxfp8() + if mxfp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not mxfp8.is_supported(): + return False, "b12x.gemm.mxfp8_linear is not supported" + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + del c + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight.data + assert weight.dtype == MXFP8_VALUE_DTYPE, ( + f"b12x MXFP8 requires {MXFP8_VALUE_DTYPE}, got {weight.dtype}" + ) + assert weight.ndim == 2, f"b12x MXFP8 weight must be 2D, got {weight.ndim}D" + assert hasattr(layer, "weight_scale"), "b12x MXFP8 linear requires weight_scale" + + out_features, in_features = map(int, weight.shape) + assert in_features % MXFP8_BLOCK_SIZE == 0, ( + "b12x MXFP8 requires input features divisible by " + f"{MXFP8_BLOCK_SIZE}, got {in_features}" + ) + weight_scale = layer.weight_scale.data + assert weight_scale.dtype == MXFP8_SCALE_DTYPE, ( + f"b12x MXFP8 requires {MXFP8_SCALE_DTYPE} weight_scale, " + f"got {weight_scale.dtype}" + ) + assert weight_scale.ndim == 2, ( + f"b12x MXFP8 weight_scale must be 2D, got {weight_scale.ndim}D" + ) + + mxfp8 = _import_b12x_mxfp8() + assert mxfp8 is not None + scale_k = in_features // MXFP8_BLOCK_SIZE + packed_weight = mxfp8.pack_weight( + weight[:out_features, :in_features].detach(), + weight_scale[:out_features, :scale_k].detach(), + ) + layer.b12x_mxfp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_mxfp8_packed_weight", None), + packed_weight, + ) + replace_parameter(layer, "weight", weight.new_empty((0,))) + replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,))) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + packed_weight = layer.b12x_mxfp8_packed_weight + device = torch.device(packed_weight.weight.values.device) + + def compile() -> None: + mxfp8 = _import_b12x_mxfp8() + assert mxfp8 is not None + for tokens in token_counts: + source = torch.zeros( + (tokens, int(packed_weight.in_features)), + dtype=output_dtype, + device=device, + ) + mxfp8.mm( + source, + packed_weight, + expected_m=max(1, int(tokens)), + stream=current_stream().cuda_stream, + ) + + return B12xWarmupUnit( + name="MXFP8", + key=( + type(self), + device, + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return _apply_b12x_mxfp8_packed_linear(layer, x, bias) diff --git a/vllm/model_executor/kernels/linear/nvfp4/b12x.py b/vllm/model_executor/kernels/linear/nvfp4/b12x.py new file mode 100644 index 000000000000..4ff531dd48be --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/b12x.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm._custom_ops import scaled_fp4_quant +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit +from vllm.utils.b12x import ( + get_b12x_blockscaled as _import_b12x_blockscaled, +) +from vllm.utils.b12x import get_b12x_intrinsics as _import_b12x_intrinsics + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + + +def _apply_b12x_nvfp4_linear( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale_storage: torch.Tensor, + input_global_scale_inv: torch.Tensor, + alpha: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + + output_size = int(weight.shape[0]) + output_shape = [*x.shape[:-1], output_size] + x_2d = x.reshape(-1, x.shape[-1]) + x_packed, x_scale_swizzled = scaled_fp4_quant( + x_2d, + input_global_scale_inv, + is_sf_swizzled_layout=True, + ) + output = blockscaled.mm_nvfp4( + x_packed, + x_scale_swizzled, + weight, + weight_scale_storage, + alpha, + out_dtype=x.dtype, + ) + if bias is not None: + output = output + bias + return output.view(*output_shape) + + +class B12xNvFp4LinearKernel(NvFp4LinearKernel): + """ModelOpt NVFP4 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X NVFP4 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X NVFP4 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None or _import_b12x_intrinsics() is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "b12x native NVFP4 GEMM is not supported" + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + del config + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + intrinsics = _import_b12x_intrinsics() + assert intrinsics is not None + replace_parameter( + layer, + "weight_scale", + intrinsics.swizzle_block_scale(layer.weight_scale.data), + ) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = layer.weight_scale + n, packed_k = map(int, weight.shape) + k = packed_k * 2 + + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_nvfp4_linear( + source, + weight, + weight_scale, + layer.input_global_scale_inv, + layer.alpha, + None, + ) + + return B12xWarmupUnit( + name="NVFP4", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return _apply_b12x_nvfp4_linear( + x, + layer.weight, + layer.weight_scale, + layer.input_global_scale_inv, + layer.alpha, + bias, + ) + + +__all__ = ["B12xNvFp4LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py new file mode 100644 index 000000000000..62e839c7291c --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage +from vllm.utils.b12x import ( + get_b12x_blockscaled as _import_b12x_blockscaled, +) +from vllm.utils.b12x import ( + get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8, +) +from vllm.utils.torch_utils import current_stream + +from .BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) +from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel + + +def _run_b12x_fp8_block_scaled_mm( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + + return blockscaled.mm_block_fp8( + a, + a_scale, + weight, + weight_scale, + out_dtype=out_dtype, + ) + + +class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """K128 block-FP8 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X FP8 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "B12X regular block-FP8 GEMM is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + can_implement_base, reason = super().can_implement(config) + if not can_implement_base: + return can_implement_base, reason + + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "Supports only bf16/fp16 input dtype" + if config.input_dtype != config.out_dtype: + return False, "Input and output dtype must match" + + act_group_shape = config.activation_quant_key.scale.group_shape + if act_group_shape != GroupShape(1, 128): + return ( + False, + "Supports only dynamic per-token group activation quantization " + "with group_shape=(1,128)", + ) + weight_group_shape = config.weight_quant_key.scale.group_shape + if weight_group_shape != GroupShape(128, 128): + return False, "Supports only 128x128 block-scaled FP8 weights" + + out_features, in_features = config.weight_shape + if in_features <= 0 or in_features % 128 != 0: + return False, "Input features must be a positive multiple of 128" + if out_features <= 0 or out_features % 128 != 0: + return False, "Output features must be a positive multiple of 128" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + params = self._get_layer_params(layer) + if params.weight_scale_inv is not None: + weight_scale = params.weight_scale_inv + scale_attr = params.WEIGHT_SCALE_INV + else: + weight_scale = params.weight_scale + scale_attr = params.WEIGHT_SCALE + if weight_scale is not None and weight_scale.dtype in ( + torch.float8_e8m0fnu, + torch.uint8, + ): + # TODO: Remove once B12X supports 128x128 UE8M0 block scales. + replace_parameter( + layer, + scale_attr, + _upcast_e8m0_to_fp32(weight_scale).contiguous(), + ) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = getattr(layer, "weight_scale_inv", None) + if weight_scale is None: + weight_scale = layer.weight_scale + n, k = map(int, weight.shape) + + def compile() -> None: + for tokens in token_counts: + a = torch.empty((tokens, k), dtype=weight.dtype, device=weight.device) + a_scale = torch.empty( + (tokens, k // 128), + dtype=torch.float32, + device=weight.device, + ) + _run_b12x_fp8_block_scaled_mm( + a, weight, a_scale, weight_scale, output_dtype + ) + + return B12xWarmupUnit( + name="block-FP8", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + return _run_b12x_fp8_block_scaled_mm( + A, + B, + As, + Bs, + self.config.out_dtype, + ) + + +def _apply_b12x_tensor_fp8_packed_linear( + layer: torch.nn.Module, + x_q: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + packed_weight = layer.b12x_tensor_fp8_packed_weight + + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + + input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous() + output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)] + output = tensor_fp8.mm( + input_2d, + packed_weight, + bias=bias, + out_dtype=out_dtype, + expected_m=max(1, int(input_2d.shape[0])), + ) + return output.view(*output_shape) + + +class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """Static per-tensor FP8 linear through the B12X SM12x dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "b12x tensor FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x tensor FP8 kernels require a Blackwell 12x device" + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not tensor_fp8.is_supported(): + return False, "b12x.gemm.tensor_fp8_linear is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + activation_scale = config.activation_quant_key.scale + weight_scale = config.weight_quant_key.scale + if ( + not activation_scale.static + or not activation_scale.group_shape.is_per_tensor() + ): + return False, "requires static per-tensor activation scales" + if not weight_scale.static or not weight_scale.group_shape.is_per_tensor(): + return False, "requires static per-tensor weight scales" + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 input dtype" + if config.out_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 output dtype" + out_features, in_features = config.weight_shape + if out_features <= 0 or in_features <= 0 or in_features % 32 != 0: + return False, "weight dimensions must be positive with K divisible by 32" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight, weight_scale, input_scale, _ = self._get_layer_params(layer) + assert weight.dtype == torch.float8_e4m3fn + assert input_scale is not None + assert weight_scale.numel() == input_scale.numel() == 1 + + out_features, in_features = map(int, self.config.weight_shape) + assert tuple(weight.shape) == (in_features, out_features) + + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + output_scale = ( + input_scale.detach().to(torch.float32).reshape(1) + * weight_scale.detach().to(torch.float32).reshape(1) + ).contiguous() + packed_weight = tensor_fp8.pack_weight( + weight.detach().T.contiguous(), + output_scale, + ) + layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_tensor_fp8_packed_weight", None), + packed_weight, + ) + weight_name, weight_scale_name, _, _ = self.layer_param_names + replace_parameter(layer, weight_name, weight.new_empty((0,))) + replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,))) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + packed_weight = layer.b12x_tensor_fp8_packed_weight + device = torch.device(packed_weight.values.device) + + def compile() -> None: + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + tensor_fp8.prewarm( + packed_weight, + token_counts, + out_dtype=output_dtype, + stream=current_stream().cuda_stream, + ) + + return B12xWarmupUnit( + name="tensor FP8", + key=( + type(self), + device, + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert isinstance(x, torch.Tensor) + _, _, input_scale, input_scale_ub = self._get_layer_params(layer) + input_2d = x.reshape(-1, x.shape[-1]) + x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub) + out_dtype = self.config.out_dtype + output = _apply_b12x_tensor_fp8_packed_linear( + layer, + x_q, + bias, + out_dtype, + ) + return output.view(*x.shape[:-1], output.shape[-1]) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + del A, B, out_dtype, As, Bs, bias, output_shape + raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights") + + +__all__ = [ + "B12xFp8BlockScaledMMKernel", + "B12xTensorFP8ScaledMMLinearKernel", +] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index e6ec7926c4bd..92d6458bb3df 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from collections.abc import Sequence import torch @@ -204,7 +205,34 @@ def process_weights_after_loading(self, layer: torch.nn.Module): ) scale = getattr(layer, scale_attr) - # Checkpoint scale is [n_blocks, k_blocks] (one value per 128x128 tile). + # Ragged N (N % block_n != 0): oneDNN needs n_blocks to divide N. + # Weight untouched; only repeat scale rows to a finer N-group gn that + # divides both N and block_n (gn = gcd(N, block_n)): + # scale [ceil(N/block_n), K/block_k] --> [N/gn, K/block_k] + # oneDNN only accepts gn that is a multiple of 16, and gcd(N, block_n) + # is a power of two (block_n=128), so gn must be >= 16. No-op when + # N % block_n == 0. + block_n, block_k = self.weight_group_shape + N, K = layer.weight.shape + if N % block_n != 0: + gn = math.gcd(N, block_n) + assert gn % 16 == 0, ( + f"XPU block-scaled FP8: N ({N}) yields group width {gn}, but " + f"oneDNN only supports multiples of 16; this weight shape is " + f"unsupported." + ) + col_start = torch.arange(N // gn, device=scale.device) * gn + src_idx = torch.div(col_start, block_n, rounding_mode="floor") + scale = scale.index_select(0, src_idx).contiguous() + + # Ragged K needs the runtime activation scale expanded too, which we + # don't handle; DeepSeek/GLM keep K block-aligned, so fail loudly. + assert K % block_k == 0, ( + f"XPU block-scaled FP8 requires K ({K}) to be a multiple of the " + f"weight block size ({block_k}); ragged-K weights are unsupported." + ) + + # Checkpoint scale is [n_blocks, k_blocks] (one value per block tile). # oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout. # We store the transposed contiguous buffer as a .t() view so that: # - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 3afd7069aa3e..925f4631a515 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -1,29 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + import math from functools import cache -from typing import TYPE_CHECKING, Any import torch from vllm.platforms import current_platform -from vllm.utils.import_utils import has_tilelang +from vllm.tilelang_utils import T, tilelang, tilelang_jit from vllm.utils.math_utils import cdiv -# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so -# registering the Python wrapper modules does not require TileLang everywhere. -if TYPE_CHECKING or current_platform.is_cuda_alike(): - if not has_tilelang(): - raise ImportError( - "tilelang is required for mhc but is not installed. Install it with " - "`pip install tilelang`." - ) - import tilelang - import tilelang.language as T -else: - tilelang = None # type: ignore[assignment] - T = None # type: ignore[assignment] - ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() @@ -40,18 +28,7 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k -pass_configs: dict[tilelang.PassConfigKey, Any] = { - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, -} - -if current_platform.is_cuda(): - pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10 - - -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def mhc_pre_big_fuse_tilelang( gemm_out_mul, gemm_out_sqrsum, @@ -191,9 +168,7 @@ def mhc_pre_big_fuse_tilelang( # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478 -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def mhc_pre_big_fuse_with_norm_tilelang( gemm_out_mul, gemm_out_sqrsum, @@ -353,9 +328,7 @@ def mhc_pre_big_fuse_with_norm_tilelang( T.pdl_trigger() -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def mhc_pre_big_fuse_broadcast_with_norm_tilelang( gemm_out_mul, gemm_out_sqrsum, @@ -517,9 +490,7 @@ def mhc_pre_big_fuse_broadcast_with_norm_tilelang( T.pdl_trigger() -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def mhc_fused_tilelang( comb_mix, residual_in, @@ -640,9 +611,7 @@ def mhc_fused_tilelang( T.pdl_trigger() -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def mhc_post_tilelang( a, b, @@ -695,9 +664,7 @@ def mhc_post_tilelang( T.pdl_trigger() -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def hc_prenorm_gemm_tilelang( x, fn, @@ -781,9 +748,7 @@ def hc_prenorm_gemm_tilelang( T.pdl_trigger() -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def hc_prenorm_gemm_block_m_tilelang( x, fn, @@ -878,9 +843,7 @@ def hc_prenorm_gemm_block_m_tilelang( T.pdl_trigger() -@tilelang.jit( - pass_configs=pass_configs, -) +@tilelang_jit def hc_head_fuse_tilelang( residual, fn, diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 8eb8aaa86695..b4831e2a0b41 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -8,7 +8,10 @@ import vllm.envs as envs from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import CacheConfig, get_current_vllm_config +from vllm.config import ( + CacheConfig, + get_current_vllm_config, +) from vllm.config.vllm import VllmConfig from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger @@ -94,6 +97,7 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool: def _largest_kernel_block_within( attn_backend: "type[AttentionBackend]", + vllm_config: VllmConfig, per_token_bytes: int, page_budget: int | None, fallback: int, @@ -108,7 +112,7 @@ def _largest_kernel_block_within( """ from vllm.v1.attention.backend import MultipleOf - sizes = attn_backend.get_supported_kernel_block_sizes() + sizes = attn_backend.get_supported_kernel_block_sizes_for_config(vllm_config) candidates = [s for s in sizes if isinstance(s, int)] if not candidates: candidates = [s.base for s in sizes if isinstance(s, MultipleOf)] @@ -619,17 +623,24 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: # bytes per block. Otherwise (page_size_padded is None) the smallest # block is fine — ``unify`` scales it up by an integer ratio. shared_page = vllm_config.cache_config.skip_page_size_padded - sw_per_token = SlidingWindowSpec( - block_size=1, - num_kv_heads=self.num_kv_heads, - head_size=self.head_size, - head_size_v=self.head_size_v, - dtype=self.kv_cache_torch_dtype, - kv_quant_mode=quant_mode, - sliding_window=self.sliding_window, + # The backend owns its packing + sw_per_token = self.attn_backend.customize_spec( + SlidingWindowSpec( + block_size=1, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size_v, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=quant_mode, + sliding_window=self.sliding_window, + ) ).real_page_size_bytes sw_block_size = _largest_kernel_block_within( - self.attn_backend, sw_per_token, shared_page, block_size + self.attn_backend, + vllm_config, + sw_per_token, + shared_page, + block_size, ) return SlidingWindowSpec( block_size=sw_block_size, @@ -641,24 +652,6 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: sliding_window=self.sliding_window, page_size_padded=shared_page, ) - elif self.kv_cache_dtype.startswith("turboquant_"): - from vllm.model_executor.layers.quantization.turboquant.config import ( - TurboQuantConfig, - ) - from vllm.v1.kv_cache_interface import TQFullAttentionSpec - - tq_config = TurboQuantConfig.from_cache_dtype( - self.kv_cache_dtype, self.head_size - ) - return TQFullAttentionSpec( - block_size=block_size, - num_kv_heads=self.num_kv_heads, - head_size=self.head_size, - head_size_v=self.head_size, - dtype=self.kv_cache_torch_dtype, - kv_quant_mode=quant_mode, - tq_slot_size=tq_config.slot_size_aligned, - ) else: return FullAttentionSpec( block_size=block_size, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index c57bd0ea3176..72d65c15748d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -204,7 +204,7 @@ import math from abc import abstractmethod from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import Enum from math import lcm from typing import ClassVar, Generic, TypeVar, cast @@ -270,6 +270,7 @@ _encode_layer_name, _resolve_layer_name, direct_register_custom_op, + get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, np_to_pinned_tensor, @@ -298,6 +299,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheSpec, + KVQuantMode, MLAAttentionSpec, SlidingWindowMLASpec, get_kv_quant_mode, @@ -1152,6 +1154,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: dtype=kv_cache_dtype, cache_dtype_str=self.kv_cache_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + # fp8_ds_mla: 656-byte custom layout (kv_lora_rank=512 + + # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py. + state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None, ) if self.sliding_window is not None: return SlidingWindowMLASpec( @@ -1370,6 +1375,20 @@ def forward( class MLACommonBackend(AttentionBackend): + @classmethod + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": + """Per-token-head modes pack an inline fp32 scale pair after the + latent data (single-sided: ``head_size_v == 0`` for MLA).""" + mode = spec.kv_quant_mode + if spec.state_content_bytes is not None or not mode.is_per_token_head: + return spec + head_size = spec.head_size + if mode == KVQuantMode.INT4_PER_TOKEN_HEAD: + head_size //= 2 + scale_bytes = get_dtype_size(torch.float32) + content = head_size * get_dtype_size(spec.dtype) + 2 * scale_bytes + return replace(spec, state_content_bytes=content) + @staticmethod def get_name() -> str: return "TRITON_MLA" @@ -1965,6 +1984,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # Whether this builder can flatten a non-causal query block into decode rows. supports_non_causal_multi_token_decode: ClassVar[bool] = False + # Whether can support non-causal multi-token decode with DCP KV cache. + supports_non_causal_multi_token_dcp: ClassVar[bool] = False + # The threshold for reordering the batch into decode and prefill requests. # If > 1, the batch will be reordered such that requests with # query length <= threshold are classified as decode requests. @@ -1972,6 +1994,31 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # when speculative decoding is enabled. reorder_batch_threshold: int = 1 + def _validate_dspark_dcp_support(self, supports_dcp_with_varlen: bool) -> None: + speculative_config = getattr(self.vllm_config, "speculative_config", None) + parallel_config = self.vllm_config.parallel_config + if ( + speculative_config is None + or getattr(speculative_config, "method", None) != "dspark" + or parallel_config.decode_context_parallel_size <= 1 + ): + return + + if self.non_causal_multi_token_decode: + supported = self.supports_non_causal_multi_token_dcp + query_mode = "non-causal draft" + else: + supported = supports_dcp_with_varlen + query_mode = "causal multi-token" + + if not supported: + raise ValueError( + f"{type(self).__name__} does not support {query_mode} MLA " + "attention for DSpark with decode context parallelism. Select " + "a backend with explicit DSpark DCP support or set " + "decode_context_parallel_size=1." + ) + @staticmethod def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int: scheduler_config = vllm_config.scheduler_config @@ -2065,6 +2112,7 @@ def __init__( self.non_causal_multi_token_decode = getattr( kv_cache_spec, "non_causal_multi_token_decode", False ) + self._validate_dspark_dcp_support(supports_dcp_with_varlen) # A draft cache group can have a different head count from the target. self.num_heads = get_num_attention_heads_from_layers( diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 000000000000..f47e4c67b456 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.5.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 000000000000..f47e4c67b456 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=512,N=512,device_name=NVIDIA_GB10,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.5.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4 + } +} diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index 77afcf0d66b0..1614e41b3303 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -111,21 +111,15 @@ def _per_expert(value: float | None) -> torch.Tensor | None: self.gemm1_alpha = _per_expert(quant_config.gemm1_alpha) self.gemm1_beta = _per_expert(quant_config.gemm1_beta) - if quant_config.weight_quant_dtype == "mxfp4": - # This value is used specifically for gpt-oss, - # Need to revisit this for other models - if self.gemm1_alpha is None: - self.gemm1_alpha = _per_expert(1.702) - if self.gemm1_beta is None: - self.gemm1_beta = _per_expert(1.0) - if self.gemm1_clamp_limit is None: - self.gemm1_clamp_limit = _per_expert(7.0) - if quant_config.quant_dtype == "mxfp8": - self.fake_input_scale = torch.ones( - self.num_experts, - device=self.device, - dtype=torch.float32, - ) + if ( + quant_config.weight_quant_dtype == "mxfp4" + and quant_config.quant_dtype == "mxfp8" + ): + self.fake_input_scale = torch.ones( + self.num_experts, + device=self.device, + dtype=torch.float32, + ) @property def expects_unquantized_inputs(self) -> bool: @@ -324,9 +318,6 @@ def apply( elif self.weight_quant_dtype == "mxfp4": assert self.w1_scale is not None and self.w2_scale is not None assert w1.is_contiguous() and w2.is_contiguous() - assert self.gemm1_alpha is not None - assert self.gemm1_beta is not None - assert self.gemm1_clamp_limit is not None assert topk_ids.is_contiguous() fc1_expert_biases = self.w1_bias 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 3b80b5b1e3fa..a433b18487dd 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 @@ -484,6 +484,7 @@ def _supports_activation(activation: MoEActivation) -> bool: return activation in [ MoEActivation.SILU, MoEActivation.GELU, + MoEActivation.SITU, MoEActivation.SWIGLUOAI, MoEActivation.SWIGLUOAI_UNINTERLEAVE, ] diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index be4930052a9a..8bfebcad1193 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -26,8 +26,10 @@ moe_align_block_size, ) from vllm.model_executor.layers.fused_moe.utils import ( + TD_MIN_GATHER_ROWS, enable_swap_ab, moe_kernel_quantize_input, + moe_use_td_hw_supported, resolve_moe_use_td, warn_if_moe_use_td_ineffective, ) @@ -109,6 +111,11 @@ def fused_moe_kernel_gptq_awq( has_zp: tl.constexpr, use_int4_w4a16: tl.constexpr, use_int8_w8a16: tl.constexpr, + # Tensor-descriptor path for the A gather and B load in the K-loop, for + # int4_w4a16 only: B is packed 2 nibbles/byte along K, so its descriptor is + # half-K wide and the tile is rebuilt with tl.interleave. False keeps the + # pointer path. + USE_TD: tl.constexpr = False, ): """ Implements the fused computation for a Mixture of Experts (MOE) using @@ -184,25 +191,47 @@ def fused_moe_kernel_gptq_awq( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + ( - offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak - ) - if use_int4_w4a16: - b_ptrs = ( - b_ptr - + off_experts * stride_be - + (offs_k[:, None] // 2) * stride_bk - + offs_bn[None, :] * stride_bn + if USE_TD: + # Activations are never quantized here, so A's descriptor does not + # depend on the weight layout (and matches fused_moe_kernel's). + m_td = num_valid_tokens // top_k + a_desc = tl.make_tensor_descriptor( + base=a_ptr, + shape=(m_td, K), + strides=(stride_am, stride_ak), + block_shape=(1, BLOCK_SIZE_K), ) - b_shifter = (offs_k[:, None] % 2) * 4 - elif use_int8_w8a16: - b_ptrs = ( - b_ptr - + off_experts * stride_be - + offs_k[:, None] * stride_bk - + offs_bn[None, :] * stride_bn + # gather() requires i32 indices; a row index fits int32 even though the + # stride products elsewhere are kept int64 against overflow. + gather_idx = (offs_token // top_k).to(tl.int32) + # Physical byte shape is (E, N, K // 2), so the descriptor is built at + # byte granularity and the nibbles unpacked in the K-loop. + b_desc = tl.make_tensor_descriptor( + base=b_ptr + off_experts * stride_be, + shape=(N, K // 2), + strides=(stride_bn, stride_bk), + block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K // 2), ) + else: + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + if use_int4_w4a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] // 2) * stride_bk + + offs_bn[None, :] * stride_bn + ) + b_shifter = (offs_k[:, None] % 2) * 4 + elif use_int8_w8a16: + b_ptrs = ( + b_ptr + + off_experts * stride_be + + offs_k[:, None] * stride_bk + + offs_bn[None, :] * stride_bn + ) if not has_zp and use_int4_w4a16: b_zp_num = 8 @@ -219,8 +248,9 @@ def fused_moe_kernel_gptq_awq( accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # Load the next block of A and B, generate a mask by checking the - # K dimension. - + # K dimension. b_scale/b_zp stay on the pointer path even under TD + # (broadcast-mod access, not a dense tile), so k_mask/k_other are + # needed either way. if not block_k_diviable: k_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K k_other = 0.0 @@ -228,14 +258,28 @@ def fused_moe_kernel_gptq_awq( k_mask = None k_other = None - a = tl.load( - a_ptrs, - mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), - other=0.0, - ) - b = tl.load(b_ptrs) - if use_int4_w4a16: - b = (b >> b_shifter) & 0xF + if USE_TD: + # The K tail needs no B mask on either path: b_scale (and b_zp) are + # loaded with k_mask/other=0.0 whenever a tail exists, so OOB B + # dequantizes to exactly 0 regardless of the bytes read. A is zeroed + # independently -- by its mask on the pointer path, by the + # descriptor's zero-fill here. + a = a_desc.gather(gather_idx, k * BLOCK_SIZE_K) + # Interleave before transposing: K // 2 must still be the last axis + # for the nibbles to reconstruct K in the right order. + b_packed = b_desc.load([pid_n * BLOCK_SIZE_N, (k * BLOCK_SIZE_K) // 2]) + b_lo = b_packed & 0xF # even k + b_hi = (b_packed >> 4) & 0xF # odd k + b = tl.interleave(b_lo, b_hi).T + else: + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs) + if use_int4_w4a16: + b = (b >> b_shifter) & 0xF b_scale_ptrs = ( b_scale_ptr @@ -275,12 +319,14 @@ def fused_moe_kernel_gptq_awq( b = ((b.to(tl.float32) - b_zp_num) * b_scale).to(compute_type) accumulator = tl.dot(a, b, acc=accumulator) - # Advance the ptrs to the next K block. - a_ptrs += BLOCK_SIZE_K * stride_ak - if use_int4_w4a16: - b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk - else: - b_ptrs += BLOCK_SIZE_K * stride_bk + if not USE_TD: + # Advance the ptrs to the next K block. TD recomputes absolute + # offsets each iteration instead. + a_ptrs += BLOCK_SIZE_K * stride_ak + if use_int4_w4a16: + b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk + else: + b_ptrs += BLOCK_SIZE_K * stride_bk if MUL_ROUTED_WEIGHT: moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0) @@ -721,6 +767,58 @@ def invoke_fused_moe_wna16_triton_kernel( ) ) + # Same hardware policy as fused_moe_kernel's TD path (moe_use_td_hw_supported: + # XPU, or CUDA sm100+ for tile::gather4); int4_w4a16 only, since the int8 + # branch of this kernel has no TD path. CUDA stays opt-in: only XPU is + # validated. + use_td = resolve_moe_use_td() and moe_use_td_hw_supported() and use_int4_w4a16 + if use_td and config["BLOCK_SIZE_M"] < TD_MIN_GATHER_ROWS: + # tensor_descriptor.gather() asserts at least TD_MIN_GATHER_ROWS rows, so + # a smaller tile aborts the launch. Two config sources can produce one: + # get_default_config's use_moe_wna16_cuda branch, which is what crashed + # --moe-backend triton on B200 via TritonWNA16Experts.apply, and + # override_config, which try_get_optimal_moe_config honours verbatim for + # every caller including fused_experts_impl. + use_td = False + if use_td and M == 1: + # Descriptor setup is a fixed per-launch cost a single row cannot + # amortize: on B70 at m=1 TD costs ~155us extra, while winning 1.1-1.6x + # from a few rows up. GEMM2 still takes TD, as its M is m * top_k. + use_td = False + if ( + use_td + and not current_platform.is_xpu() + and A.size(1) % config["BLOCK_SIZE_K"] != 0 + ): + # Mirrors invoke_fused_moe_triton_kernel's bail-out, which blames a Triton + # codegen bug observed on CUDA ("~74% of output elements wrong") rather + # than a maskable boundary gap -- a claim about the compiler, which this + # kernel's different B-masking does not neutralize. XPU is exempt: + # test_fused_moe_wn16_td_k_tail_matches_pointer covers the tail there. + logger.warning_once( + "Disabling VLLM_TRITON_USE_TD for this MoE launch: K=%d is not a " + "multiple of BLOCK_SIZE_K=%d, which triggers a known Triton " + "tensor-descriptor + tl.dot miscompilation on this platform.", + A.size(1), + config["BLOCK_SIZE_K"], + ) + use_td = False + if use_td and config["BLOCK_SIZE_K"] < 32: + # The packed-byte descriptor's innermost dim is BLOCK_SIZE_K // 2 bytes + # and needs at least 16. get_moe_wna16_block_config only returns 32 or 64, + # but override_config is a public tuning path, so fall back rather than + # assert -- TD is a perf opt-in, never a correctness requirement. + logger.warning_once( + "Disabling the int4 tensor-descriptor path: BLOCK_SIZE_K=%d is " + "below the 32 its packed-byte descriptor needs.", + config["BLOCK_SIZE_K"], + ) + use_td = False + if use_td: + # In-kernel descriptor construction needs a PyTorch-backed scratch + # allocator registered. + set_triton_allocator(A.device) + fused_moe_kernel_gptq_awq[grid]( A, B, @@ -756,6 +854,7 @@ def invoke_fused_moe_wna16_triton_kernel( has_zp=B_zp is not None, use_int4_w4a16=use_int4_w4a16, use_int8_w8a16=use_int8_w8a16, + USE_TD=use_td, **config, ) @@ -796,7 +895,11 @@ def invoke_fused_moe_triton_kernel( is_quantized = B_scale is not None warn_if_moe_use_td_ineffective("TRITON", is_quantized=is_quantized) - # TD path is unvalidated under quantization; fall back to the pointer path. + # This kernel has no TD path for any of its quantized branches -- fp8_w8a8, + # int8_w8a8, and the ungrouped int8_w8a16 that lands here all fall back to + # pointer arithmetic. Grouped WNA16 has one, in fused_moe_kernel_gptq_awq. + # (batched_triton_kernel also runs TD under quantization, via a separate + # gate in fused_batched_moe.py that has no is_quantized check.) use_td = resolve_moe_use_td() and not is_quantized if use_td: # The TD path builds a tensor descriptor inside the kernel, which diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index edb54bcf4080..e11833197032 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -130,14 +130,6 @@ class Mxfp4MoeBackend(Enum): HUMMING = "HUMMING" -# AITER backends group -AITER_BACKENDS = ( - Mxfp4MoeBackend.AITER_MXFP4_BF16, - Mxfp4MoeBackend.AITER_MXFP4_FP8, - Mxfp4MoeBackend.AITER_MXFP4_MXFP4, -) - - # Backends that share the same TRTLLM weight format TRTLLM_BACKENDS = ( Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, @@ -644,20 +636,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( activation: MoEActivation | None = None, ) -> tuple[int, int]: """Round up hidden_size and intermediate_size based on backend requirements.""" - if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SITU: - # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size - # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a - # no-op for K3's shapes) rather than the generic ROCm 256 round-up, - # which would inflate weights and OOM. - intermediate_size = round_up(intermediate_size, 128) - hidden_size = round_up(hidden_size, 128) - elif ( - backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and activation == MoEActivation.SILU - ): - # AITER's A16W4 SiLU kernel handles native dimensions aligned to 128. - intermediate_size = round_up(intermediate_size, 128) - hidden_size = round_up(hidden_size, 128) - elif backend == Mxfp4MoeBackend.EMULATION: + if backend == Mxfp4MoeBackend.EMULATION: # Emulation has no kernel tile; it only needs OCP MX block alignment so the # per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated # by a non-block-aligned TP/DP shard (e.g. 2880 // 4 = 720). @@ -683,8 +662,18 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( intermediate_size = round_up(intermediate_size, 128) hidden_size = round_up(hidden_size, 128) elif current_platform.is_rocm(): - intermediate_size = round_up(intermediate_size, 256) - hidden_size = round_up(hidden_size, 256) + if backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and ( + activation == MoEActivation.SITU or activation == MoEActivation.SILU + ): + # K3's AITER A16W4 SiTU kernel handles K3's native intermediate size + # (moe_intermediate 3072; e.g. 384/partition at TP8). Align to 128 (a + # no-op for K3's shapes) rather than the generic ROCm 256 round-up, + # which would inflate weights and OOM. + intermediate_size = round_up(intermediate_size, 128) + hidden_size = round_up(hidden_size, 128) + else: + intermediate_size = round_up(intermediate_size, 256) + hidden_size = round_up(hidden_size, 256) elif backend == Mxfp4MoeBackend.CPU: # CPU AMX kernel uses BLOCK_N=32, align to 32 intermediate_size = round_up(intermediate_size, 32) @@ -1280,6 +1269,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, + activation: MoEActivation | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -1479,13 +1469,41 @@ def convert_weight_to_mxfp4_moe_kernel_format( import os - from aiter.ops.shuffle import shuffle_scale as _shuf_s - from aiter.ops.shuffle import shuffle_weight as _shuf_w - # 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 + + from vllm._aiter_ops import rocm_aiter_ops + + 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() + w13 = rocm_aiter_ops.shuffle_weight_a16w4( + w13_weight.data.view(fp4_dtype), 16, guinterleave + ) + w2 = rocm_aiter_ops.shuffle_weight_a16w4( + w2_weight.data.view(fp4_dtype), 16, False + ) + w13_scale_raw = w13_weight_scale.data.view(e8m0_dtype) + w2_scale_raw = w2_weight_scale.data.view(e8m0_dtype) + w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( + w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), + num_experts, + guinterleave, + ) + w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1])) + w13.is_shuffled = True + w2.is_shuffled = True + return (w13, w2, w13_scale, w2_scale, w13_bias, w2_bias) + + from aiter.ops.shuffle import shuffle_scale as _shuf_s + from aiter.ops.shuffle import shuffle_weight as _shuf_w + w13_weight = torch.nn.Parameter( _shuf_w( w13_weight.data.view(torch.float4_e2m1fn_x2), 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 129e3b5d5c27..5e950fdf606e 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 @@ -80,9 +80,23 @@ def __init__( # DBO microbatching: one handle slot per micro-batch. self.handles: list[deep_ep.EPHandle | None] = [None, None] + # arange(num_local_experts) + rank_expert_offset. Rank-constant, so it + # is built once per device instead of once per layer per step. + self._global_expert_ids_cache: torch.Tensor | None = None + def num_dispatchers(self) -> int: return self.num_dispatchers_ + def _global_expert_ids(self, num_local: int, device: torch.device) -> torch.Tensor: + ids = self._global_expert_ids_cache + if ids is None or ids.numel() != num_local or ids.device != device: + ids = ( + torch.arange(num_local, dtype=torch.int64, device=device) + + self.rank_expert_offset + ) + self._global_expert_ids_cache = ids + return ids + def output_is_reduced(self) -> bool: return True @@ -196,23 +210,23 @@ def _receiver( else: expert_x, expert_x_scale = recv_x, None + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( + recv_expert_num_tokens, + device=expert_x.device, + ) + if recv_topk_idx is None: # do_expand=True (prefill mode): build topk_ids from # per-expert token counts. total_tokens = sum(recv_expert_num_tokens) if total_tokens > 0: - recv_topk_idx = torch.empty( - total_tokens, - dtype=torch.int64, - device=expert_x.device, + recv_topk_idx = torch.repeat_interleave( + self._global_expert_ids( + len(recv_expert_num_tokens), expert_x.device + ), + expert_tokens_meta.expert_num_tokens, + output_size=total_tokens, ) - offset = 0 - for i, count in enumerate(recv_expert_num_tokens): - if count > 0: - recv_topk_idx[offset : offset + count].fill_( - i + self.rank_expert_offset - ) - offset += count else: recv_topk_idx = torch.empty( 0, @@ -243,11 +257,6 @@ def _receiver( if recv_topk_weights is not None and recv_topk_weights.ndim == 1: recv_topk_weights = recv_topk_weights.unsqueeze(1) - expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list( - recv_expert_num_tokens, - device=expert_x.device, - ) - if not quant_config.is_block_quantized and not defer_input_quant: expert_x_scale = None if expert_x.numel() != 0: diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 87fc3b1c3653..b5557f4b59e9 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -1060,8 +1060,8 @@ def build_expert_params_mapping( if routed_experts_prefix != "": routed_experts_prefix = f"{routed_experts_prefix}." - w13 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w13_" - w2 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w2_" + w13 = f"experts.{lora_base_layer_prefix}{routed_experts_prefix}w13_" + w2 = f"experts.{lora_base_layer_prefix}{routed_experts_prefix}w2_" fused_mapping = [] if include_fused: diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index cce8ccd073fc..f6d14e1e7043 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -591,6 +591,11 @@ def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool: ) +# tensor_descriptor.gather() asserts at least this many rows in the gathered +# tile, so a TD launch needs BLOCK_SIZE_M >= 8. +TD_MIN_GATHER_ROWS = 8 + + def moe_use_td_hw_supported() -> bool: """Whether the current device can run the TD (gather) path of ``fused_moe_kernel`` (ignores the ``VLLM_TRITON_USE_TD`` override). @@ -637,9 +642,11 @@ def warn_if_moe_use_td_ineffective( """One-shot warning when ``VLLM_TRITON_USE_TD`` is set but ignored. Fires when the user set the env explicitly and either (a) the active - MoE backend is not the fused Triton kernel, or (b) the model is - quantized (the TD path falls back to the pointer path under any - quantization). + MoE backend is not the fused Triton kernel, or (b) the weights are + quantized in a scheme ``fused_moe_kernel`` has no TD path for -- which is + every scheme reaching it, including the ungrouped ``int8_w8a16`` produced + by ``oracle/int8.py``. Grouped WNA16 has a TD path, but it lives in + ``fused_moe_kernel_gptq_awq`` and is warned about by its own launcher. """ global _warned_moe_use_td_ineffective if _warned_moe_use_td_ineffective: @@ -656,9 +663,8 @@ def warn_if_moe_use_td_ineffective( ) else: reason = ( - "the model uses quantized MoE weights; the TD path is " - "currently restricted to non-quantized weights and falls " - "back to the pointer path" + "this MoE layer's quantization scheme has no tensor-descriptor " + "path in fused_moe_kernel and falls back to the pointer path" ) logger.warning( "VLLM_TRITON_USE_TD is set to %s but %s.", diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index d06916f697c7..ef350237dc70 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -71,10 +71,12 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: page_size_padded=page_size_padded, mamba_type=self.mamba_type, mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, + # RecoverSSM verifies the whole window off one checkpoint, so it + # never writes the baseline's per-draft-token state slots. num_speculative_blocks=( - vllm_config.speculative_config.num_speculative_tokens - if vllm_config.speculative_config - else 0 + 0 + if vllm_config.cache_config.use_kda_recoverssm + else vllm_config.num_speculative_tokens ), ) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index f3ebcd63255b..27dc4d90d7e0 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen3-Next/Qwen3.5 model.""" +import os from typing import Literal import torch from einops import rearrange from torch import nn +from vllm import _custom_ops as ops from vllm import envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import ( @@ -84,6 +86,9 @@ logger = init_logger(__name__) +MAX_FUSED_GDN_MTP_TOKENS = 8 +FUSED_GDN_STATE_DTYPES = (torch.float32, torch.bfloat16) + def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, @@ -484,12 +489,50 @@ def __init__( self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) + self.gdn_decode_kernel = envs.VLLM_GDN_DECODE_KERNEL.strip().lower() + if self.gdn_decode_kernel == "cuda": + reason = self._fused_gdn_decode_unsupported_reason(vllm_config) + if reason is not None: + if "VLLM_GDN_DECODE_KERNEL" in os.environ: + raise ValueError( + f"VLLM_GDN_DECODE_KERNEL=cuda is not supported: {reason}" + ) + logger.info_once( + "Falling back to the Triton GDN decode path: %s", reason + ) + self.gdn_decode_kernel = "triton" + self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" + logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel) compilation_config = get_current_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 + def _fused_gdn_decode_unsupported_reason( + self, vllm_config: VllmConfig + ) -> str | None: + conv_state_dtype, recurrent_state_dtype = self.get_state_dtype() + if ( + self.gqa_interleaved_layout + or self.head_k_dim != 128 + or self.head_v_dim != 128 + or self.norm.activation != "silu" + or vllm_config.model_config.dtype != torch.bfloat16 + or conv_state_dtype != torch.bfloat16 + or recurrent_state_dtype not in FUSED_GDN_STATE_DTYPES + or not current_platform.has_device_capability(80) + ): + return ( + "the fused CUDA kernel requires a BF16 GDN model with " + "K=V=128, SiLU gating, non-interleaved GQA layout, BF16 " + "convolution cache, BF16 or FP32 recurrent state, and a " + "GPU with compute capability 8.0+" + ) + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): + return "torch.ops._C.fused_gdn_decode_post_conv_mtp is not built" + return None + def create_qkvz_proj( self, hidden_size: int, @@ -846,6 +889,26 @@ def forward_cuda( mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) ba, _ = self.in_proj_ba(hidden_states) + use_fused_gdn_decode = ( + self.enable_fused_gdn_decode + and hidden_states.dtype == torch.bfloat16 + and self.norm.weight.dtype in (torch.bfloat16, torch.float32) + ) + if use_fused_gdn_decode: + core_attn_out = torch.zeros( + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + torch.ops.vllm.qwen_gdn_attention_core_fused_norm_packed( + mixed_qkvz, + ba, + core_attn_out, + layer_name=_encode_layer_name(self.prefix), + ) + output, _ = self.out_proj(core_attn_out.flatten(-2)) + return output + if self.gqa_interleaved_layout: # Qwen3-Next: unpack the interleaved GQA layout query, key, value, z, b, a = self.fix_query_key_value_ordering( @@ -862,8 +925,6 @@ def forward_cuda( mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1) z = z.reshape(z.size(0), -1, self.head_v_dim) b, a = self.split_ba(ba) - b = b.contiguous() - a = a.contiguous() # ============================================================ # Part 2: Core Attention (Custom Op) @@ -878,8 +939,8 @@ def forward_cuda( torch.ops.vllm.qwen_gdn_attention_core( mixed_qkv, - b, - a, + b.contiguous(), + a.contiguous(), core_attn_out, layer_name=_encode_layer_name(self.prefix), ) @@ -1622,6 +1683,211 @@ def _forward_core_decode_non_spec( ) return + def _forward_core_decode_spec_fused_norm( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + output_gate: torch.Tensor, + core_attn_out: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + ) -> None: + state_indices = attn_metadata.spec_state_indices_tensor + cu_seqlens = attn_metadata.spec_query_start_loc + num_accepted_tokens = attn_metadata.num_accepted_tokens + assert state_indices is not None + assert cu_seqlens is not None + assert num_accepted_tokens is not None + + num_requests = attn_metadata.num_spec_decodes + num_actual_tokens = attn_metadata.num_actual_tokens + conv_state = ( + self.kv_cache[0] + if is_conv_state_dim_first() + else self.kv_cache[0].transpose(-1, -2) + ) + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + mixed_qkv = causal_conv1d_update( + mixed_qkv[:num_actual_tokens], + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=state_indices[:num_requests, 0], + num_accepted_tokens=num_accepted_tokens[:num_requests], + query_start_loc=cu_seqlens[: num_requests + 1], + max_query_len=state_indices.size(1), + validate_data=False, + ) + self._forward_core_decode_spec_post_conv_fused_norm( + mixed_qkv=mixed_qkv, + b=b[:num_actual_tokens], + a=a[:num_actual_tokens], + output_gate=output_gate[:num_actual_tokens], + core_attn_out=core_attn_out[:num_actual_tokens], + attn_metadata=attn_metadata, + ) + + def _forward_core_decode_spec_post_conv_fused_norm( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + output_gate: torch.Tensor, + core_attn_out: torch.Tensor, + attn_metadata: GDNAttentionMetadata, + ) -> None: + state_indices = attn_metadata.spec_state_indices_tensor + cu_seqlens = attn_metadata.spec_query_start_loc + num_accepted_tokens = attn_metadata.num_accepted_tokens + assert state_indices is not None + assert cu_seqlens is not None + assert num_accepted_tokens is not None + + num_requests = attn_metadata.num_spec_decodes + ops.fused_gdn_decode_post_conv_mtp( + mixed_qkv=mixed_qkv, + a=a, + b=b, + A_log=self.A_log, + dt_bias=self.dt_bias, + state_indices=state_indices[:num_requests], + cu_seqlens=cu_seqlens[: num_requests + 1], + num_accepted_tokens=num_accepted_tokens[:num_requests], + state=self.kv_cache[1], + output_gate=output_gate, + norm_weight=self.norm.weight, + out=core_attn_out, + scale=self.head_k_dim**-0.5, + norm_eps=self.layer_norm_epsilon, + ) + + def _forward_core_fused_norm_packed( + self, + mixed_qkvz: torch.Tensor, + ba: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size + if attn_metadata_raw is None: + self._warmup_prefill_kernels(mixed_qkvz[:, :qkv_size], 0) + return + + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index] + assert isinstance(attn_metadata, GDNAttentionMetadata) + mixed_qkv, output_gate_flat = mixed_qkvz.split( + [qkv_size, self.value_dim // self.tp_size], dim=-1 + ) + output_gate = output_gate_flat.reshape( + output_gate_flat.size(0), -1, self.head_v_dim + ) + b, a = self.split_ba(ba) + self._forward_core_fused_norm( + mixed_qkv=mixed_qkv, + b=b, + a=a, + output_gate=output_gate, + core_attn_out=core_attn_out, + ) + + def _can_use_fused_gdn_mtp_decode( + self, attn_metadata: GDNAttentionMetadata + ) -> bool: + state_indices = attn_metadata.spec_state_indices_tensor + return ( + attn_metadata.spec_sequence_masks is not None + and attn_metadata.num_decodes == 0 + and attn_metadata.num_spec_decodes > 0 + and self.kv_cache[1].dtype in FUSED_GDN_STATE_DTYPES + and self.gdn_decode_kernel == "cuda" + and self.num_v_heads == 8 * self.num_k_heads + and state_indices is not None + and state_indices.size(1) <= MAX_FUSED_GDN_MTP_TOKENS + and hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp") + ) + + def _rms_norm_gated_cuda( + self, + x: torch.Tensor, + output_gate: torch.Tensor, + out: torch.Tensor, + ) -> None: + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( + layer_norm_fwd, + ) + + x_shape = x.shape + assert output_gate.shape == x_shape + assert out.shape == x_shape + x_2d = x.reshape(-1, x_shape[-1]) + output_gate_2d = output_gate.reshape(-1, x_shape[-1]) + out_2d = out.reshape(-1, x_shape[-1]) + assert x_2d.stride(-1) == 1 + assert output_gate_2d.stride(-1) == 1 + assert out_2d.stride(-1) == 1 + layer_norm_fwd( + x_2d, + self.norm.weight.contiguous(), + self.norm.bias, + self.norm.eps, + z=output_gate_2d, + out=out_2d, + group_size=( + x_shape[-1] if self.norm.group_size is None else self.norm.group_size + ), + norm_before_gate=self.norm.norm_before_gate, + is_rms_norm=True, + activation=self.norm.activation, + ) + + def _forward_core_fused_norm( + self, + mixed_qkv: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + output_gate: torch.Tensor, + core_attn_out: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata_raw = forward_context.attn_metadata + if attn_metadata_raw is None: + self._warmup_prefill_kernels(mixed_qkv, 0) + return + + assert isinstance(attn_metadata_raw, dict) + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index] + assert isinstance(attn_metadata, GDNAttentionMetadata) + if ( + self._can_use_fused_gdn_mtp_decode(attn_metadata) + and attn_metadata.num_prefills == 0 + ): + self._forward_core_decode_spec_fused_norm( + mixed_qkv=mixed_qkv, + b=b, + a=a, + output_gate=output_gate, + core_attn_out=core_attn_out, + attn_metadata=attn_metadata, + ) + return + self._forward_core( + mixed_qkv=mixed_qkv, + b=b.contiguous(), + a=a.contiguous(), + core_attn_out=core_attn_out, + ) + num_actual_tokens = attn_metadata.num_actual_tokens + self._rms_norm_gated_cuda( + core_attn_out[:num_actual_tokens], + output_gate[:num_actual_tokens], + core_attn_out[:num_actual_tokens], + ) + def qwen_gdn_attention_core( qkv_or_qkvz: torch.Tensor, @@ -1683,6 +1949,39 @@ def gdn_attention_core_fake( ) +def qwen_gdn_attention_core_fused_norm_packed( + mixed_qkvz: torch.Tensor, + ba: torch.Tensor, + core_attn_out: torch.Tensor, + layer_name: LayerNameType, +) -> None: + layer_name = _resolve_layer_name(layer_name) + forward_context: ForwardContext = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + self._forward_core_fused_norm_packed( + mixed_qkvz=mixed_qkvz, + ba=ba, + core_attn_out=core_attn_out, + ) + + +def gdn_attention_core_fused_norm_packed_fake( + mixed_qkvz: torch.Tensor, + ba: torch.Tensor, + core_attn_out: torch.Tensor, + layer_name: LayerNameType, +) -> None: + return + + +direct_register_custom_op( + op_name="qwen_gdn_attention_core_fused_norm_packed", + op_func=qwen_gdn_attention_core_fused_norm_packed, + mutates_args=["core_attn_out"], + fake_impl=gdn_attention_core_fused_norm_packed_fake, +) + + @triton.jit def fused_gdn_gating_kernel( g, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 7538a2b6b49d..14b112138998 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -48,6 +48,7 @@ from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -904,17 +905,21 @@ def conv_ssm_forward( # then chunk_stride = 2 chunk_stride = mamba_block_size // chunk_size + # The per-sequence loop below uses these as Python scalars. + # TODO avoid sync here? + with gpu_sync_allowed(): + block_idx_first_cpu = block_idx_first_scheduled_token_p.tolist() + block_idx_last_cpu = block_idx_last_scheduled_token_p.tolist() + num_computed_tokens_p_cpu = num_computed_tokens_p.tolist() + last_chunk_indices_p_cpu = last_chunk_indices_p.tolist() + # Save state for sequences with more than just final state for seq_idx in range(num_prefills): # Block index for the first scheduled token - block_idx_first_scheduled_token = block_idx_first_scheduled_token_p[ - seq_idx - ] + block_idx_first_scheduled_token = block_idx_first_cpu[seq_idx] # Block index for the last scheduled token - block_idx_last_scheduled_token = block_idx_last_scheduled_token_p[ - seq_idx - ] + block_idx_last_scheduled_token = block_idx_last_cpu[seq_idx] # Number of blocks that need to be written n_blocks_to_fill = ( @@ -935,7 +940,7 @@ def conv_ssm_forward( if seq_idx == 0: first_chunk = 0 else: - first_chunk = 1 + last_chunk_indices_p[seq_idx - 1] + first_chunk = 1 + last_chunk_indices_p_cpu[seq_idx - 1] # First chunk that is aligned on the mamba block boundary first_aligned_chunk = first_chunk + chunk_stride - 1 @@ -943,7 +948,7 @@ def conv_ssm_forward( # Calculate the number of computed tokens that were not # already cached num_unaligned_computed_tokens = ( - num_computed_tokens_p[seq_idx] % mamba_block_size + num_computed_tokens_p_cpu[seq_idx] % mamba_block_size ) if num_unaligned_computed_tokens > 0: diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 50f6e059c0f4..18a73c6554e6 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -136,6 +136,15 @@ def kda_state_dtype( state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype) return (state_dtype, torch.float32) + @classmethod + def append_kda_recoverssm_record( + cls, + base_dtypes: tuple[torch.dtype, ...], + model_dtype: ModelDType | torch.dtype, + ) -> tuple[torch.dtype, ...]: + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) + return (*base_dtypes, torch.float32, activation_dtype) + class MambaStateShapeCalculator: @classmethod @@ -293,6 +302,27 @@ def kda_state_shape( recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim) return (conv_state_shape, recurrent_state_shape) + @classmethod + def append_kda_recoverssm_record( + cls, + base_shapes: tuple[tuple[int, int], tuple[int, int, int]], + num_heads: int, + head_dim: int, + tp_world_size: int, + spec_query_len: int, + ) -> tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int, int], + ]: + local_num_heads = divide(num_heads, tp_world_size) + return ( + *base_shapes, + (local_num_heads, spec_query_len, head_dim), + (local_num_heads, spec_query_len, 2 * head_dim), + ) + @dataclass class MambaCopySpec: diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..edc77b17c5d7 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..6e01b1b9a6b6 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=128,dstate=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,51 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 2 + }, + "256": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "512": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "1024": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "8192": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json new file mode 100644 index 000000000000..8735bb1c6fb7 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float16.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 64, + "num_warps": 8 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 2 + }, + "32768": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "131072": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json new file mode 100644 index 000000000000..a382f024f96d --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,cache_dtype=float32.json @@ -0,0 +1,63 @@ +{ + "triton_version": "3.7.1", + "8": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "16": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "64": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "128": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "256": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "512": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "1024": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "2048": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "4096": { + "BLOCK_SIZE_M": 4, + "num_warps": 1 + }, + "8192": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "16384": { + "BLOCK_SIZE_M": 8, + "num_warps": 1 + }, + "32768": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "65536": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + }, + "131072": { + "BLOCK_SIZE_M": 16, + "num_warps": 8 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py index b952e3ebbce0..02c89823ee06 100644 --- a/vllm/model_executor/layers/mamba/ops/gather_initial_states.py +++ b/vllm/model_executor/layers/mamba/ops/gather_initial_states.py @@ -49,7 +49,7 @@ def gather_initial_states( ) -> torch.Tensor: """Gather dense state rows, replacing uninitialized rows with zeros.""" assert state.ndim >= 2 - assert state.is_cuda + assert state.is_cuda or state.is_xpu assert indices.ndim == 1 and has_initial_state.ndim == 1 assert indices.shape == has_initial_state.shape assert indices.device == state.device diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index ef81b0b876df..713bd7bd5e8d 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -358,31 +358,6 @@ def get_quant_method( return None - @classmethod - def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]): - # Extract data from quant config. - quant_method = quant_config.get("quant_method", "").lower() - num_bits = quant_config.get("bits") - group_size = quant_config.get("group_size") - zero_point = quant_config.get("zero_point") - - if not (current_platform.is_cuda_alike() or current_platform.is_cpu()): - return False - - if quant_method != "awq": - return False - - # If we cannot find the info needed in the config, cannot convert. - if num_bits is None or group_size is None or zero_point is None: - return False - - if num_bits not in cls.TYPE_MAP: - return False - - return check_marlin_supported( - quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point - ) - def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): if self.modules_to_not_convert: self.modules_to_not_convert = hf_to_vllm_mapper.apply_list( diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index ad7aea175def..c76259616cd9 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -258,19 +258,3 @@ def maybe_update_config( # noqa: B027 # TODO: revision is never passed currently in vllm.py, # but is used in subclasses, should we remove this parameter? pass - - def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool: - """ - Determine if mxfp4 quantization will be used for this config. - - This allows hidden_size rounding to happen before moe_config creation - without needing to instantiate quant_method first. - - Args: - prefix: The layer prefix/name in the model - layer: The layer module - - Returns: - True if this config uses MXFP4 quantization, False otherwise - """ - return False diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py index 22c3539e9aec..2c121c6edfe9 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py @@ -32,7 +32,6 @@ W4A8_SUPPORTED_TYPES_MAP = { 4: scalar_types.int4, } -W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys()) class CompressedTensorsW4A8Fp8(CompressedTensorsScheme): diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py index 77933ea2c736..a26a35f4e313 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py @@ -27,7 +27,6 @@ W4A8_SUPPORTED_TYPES_MAP = { 4: scalar_types.int4, } -W4A8_SUPPORTED_BITS = list(W4A8_SUPPORTED_TYPES_MAP.keys()) class CompressedTensorsW4A8Int(CompressedTensorsScheme): diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index b814205922c6..ea00c0f3c408 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import math from collections.abc import Callable from typing import TYPE_CHECKING, Any @@ -61,11 +60,6 @@ ) -def prepare_padded_shape(shape, x): - padded_shape = math.ceil(shape / x) * x - return padded_shape, padded_shape - shape - - def prepare_param(tensor, name, extra_attrs): extra_attrs = extra_attrs.copy() scale_type = extra_attrs.pop("scale_type", None) @@ -129,22 +123,6 @@ def prepare_moe_param(tensor: torch.Tensor, name: str, extra_attrs: dict[str, An return param -def may_pad_loaded_weight(param, loaded_weight): - pad_shape = getattr(param, "pad_shape", None) - if pad_shape is None: - return loaded_weight - value = 1 if loaded_weight.dtype == torch.float8_e8m0fnu else 0 - padding = [] - for x in pad_shape[::-1][: loaded_weight.ndim]: - padding += [0, x] - loaded_weight = torch.nn.functional.pad( - input=loaded_weight, - pad=padding, - value=value, - ) - return loaded_weight - - def compressed_tensors_get_config(config: dict[str, Any], key: str): assert key in ["weights", "input_activations"] target_group_config = None diff --git a/vllm/model_executor/layers/quantization/inc/inc.py b/vllm/model_executor/layers/quantization/inc/inc.py index 219fc2b04692..ff8eae3a4649 100644 --- a/vllm/model_executor/layers/quantization/inc/inc.py +++ b/vllm/model_executor/layers/quantization/inc/inc.py @@ -53,7 +53,6 @@ class INCConfig(QuantizationConfig): MXFP8_GROUP_SIZE = 32 MXFP8_DATA_TYPE = "mx_fp" MXFP8_PACKING_FORMAT = "auto_round:llm_compressor" - MXFP8_SUPPORTED_ACT_DTYPES = {"mx_fp", "mx_fp_rceil"} def __init__( self, diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py index 5c99fd98b54d..42f053b471a9 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py @@ -458,7 +458,3 @@ def apply_weights( layer.ark_scale_type, not self.sym, ) - - -class INCXPUW4A16LinearScheme(INCXPULinearMethod): - pass diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index db98b76b3f5d..e24c5656add3 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -112,7 +112,6 @@ # MIXED_PRECISION, "MIXED_PRECISION", ] -KV_CACHE_QUANT_ALGOS = ["FP8", "NVFP4"] class ModelOptKVCacheMethod(BaseKVCacheMethod): diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index de659e4f5084..6506c1fce1df 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import os - import torch from vllm.logger import init_logger @@ -22,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, - backend_to_kernel_cls, convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, convert_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, @@ -104,10 +101,6 @@ def get_quant_method( ) return None - def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool: - """MXFP4 config always uses MXFP4 quantization.""" - return True - class GptOssMxfp4Config(Mxfp4Config): """MXFP4 config for GPT-OSS checkpoints. @@ -476,51 +469,14 @@ def apply_monolithic( ) -def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool: - """Whether Kimi-K3's SiTU MXFP4 MoE should use the AITER A16W4 kernel. - - K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic - MXFP4 backend selector does not cover; route it to AITER on gfx950. - """ - if not current_platform.is_rocm(): - return False - from vllm._aiter_ops import rocm_aiter_ops - from vllm.model_executor.layers.fused_moe.activation import MoEActivation - from vllm.platforms.rocm import on_gfx950 - - return ( - rocm_aiter_ops.is_fused_moe_enabled() - and on_gfx950() - and moe.activation == MoEActivation.SITU - and moe.activation_situ_linear_beta is not None - and rocm_aiter_ops.get_aiter_activation_type("situ") is not None - ) - - class Mxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" def __init__(self, moe: FusedMoEConfig): super().__init__(moe) + self.weight_dtype = "mxfp4" - self.is_k3_situ_aiter = _use_k3_situ_aiter(moe) - self.experts_cls: type[mk.FusedMoEExperts] | None - if self.is_k3_situ_aiter: - self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 - self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] - logger.info_once("Using AITER_MXFP4_BF16 for Kimi-K3 SiTU MXFP4 MoE.") - from vllm._aiter_ops import rocm_aiter_ops - - if rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled(): - # AITER keeps bf16 activations below this token count, which - # would not match the fp8 a8w4 kernels the interleaved SiTU - # path is tuned for. The a16w4 path never reads it. - # TODO: Remove once AITER takes this as a kernel argument. - os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0" - else: - self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend( - moe - ) + self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe) self.max_capture_size = moe.max_capture_size @@ -718,62 +674,62 @@ def _setup_kernel( hidden_size = self.hidden_size sf_block_size = 32 - # Shape assertions - assert ( - w13.dim() == 3 - and w13.shape[0] == num_experts - and w13.shape[1] == intermediate_size * self.moe.w13_num_shards - and w13.shape[2] == hidden_size // 2 - ) - assert ( - w13_scale.dim() == 3 - and w13_scale.shape[0] == num_experts - and w13_scale.shape[1] == intermediate_size * self.moe.w13_num_shards - and w13_scale.shape[2] == hidden_size // sf_block_size - ) - assert ( - w2.dim() == 3 - and w2.shape[0] == num_experts - and w2.shape[1] == hidden_size - and w2.shape[2] == intermediate_size // 2 - ) - assert ( - w2_scale.dim() == 3 - and w2_scale.shape[1] == hidden_size - and w2_scale.shape[2] == intermediate_size // sf_block_size - ) - if w13_bias is not None: + # Shape assertions — skipped for SITU since its kernel handles native + # (non-256-aligned) intermediate sizes without prior round-up. + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + + if self.moe.activation != MoEActivation.SITU: assert ( - w13_bias.dim() == 2 - and w13_bias.shape[0] == num_experts - and w13_bias.shape[1] == intermediate_size * self.moe.w13_num_shards + w13.dim() == 3 + and w13.shape[0] == num_experts + and w13.shape[1] == intermediate_size * self.moe.w13_num_shards + and w13.shape[2] == hidden_size // 2 ) - if w2_bias is not None: assert ( - w2_bias.dim() == 2 - and w2_bias.shape[0] == num_experts - and w2_bias.shape[1] == hidden_size + w13_scale.dim() == 3 + and w13_scale.shape[0] == num_experts + and w13_scale.shape[1] == intermediate_size * self.moe.w13_num_shards + and w13_scale.shape[2] == hidden_size // sf_block_size ) - - # Convert weights to kernel format - if self.is_k3_situ_aiter: - w13, w2, w13_scale, w2_scale = ( - self._convert_k3_situ_weight_to_kernel_format(layer) + assert ( + w2.dim() == 3 + and w2.shape[0] == num_experts + and w2.shape[1] == hidden_size + and w2.shape[2] == intermediate_size // 2 ) - else: - w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( - convert_weight_to_mxfp4_moe_kernel_format( - mxfp4_backend=self.mxfp4_backend, - layer=layer, - w13_weight=w13, - w2_weight=w2, - w13_weight_scale=w13_scale, - w2_weight_scale=w2_scale, - w13_bias=w13_bias, - w2_bias=w2_bias, - _cache_permute_indices=self._cache_permute_indices, + assert ( + w2_scale.dim() == 3 + and w2_scale.shape[1] == hidden_size + and w2_scale.shape[2] == intermediate_size // sf_block_size + ) + if w13_bias is not None: + assert ( + w13_bias.dim() == 2 + and w13_bias.shape[0] == num_experts + and w13_bias.shape[1] == intermediate_size * self.moe.w13_num_shards ) + if w2_bias is not None: + assert ( + w2_bias.dim() == 2 + and w2_bias.shape[0] == num_experts + and w2_bias.shape[1] == hidden_size + ) + + # Convert weights to kernel format + w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( + convert_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=w13, + w2_weight=w2, + w13_weight_scale=w13_scale, + w2_weight_scale=w2_scale, + w13_bias=w13_bias, + w2_bias=w2_bias, + _cache_permute_indices=self._cache_permute_indices, + activation=self.moe.activation, ) + ) # For TRITON backends, weights are wrapped tensors from triton_kernels # that don't support .detach(). Manually assign parameters. @@ -814,42 +770,6 @@ def _setup_kernel( routing_tables=layer._expert_routing_tables(), ) - def _convert_k3_situ_weight_to_kernel_format( - self, layer: RoutedExperts - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - # K3's AITER A16W4 kernel wants the separated ([gate_all, up_all]) - # stage-1 layout, unlike the interleaved gpt-oss/DeepSeek path in - # convert_weight_to_mxfp4_moe_kernel_format. Preshuffle once here. - from aiter.utility.fp4_utils import e8m0_shuffle - - from vllm._aiter_ops import rocm_aiter_ops - - fp4_dtype = torch.float4_e2m1fn_x2 - e8m0_dtype = torch.float8_e8m0fnu - num_experts = layer.w13_weight.shape[0] - - # a8w4 (VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1) uses the gate/up- - # interleaved (_gui_) fp8 flydsl kernels, which need w13 weight+scale - # in interleave layout. Default a16w4 keeps the separated layout. - guinterleave = rocm_aiter_ops.is_fused_moe_situv2_a8w4_enabled() - w13 = rocm_aiter_ops.shuffle_weight_a16w4( - layer.w13_weight.data.view(fp4_dtype), 16, guinterleave - ) - w2 = rocm_aiter_ops.shuffle_weight_a16w4( - layer.w2_weight.data.view(fp4_dtype), 16, False - ) - w13_scale_raw = layer.w13_weight_scale.data.view(e8m0_dtype) - w2_scale_raw = layer.w2_weight_scale.data.view(e8m0_dtype) - w13_scale = rocm_aiter_ops.shuffle_scale_a16w4( - w13_scale_raw.view(-1, w13_scale_raw.shape[-1]), num_experts, guinterleave - ) - w2_scale = e8m0_shuffle(w2_scale_raw.view(-1, w2_scale_raw.shape[-1])) - - w13.is_shuffled = True - w2.is_shuffled = True - - return w13, w2, w13_scale, w2_scale - def process_weights_after_loading(self, layer): w13 = layer.w13_weight w2 = layer.w2_weight diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index 7d3b0c4f931c..9a65432c5c42 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -576,22 +576,6 @@ def _is_w_ocp_mx_a_x( return True - def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool: - """ - For Quark, determine if it's OCP MXFP4 by checking config directly. - This allows hidden_size rounding to happen before moe_config creation. - """ - layer_quant_config = self._find_matched_config(prefix, layer) - weight_config = layer_quant_config.get("weight") - input_config = layer_quant_config.get("input_tensors") - - return ( - self._is_w_ocp_mx_a_x(weight_config, input_config) - and weight_config is not None - and weight_config.get("dtype") == "fp4" - and getattr(torch, "float4_e2m1fn_x2", None) is not None - ) - def _find_matched_config( self, layer_name: str, module: torch.nn.Module ) -> dict[str, Any]: diff --git a/vllm/model_executor/layers/quantization/torchao.py b/vllm/model_executor/layers/quantization/torchao.py index 15399cfd39b4..8862578d6f7e 100644 --- a/vllm/model_executor/layers/quantization/torchao.py +++ b/vllm/model_executor/layers/quantization/torchao.py @@ -283,9 +283,6 @@ def get_quant_method( return TorchAOLinearMethod(self) - def get_scaled_act_names(self) -> list[str]: - return [] - def torchao_quantize_param_data( param: torch.Tensor, torchao_config: Any diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index 4f624cf49630..db95d3588a2a 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -2,17 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from https://github.com/sgl-project/sglang/blob/4cb53ecd0cffceb6dee5c011a58f65997a86f151/python/sglang/srt/layers/quantization/int8_kernel.py -import functools -import json import logging -import os -from typing import Any import torch from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.platform_utils import get_device_name_as_file_name logger = logging.getLogger(__name__) @@ -228,221 +223,3 @@ def per_token_group_quant_int8( ) return x_q, x_s - - -@triton.jit -def _w8a8_block_int8_matmul( - # Pointers to inputs and output - A, - B, - C, - As, - Bs, - # Shape for matmul - M, - N, - K, - # Block size for block-wise quantization - group_n, - group_k, - # Stride for inputs and output - stride_am, - stride_ak, - stride_bk, - stride_bn, - stride_cm, - stride_cn, - stride_As_m, - stride_As_k, - stride_Bs_k, - stride_Bs_n, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, -): - """Triton-accelerated function used to perform linear operations (dot - product) on input tensors `A` and `B` with block-wise quantization, and - store the result in output tensor `C`. - """ - - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + (pid % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - As_ptrs = As + offs_am * stride_As_m - offs_bsn = offs_bn // group_n - Bs_ptrs = Bs + offs_bsn * stride_Bs_n - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) - - k_start = k * BLOCK_SIZE_K - offs_ks = k_start // group_k - a_s = tl.load(As_ptrs + offs_ks * stride_As_k) - b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) - - accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :] - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - - if C.dtype.element_ty == tl.bfloat16: - c = accumulator.to(tl.bfloat16) - elif C.dtype.element_ty == tl.float16: - c = accumulator.to(tl.float16) - else: - c = accumulator.to(tl.float32) - - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=c_mask) - - -@functools.lru_cache -def get_w8a8_block_int8_configs( - N: int, K: int, block_n: int, block_k: int -) -> dict[int, Any] | None: - """ - Return optimized configurations for the w8a8 block fp8 kernel. - - The return value will be a dictionary that maps an irregular grid of - batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the - kernel on a given batch size bs, the closest batch size in the grid should - be picked and the associated configuration chosen to invoke the kernel. - """ - - # First look up if an optimized configuration is available in the configs - # directory - device_name = get_device_name_as_file_name() - json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" # noqa: E501 - - config_file_path = os.path.join( - os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name - ) - if os.path.exists(config_file_path): - with open(config_file_path) as f: - logger.info( - "Using configuration from %s for W8A8 Block INT8 kernel.", - config_file_path, - ) - # If a configuration has been found, return it - return {int(key): val for key, val in json.load(f).items()} - - # If no optimized configuration is available, we will use the default - # configuration - logger.warning( - ( - "Using default W8A8 Block INT8 kernel config. Performance might " - "be sub-optimal! Config file not found at %s" - ), - config_file_path, - ) - return None - - -def w8a8_block_int8_matmul( - A: torch.Tensor, - B: torch.Tensor, - As: torch.Tensor, - Bs: torch.Tensor, - block_size: list[int], - output_dtype: torch.dtype = torch.float16, -) -> torch.Tensor: - """This function performs matrix multiplication with block-wise - quantization. - - It takes two input tensors `A` and `B` with scales `As` and `Bs`. - The output is returned in the specified `output_dtype`. - - Args: - A: The input tensor, e.g., activation. - B: The input tensor, e.g., weight. - As: The per-token-group quantization scale for `A`. - Bs: The per-block quantization scale for `B`. - block_size: The block size for per-block quantization. It should be - 2-dim, e.g., [128, 128]. - output_dtype: The dtype of the returned tensor. - - Returns: - torch.Tensor: The result of matmul. - """ - assert len(block_size) == 2 - block_n, block_k = block_size[0], block_size[1] - - assert A.shape[-1] == B.shape[-1] - assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() - assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] - M = A.numel() // A.shape[-1] - - assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 - N, K = B.shape - assert triton.cdiv(N, block_n) == Bs.shape[0] - assert triton.cdiv(K, block_k) == Bs.shape[1] - - C_shape = A.shape[:-1] + (N,) - C = A.new_empty(C_shape, dtype=output_dtype) - - configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1]) - if configs: - # If an optimal configuration map has been found, look up the - # optimal config - config = configs[min(configs.keys(), key=lambda x: abs(x - M))] - else: - # Default config - # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1] - config = { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": block_size[0], - "BLOCK_SIZE_K": block_size[1], - "GROUP_SIZE_M": 32, - "num_warps": 4, - "num_stages": 3, - } - - def grid(META): - return ( - triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - ) - - _w8a8_block_int8_matmul[grid]( - A, - B, - C, - As, - Bs, - M, - N, - K, - block_n, - block_k, - A.stride(-2), - A.stride(-1), - B.stride(1), - B.stride(0), - C.stride(-2), - C.stride(-1), - As.stride(-2), - As.stride(-1), - Bs.stride(1), - Bs.stride(0), - **config, - ) - - return C diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 65348a822214..2d0a4fa48c0d 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -12,13 +12,6 @@ logger = init_logger(__name__) -# CK's pre-compiled MXFP4 MoE GEMM kernel instances require the -# intermediate_size (after TP split) to be a multiple of this value. -# This arises from FP4 packing (2 values per byte) combined with CK -# tile size constraints. When violated, AITER raises: -# "device_gemm ... does not support this GEMM problem". -CK_MXFP4_MOE_DIM_ALIGNMENT = 256 - def should_use_cdna4_mx_scale_swizzle() -> bool: """Whether to use the CDNA4 swizzled scale layout for mxfp4 on gfx950. diff --git a/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py b/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py index a9157cbfb08b..0d56f179b93f 100644 --- a/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py +++ b/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py @@ -16,7 +16,6 @@ "mxfp8_e5m2", "mxint8", } -SUPPORTED_OCP_MX_DTYPES = {"mxfp4", "mxfp6_e3m2", "mxfp6_e2m3"} class OCP_MX_Scheme(str, Enum): diff --git a/vllm/model_executor/layers/rotary_embedding/mrope.py b/vllm/model_executor/layers/rotary_embedding/mrope.py index 29ce9e5000d9..4da83501713c 100644 --- a/vllm/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm/model_executor/layers/rotary_embedding/mrope.py @@ -226,12 +226,14 @@ def triton_mrope( def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.Tensor: """Apply interleaved MRoPE to 3D rotary embeddings. Reorganizes frequency layout from chunked [TTT...HHH...WWW] to - interleaved [THTHWHTHW...TT], preserving frequency continuity. + interleaved [THWTHWTHW...TT], preserving frequency continuity. """ - x_t = x[0].clone() - x_t[..., 1 : mrope_section[1] * 3 : 3] = x[1, ..., 1 : mrope_section[1] * 3 : 3] - x_t[..., 2 : mrope_section[2] * 3 : 3] = x[2, ..., 2 : mrope_section[2] * 3 : 3] - return x_t + channels = torch.arange(x.shape[-1], device=x.device) + is_height = (channels % 3 == 1) & (channels < mrope_section[1] * 3) + is_width = (channels % 3 == 2) & (channels < mrope_section[2] * 3) + + result = torch.where(is_height, x[1], x[0]) + return torch.where(is_width, x[2], result) class MRotaryEmbedding(RotaryEmbeddingBase): diff --git a/vllm/model_executor/models/arctic.py b/vllm/model_executor/models/arctic.py index b7e8e6796f3e..ccb3f1976cf6 100644 --- a/vllm/model_executor/models/arctic.py +++ b/vllm/model_executor/models/arctic.py @@ -557,7 +557,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.num_experts = config.num_local_experts self.num_experts_per_tok = config.num_experts_per_tok diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index 4fa03261c4b6..d8099c311a59 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -54,6 +54,7 @@ PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -306,7 +307,10 @@ def _flatten_valid_audio_embeddings( < output_lengths[:, None] ) - return audio_embeddings[valid_mask], output_lengths + # Boolean-mask indexing has a data-dependent output shape, so the count + # has to come back to the host. + with gpu_sync_allowed(): + return audio_embeddings[valid_mask], output_lengths def _count_audio_tokens_from_mask( @@ -630,10 +634,10 @@ def _group_audio_embeddings( audio_features, feature_attention_mask, ) - chunk_embeddings = torch.split( - masked_audio_features, - audio_output_lengths.tolist(), - ) + # `split` needs Python int sizes. + with gpu_sync_allowed(): + split_sizes = audio_output_lengths.tolist() + chunk_embeddings = torch.split(masked_audio_features, split_sizes) grouped_embeddings = [] current_idx = 0 diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 0834de5ee507..a3da7e39e78f 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -509,15 +509,14 @@ def __init__( self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.tie_word_embeddings: - self.lm_head = self.model.word_embeddings - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.word_embeddings) self.logits_processor = LogitsProcessor(config.vocab_size) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py index a2b105572d5f..6ddb277778c2 100644 --- a/vllm/model_executor/models/bailing_moe_v3.py +++ b/vllm/model_executor/models/bailing_moe_v3.py @@ -13,6 +13,7 @@ from math import lcm from typing import TypeGuard +import regex as re import torch import torch.nn as nn import torch.nn.functional as F @@ -214,9 +215,49 @@ def _is_block_fp8_config( def _configure_ling_fp8_quant_config( quant_config: QuantizationConfig | None, + config: PretrainedConfig, ) -> None: - if _is_block_fp8_config(quant_config): - quant_config.ignored_layers_match_mode = "suffix" + if not _is_block_fp8_config(quant_config): + return + + quant_config.ignored_layers_match_mode = "suffix" + hf_quant_config = getattr(config, "quantization_config", None) + if not isinstance(hf_quant_config, dict): + return + + quant_config.is_scale_e8m0 = ( # type: ignore[attr-defined] + hf_quant_config.get("scale_fmt") == "ue8m0" + ) + + routed_quant_method = hf_quant_config.get("routed_experts_quant_method") + if routed_quant_method is None: + return + if routed_quant_method != "mxfp4": + raise ValueError( + f"Unsupported routed experts quantization: {routed_quant_method!r}" + ) + + quant_config.store_dtype = "mxfp4" + + +_LING_MXFP4_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_regex={ + re.compile( + r"(\.mlp\.experts\.\d+\." + r"(?:gate_proj|up_proj|down_proj)\.weight_scale)_inv$" + ): r"\1" + } +) + + +def _maybe_remap_ling_mxfp4_weight_names( + weights: Iterable[tuple[str, torch.Tensor]], + quant_config: QuantizationConfig | None, +) -> Iterable[tuple[str, torch.Tensor]]: + """Map Ling's MXFP4 expert scales to Mxfp4MoEMethod parameters.""" + if isinstance(quant_config, Fp8Config) and quant_config.store_dtype == "mxfp4": + return _LING_MXFP4_WEIGHTS_MAPPER.apply(weights) + return weights def _is_fp8_module_excluded( @@ -1351,7 +1392,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config - _configure_ling_fp8_quant_config(quant_config) + _configure_ling_fp8_quant_config(quant_config, config) self.config = config self.quant_config = quant_config self.model = BailingMoeV3Model( @@ -1421,6 +1462,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = self.hf_to_vllm_mapper.apply(weights) + weights = _maybe_remap_ling_mxfp4_weight_names(weights, self.quant_config) params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() stacked_mappings = [ diff --git a/vllm/model_executor/models/bailing_moe_v3_mtp.py b/vllm/model_executor/models/bailing_moe_v3_mtp.py index ffddaa8480cd..037239b21ca5 100644 --- a/vllm/model_executor/models/bailing_moe_v3_mtp.py +++ b/vllm/model_executor/models/bailing_moe_v3_mtp.py @@ -31,6 +31,7 @@ BailingMoeV3MoE, _configure_ling_fp8_quant_config, _maybe_pad_block_fp8_shared_expert_checkpoint_tensor, + _maybe_remap_ling_mxfp4_weight_names, ) from vllm.model_executor.models.interfaces import SupportsPP from vllm.model_executor.models.utils import ( @@ -231,7 +232,7 @@ def __init__( super().__init__() self.config = _get_draft_hf_config(vllm_config) self.quant_config = vllm_config.quant_config - _configure_ling_fp8_quant_config(self.quant_config) + _configure_ling_fp8_quant_config(self.quant_config, self.config) self.model = BailingMoeV3MultiTokenPredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"), @@ -281,6 +282,7 @@ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weights = self.hf_to_vllm_mapper.apply(weights) + weights = _maybe_remap_ling_mxfp4_weight_names(weights, self.quant_config) stacked_params_mapping = [ (".fused_qkv_a_proj", ".q_a_proj", 0), (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1), diff --git a/vllm/model_executor/models/bloom.py b/vllm/model_executor/models/bloom.py index cdcf82f385f7..1347b0e60538 100644 --- a/vllm/model_executor/models/bloom.py +++ b/vllm/model_executor/models/bloom.py @@ -324,14 +324,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.transformer = BloomModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "transformer") ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.config.tie_word_embeddings: - self.lm_head = self.transformer.word_embeddings - else: - self.lm_head = ParallelLMHead( - self.config.vocab_size, - self.config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.transformer.word_embeddings) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( diff --git a/vllm/model_executor/models/chameleon.py b/vllm/model_executor/models/chameleon.py index 94b6b4d34588..03c9fe1179e4 100644 --- a/vllm/model_executor/models/chameleon.py +++ b/vllm/model_executor/models/chameleon.py @@ -977,7 +977,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale) diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index b124f30d520f..0f9bd86dd09a 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -56,6 +56,7 @@ CohereASRProcessor, ) from vllm.utils.collection_utils import is_list_of +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.attention.backend import ( AttentionType, ) @@ -1770,10 +1771,9 @@ def get_encoder_outputs( out = self.encoder_decoder_proj(out) # Convert padded tensor to packed - outs = [] - for i, feat in enumerate(out): - feat_len = encoder_output_length[i] - outs.append(feat[:feat_len, :]) + with gpu_sync_allowed(): + lengths = encoder_output_length.tolist() + outs = [feat[:length, :] for feat, length in zip(out, lengths)] return outs else: diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 6395e4fddf17..0fc32b99d8e9 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -867,7 +867,7 @@ def add_request(self, req_index: int, new_req_data: Any) -> None: self.diffusion_states.add_request(req_index) if not new_req_data.req_id.startswith("_warmup_"): prompt_len = len(new_req_data.prompt_token_ids) - self.diffusion_states.prompt_len[req_index] = prompt_len + self.diffusion_states.prompt_len[req_index].fill_(prompt_len) def remove_request(self, req_id: str) -> None: idx = self._req_id_to_index.pop(req_id, None) @@ -1276,7 +1276,10 @@ def __call__( # before canvas padding so phantom positions stay uniform. if num_decode > 0: top_k, top_p = self.sampling_states.get_top_k_top_p( - decode_slots.repeat_interleave(valid_canvas_len), decode_slots_np + decode_slots.repeat_interleave( + valid_canvas_len, output_size=int(valid_canvas_len_np.sum()) + ), + decode_slots_np, ) if top_k is not None or top_p is not None: logits = apply_top_k_top_p(logits.float(), top_k, top_p) diff --git a/vllm/model_executor/models/ernie45_moe.py b/vllm/model_executor/models/ernie45_moe.py index 89f0b32c66fd..5cc351a6aa75 100644 --- a/vllm/model_executor/models/ernie45_moe.py +++ b/vllm/model_executor/models/ernie45_moe.py @@ -543,7 +543,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.lm_head = PPMissingLayer() if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index f27e047b476c..0c1781b04fa8 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -71,7 +71,9 @@ PromptUpdate, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum from .ernie45_vl_moe import Ernie4_5_VLMoeForCausalLM @@ -446,6 +448,9 @@ def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: pos_ids = torch.cat(pos_ids, dim=0) max_grid_size = grid_thw[:, 1:].max() rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + # `pos_ids` is built on the host; stage it over non-blocking so the + # gather below doesn't index a device tensor with a CPU one. + pos_ids = pos_ids.to(rotary_pos_emb_full.device, non_blocking=True) rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb @@ -530,7 +535,7 @@ def prepare_encoder_metadata( else: max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) - cu_seqlens = cu_seqlens.to(device) + cu_seqlens = cu_seqlens.to(device, non_blocking=True) return { "rotary_pos_emb": rotary_pos_emb, @@ -784,7 +789,10 @@ def fwd_spatial(x): return x def fwd_placeholder(x, grid_thw, to_tensor=False): - grid_thw_cpu = grid_thw.cpu().numpy() + # The per-image grids drive the Python-level offset arithmetic + # below; the same read is wrapped on the cudagraph path. + with gpu_sync_allowed(): + grid_thw_cpu = grid_thw.cpu().numpy() grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:] grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**2) @@ -806,8 +814,8 @@ def fwd_placeholder(x, grid_thw, to_tensor=False): b_offset + (temp_offset + 1) * spatial_size, ) ) - slice_offsets = torch.tensor(np.concatenate(slice_offsets, axis=-1)).to( - x.device + slice_offsets = async_tensor_h2d( + np.concatenate(slice_offsets, axis=-1), device=x.device ) slice_offsets2 = [] @@ -823,8 +831,8 @@ def fwd_placeholder(x, grid_thw, to_tensor=False): b_offset + (temp_offset + 1) * spatial_size, ) ) - slice_offsets2 = torch.tensor(np.concatenate(slice_offsets2, axis=-1)).to( - x.device + slice_offsets2 = async_tensor_h2d( + np.concatenate(slice_offsets2, axis=-1), device=x.device ) x_timestep_1 = torch.index_select(x, dim=0, index=slice_offsets) @@ -1693,8 +1701,11 @@ def prepare_encoder_cudagraph_replay_buffers( EncoderCudaGraphReplayBuffers, ) + # The per-image grids are needed as Python ints to size the buffers. + with gpu_sync_allowed(): + grid_thw_list = mm_kwargs["image_grid_thw"].tolist() metadata = self.vision_model.prepare_encoder_metadata( - mm_kwargs["image_grid_thw"].tolist(), max_batch_size=max_batch_size + grid_thw_list, max_batch_size=max_batch_size ) values = metadata | { "pixel_values": mm_kwargs["pixel_values"], @@ -1734,7 +1745,10 @@ def postprocess_encoder_output( # Ernie only uses the single "default" encoder path. output = outputs["default"] grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device) - num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum()) + # The valid token count slices the graph output for the eager + # resampler call, so it has to come back to the host. + with gpu_sync_allowed(): + num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum()) image_embeds = self.resampler_model(output[:num_valid], grid_thw) scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone) diff --git a/vllm/model_executor/models/ernie45_vl_moe.py b/vllm/model_executor/models/ernie45_vl_moe.py index 5d36ca6d3756..e18a3fec508e 100644 --- a/vllm/model_executor/models/ernie45_vl_moe.py +++ b/vllm/model_executor/models/ernie45_vl_moe.py @@ -611,7 +611,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.lm_head = PPMissingLayer() if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/ernie_mtp.py b/vllm/model_executor/models/ernie_mtp.py index b57da3698a55..88aadbdb73c1 100644 --- a/vllm/model_executor/models/ernie_mtp.py +++ b/vllm/model_executor/models/ernie_mtp.py @@ -177,7 +177,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/model_executor/models/exaone.py b/vllm/model_executor/models/exaone.py index 79314a7b9315..0742ae870b4b 100644 --- a/vllm/model_executor/models/exaone.py +++ b/vllm/model_executor/models/exaone.py @@ -410,7 +410,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.transformer.wte.weight + self.lm_head = self.lm_head.tie_weights(self.transformer.wte) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index 8e8252f45922..571aa3a42645 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -410,7 +410,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/exaone4_5_mtp.py b/vllm/model_executor/models/exaone4_5_mtp.py index 7711f72e42ca..80e1903f9dd8 100644 --- a/vllm/model_executor/models/exaone4_5_mtp.py +++ b/vllm/model_executor/models/exaone4_5_mtp.py @@ -159,7 +159,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor( self.unpadded_vocab_size, config.vocab_size ) diff --git a/vllm/model_executor/models/exaone_moe.py b/vllm/model_executor/models/exaone_moe.py index e942bf62ed05..ef01330f0c56 100644 --- a/vllm/model_executor/models/exaone_moe.py +++ b/vllm/model_executor/models/exaone_moe.py @@ -522,7 +522,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): quant_config=quant_config, ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/exaone_moe_mtp.py b/vllm/model_executor/models/exaone_moe_mtp.py index ddd570d604e7..ebe1f8cbaeb6 100644 --- a/vllm/model_executor/models/exaone_moe_mtp.py +++ b/vllm/model_executor/models/exaone_moe_mtp.py @@ -159,7 +159,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor( self.unpadded_vocab_size, config.vocab_size ) diff --git a/vllm/model_executor/models/falcon.py b/vllm/model_executor/models/falcon.py index efd24b51442a..cdc388e529b4 100644 --- a/vllm/model_executor/models/falcon.py +++ b/vllm/model_executor/models/falcon.py @@ -499,15 +499,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if config.tie_word_embeddings is not None else True ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.tie_word_embeddings: - self.lm_head = self.transformer.word_embeddings - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.transformer.word_embeddings) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.transformer.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/funasr.py b/vllm/model_executor/models/funasr.py index e128b22e8e04..56b1b68f4f7e 100644 --- a/vllm/model_executor/models/funasr.py +++ b/vllm/model_executor/models/funasr.py @@ -919,15 +919,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) logit_scale = getattr(config, "logit_scale", 1.0) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.decoder.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.decoder.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size, scale=logit_scale) def get_language_model(self) -> torch.nn.Module: diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 1ef236907fc0..aa9669aeb157 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -50,6 +50,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, + MultiModalKwargsItem, MultiModalKwargsItems, VideoItem, ) @@ -68,6 +69,7 @@ PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import async_tensor_h2d @@ -799,7 +801,7 @@ def _get_mm_fields_config( MultiModalFieldConfig.flat_from_sizes("video", vfc) ), video_frame_counts=MultiModalFieldConfig.batched( - "video", + "video", keep_on_cpu=True ), video_num_soft_tokens=MultiModalFieldConfig.batched( "video", keep_on_cpu=True @@ -1037,6 +1039,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.multimodal_config = multimodal_config self.model_dtype = vllm_config.model_config.dtype self.vllm_config = vllm_config + lora_config = vllm_config.lora_config + self._enable_mm_lora = bool( + lora_config is not None and lora_config.enable_tower_connector_lora + ) # Only quantize towers when the quant method supports their # dimensions. BNB/torchao handle arbitrary sizes; other methods @@ -1265,8 +1271,8 @@ def _process_image_input( Groups images by patch count (resolution bucket) so each encoder call processes a uniform-shape batch with no - cross-resolution padding. Pooling and projection are then - applied over a single concatenated tensor for all images. + cross-resolution padding. With MM LoRA enabled, all images are + padded into one batch so the encoder call matches the tower mapping. """ pixel_values = image_input["pixel_values"] pixel_position_ids = image_input["pixel_position_ids"] @@ -1284,24 +1290,65 @@ def _process_image_input( if isinstance(pixel_values, list) else pixel_values.shape[0] ) + pool_position_ids = pixel_position_ids - for idx in range(total_images): - pv = pixel_values[idx] - pp = pixel_position_ids[idx] - buckets.setdefault(pv.shape[0], []).append((idx, pv, pp)) + if self._enable_mm_lora: + max_soft_tokens = vision_cfg.default_output_length + mm_processor_kwargs = getattr( + getattr(self, "multimodal_config", None), + "mm_processor_kwargs", + None, + ) + if isinstance(mm_processor_kwargs, Mapping): + value, _ = _get_max_soft_tokens(mm_processor_kwargs) + if isinstance(value, int) and value in _SUPPORTED_SOFT_TOKENS: + max_soft_tokens = value + + max_patches = max_soft_tokens * pooling_k2 + padded_position_ids: list[torch.Tensor] = [] + for idx in range(total_images): + pv = pixel_values[idx] + pp = pixel_position_ids[idx] + num_patches = pv.shape[0] + if num_patches > max_patches: + raise ValueError( + f"Image {idx} has {num_patches} patches, which exceeds " + f"the MM LoRA patch limit of {max_patches}." + ) + + pad_len = max_patches - num_patches + pv = torch.cat( + (pv, pv.new_zeros((pad_len, *pv.shape[1:]))), + dim=0, + ) + pp = torch.cat( + (pp, pp.new_full((pad_len, *pp.shape[1:]), -1)), + dim=0, + ) + buckets.setdefault(max_patches, []).append((idx, pv, pp)) + padded_position_ids.append(pp) + pool_position_ids = padded_position_ids + else: + for idx in range(total_images): + pv = pixel_values[idx] + pp = pixel_position_ids[idx] + buckets.setdefault(pv.shape[0], []).append((idx, pv, pp)) # Encode each resolution bucket in memory-safe chunks. Re-read # free memory per bucket because the previous bucket's encoder # pass has already allocated activations we should account for. last_hidden_states_map: dict[int, torch.Tensor] = {} for patches, items in buckets.items(): - free, total = torch.accelerator.get_memory_info() - max_batch_size = min( - len(items), - self._encoder_chunk( - patches, free, total, vision_cfg.position_embedding_size - ), - ) + if self._enable_mm_lora: + max_batch_size = len(items) + else: + free, total = torch.accelerator.get_memory_info() + max_batch_size = min( + len(items), + self._encoder_chunk( + patches, free, total, vision_cfg.position_embedding_size + ), + ) for chunk_idx in range(0, len(items), max_batch_size): chunk_items = items[chunk_idx : chunk_idx + max_batch_size] @@ -1319,11 +1366,14 @@ def _process_image_input( pp_tensor, pad_tensor, ).to(self.model_dtype) - encoder_outputs = vt.encoder( - inputs_embeds=inputs_embeds, - attention_mask=~pad_tensor, - pixel_position_ids=pp_tensor, - ) + # HuggingFace's mask builder probes `padding_mask.all()` to + # decide whether the mask can be skipped, which syncs. + with gpu_sync_allowed(): + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_tensor, + pixel_position_ids=pp_tensor, + ) hidden_states = encoder_outputs.last_hidden_state for i, (orig_idx, _, _) in enumerate(chunk_items): @@ -1338,16 +1388,20 @@ def _process_image_input( output_length = chunk_hidden.shape[0] // pooling_k2 single_hidden = chunk_hidden.unsqueeze(0) - single_pos_ids = pixel_position_ids[orig_idx].unsqueeze(0) + single_pos_ids = pool_position_ids[orig_idx].unsqueeze(0) padding_positions = (single_pos_ids == -1).all(dim=-1) - pooled_states, valid_mask = vt.pooler( - hidden_states=single_hidden, - pixel_position_ids=single_pos_ids, - padding_positions=padding_positions, - output_length=output_length, - ) - valid_states = pooled_states[valid_mask] + # The pooler goes through HuggingFace's mask builder, which probes + # `padding_mask.all()`, and the mask indexing below needs the + # selected count on the host. + with gpu_sync_allowed(): + pooled_states, valid_mask = vt.pooler( + hidden_states=single_hidden, + pixel_position_ids=single_pos_ids, + padding_positions=padding_positions, + output_length=output_length, + ) + valid_states = pooled_states[valid_mask] if getattr(vt.config, "standardize", False): valid_states = (valid_states - vt.std_bias) * vt.std_scale @@ -1398,7 +1452,9 @@ def _process_video_input( pooling_k2 = vision_cfg.pooling_kernel_size**2 if isinstance(frame_counts, torch.Tensor): - fc_list = frame_counts.tolist() + # Per-video frame counts drive the Python-level batching below. + with gpu_sync_allowed(): + fc_list = frame_counts.tolist() else: fc_list = list(frame_counts) @@ -1428,11 +1484,13 @@ def _process_video_input( pp_chunk, pad_chunk, ).to(self.model_dtype) - encoder_outputs = vt.encoder( - inputs_embeds=inputs_embeds, - attention_mask=~pad_chunk, - pixel_position_ids=pp_chunk, - ) + # HuggingFace's mask builder probes `padding_mask.all()`. + with gpu_sync_allowed(): + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_chunk, + pixel_position_ids=pp_chunk, + ) last_hidden_states_list.append(encoder_outputs.last_hidden_state) last_hidden_states = torch.cat(last_hidden_states_list, dim=0) @@ -1447,13 +1505,15 @@ def _process_video_input( single_pos_ids = pixel_position_ids[i].unsqueeze(0) single_pad_pos = padding_positions[i].unsqueeze(0) - pooled_states, valid_mask = vt.pooler( - hidden_states=single_hidden, - pixel_position_ids=single_pos_ids, - padding_positions=single_pad_pos, - output_length=output_length, - ) - valid_states = pooled_states[valid_mask] + # As above, plus mask indexing that needs the count on the host. + with gpu_sync_allowed(): + pooled_states, valid_mask = vt.pooler( + hidden_states=single_hidden, + pixel_position_ids=single_pos_ids, + padding_positions=single_pad_pos, + output_length=output_length, + ) + valid_states = pooled_states[valid_mask] if getattr(vt.config, "standardize", False): valid_states = (valid_states - vt.std_bias) * vt.std_scale @@ -1507,9 +1567,11 @@ def _process_audio_input( # Strip padding per-batch element: only keep valid (non-padding) # tokens. + # Boolean-mask indexing needs the selected count on the host. per_audio = [] - for enc, mask in zip(audio_features, audio_mask, strict=True): - per_audio.append(enc[mask]) # [num_real, hidden_size] + with gpu_sync_allowed(): + for enc, mask in zip(audio_features, audio_mask, strict=True): + per_audio.append(enc[mask]) # [num_real, hidden_size] return per_audio @@ -2146,6 +2208,66 @@ def get_mm_mapping(self) -> MultiModelKeys: tower_model=tower_models, ) + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: MultiModalKwargsItem | None, + num_mm_embeds: int, + ) -> tuple[int, int | None]: + if modality in ("image", "video"): + vision_config = self.config.vision_config + pooling_k2 = vision_config.pooling_kernel_size**2 + + if modality == "image": + pixel_values_key = "pixel_values" + max_soft_tokens = vision_config.default_output_length + mm_processor_kwargs = getattr( + getattr(self, "multimodal_config", None), + "mm_processor_kwargs", + None, + ) + if isinstance(mm_processor_kwargs, Mapping): + val, _ = _get_max_soft_tokens(mm_processor_kwargs) + if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS: + max_soft_tokens = val + else: + pixel_values_key = "pixel_values_videos" + max_soft_tokens = _VIDEO_MAX_SOFT_TOKENS + + tower_tokens = max_soft_tokens * pooling_k2 if modality == "image" else None + connector_tokens = num_mm_embeds + if tower_tokens is None and mm_kwargs is not None: + field = mm_kwargs.get(pixel_values_key) + if field is not None: + data = field.data + if isinstance(data, torch.Tensor) and data.ndim >= 2: + tower_tokens = int(math.prod(data.shape[:-1])) + + if tower_tokens is None: + min_soft_tokens = min(_SUPPORTED_SOFT_TOKENS) + tower_tokens = ( + math.ceil(num_mm_embeds / min_soft_tokens) + * max_soft_tokens + * pooling_k2 + ) + + if modality == "audio": + tower_tokens = num_mm_embeds + connector_tokens = num_mm_embeds + + if mm_kwargs is not None: + field = mm_kwargs.get("input_features_padded") + if field is not None: + data = field.data + if isinstance(data, torch.Tensor) and data.ndim >= 2: + batch_size = math.prod(data.shape[:-2]) + audio_tokens = batch_size * math.ceil(data.shape[-2] / 4) + tower_tokens = audio_tokens + connector_tokens = audio_tokens + + return tower_tokens, connector_tokens + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality == "image": diff --git a/vllm/model_executor/models/gemma4_mtp.py b/vllm/model_executor/models/gemma4_mtp.py index 0fe972691e9e..9fd2b1f1641e 100644 --- a/vllm/model_executor/models/gemma4_mtp.py +++ b/vllm/model_executor/models/gemma4_mtp.py @@ -489,7 +489,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if getattr(config, "tie_word_embeddings", True): - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor( text_config.vocab_size, diff --git a/vllm/model_executor/models/glm4.py b/vllm/model_executor/models/glm4.py index a1fb94fb26fa..4c30e30008fe 100644 --- a/vllm/model_executor/models/glm4.py +++ b/vllm/model_executor/models/glm4.py @@ -252,15 +252,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/glm4_1v.py b/vllm/model_executor/models/glm4_1v.py index 0e1f079dfae4..679c90728bd3 100644 --- a/vllm/model_executor/models/glm4_1v.py +++ b/vllm/model_executor/models/glm4_1v.py @@ -546,14 +546,6 @@ def forward( 0, hidden_size, device=device, dtype=pos_embed_weight.dtype ) else: - # Convert inputs to tensors if needed - if isinstance(lengths, list): - lengths = async_tensor_h2d(lengths, device=device, dtype=torch.long) - if not isinstance(image_shapes, torch.Tensor): - image_shapes = async_tensor_h2d( - image_shapes, device=device, dtype=torch.long - ) - # Prepare 2D position embedding orig_size_sq = pos_embed_weight.shape[0] orig_size = int(orig_size_sq**0.5) @@ -564,32 +556,23 @@ def forward( .to(device=device, dtype=torch.float32) ) - # Calculate target dimensions for each patch - # Add bounds checking for data parallel mode - if len(lengths) > image_shapes.shape[0]: - # In data parallel mode, some GPUs might not have all - # image shapes - # Use available image shapes, cycling if necessary - target_h_list = [] - target_w_list = [] - for i in range(len(lengths)): - # Cycle through available shapes - shape_idx = i % image_shapes.shape[0] - target_h_list.append(image_shapes[shape_idx, 1].repeat(lengths[i])) - target_w_list.append(image_shapes[shape_idx, 2].repeat(lengths[i])) - target_h = torch.cat(target_h_list).to( - device=device, dtype=torch.float32 - ) - target_w = torch.cat(target_w_list).to( - device=device, dtype=torch.float32 - ) - else: - target_h = torch.cat( - [image_shapes[i, 1].repeat(lengths[i]) for i in range(len(lengths))] - ).to(device=device, dtype=torch.float32) - target_w = torch.cat( - [image_shapes[i, 2].repeat(lengths[i]) for i in range(len(lengths))] - ).to(device=device, dtype=torch.float32) + # Calculate target dimensions for each patch. `lengths` and + # `image_shapes` are host data, so expand them with numpy and move + # the result across once rather than per-image. + # Shapes are cycled: in data parallel mode some GPUs might not + # have all image shapes. + shapes_np = np.asarray(image_shapes) + shape_idx = np.arange(len(lengths)) % shapes_np.shape[0] + target_h = async_tensor_h2d( + np.repeat(shapes_np[shape_idx, 1], lengths), + device=device, + dtype=torch.float32, + ) + target_w = async_tensor_h2d( + np.repeat(shapes_np[shape_idx, 2], lengths), + device=device, + dtype=torch.float32, + ) # Normalize coordinates to [-1, 1] range for grid_sample h_coords = h_coords.to(device=device, dtype=torch.float32) @@ -818,9 +801,7 @@ def pos_embeds_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor: ) lengths = [h * w] * t - image_shapes = async_tensor_h2d( - [[t, h, w]], dtype=torch.long, device=device - ) + image_shapes = [[t, h, w]] # Build the coordinates on the host (cheap integer math) but move # them across pinned + non-blocking, so the consumer's diff --git a/vllm/model_executor/models/gpt_neox.py b/vllm/model_executor/models/gpt_neox.py index 907ab7776015..8af7eb67923e 100644 --- a/vllm/model_executor/models/gpt_neox.py +++ b/vllm/model_executor/models/gpt_neox.py @@ -285,7 +285,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "embed_out"), ) if self.config.tie_word_embeddings: - self.embed_out.weight = self.gpt_neox.embed_in.weight + self.embed_out = self.embed_out.tie_weights(self.gpt_neox.embed_in) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.gpt_neox.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index c46fefbf8894..a0bc8350b43a 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -365,7 +365,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) if hasattr(config, "logits_scaling"): diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index c6e4df2992cb..5c56717aa631 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -316,7 +316,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) if hasattr(config, "logits_scaling"): logit_scale /= config.logits_scaling diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 262eda4c2535..0e101212744e 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -60,6 +60,7 @@ from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import PIN_MEMORY @@ -788,8 +789,9 @@ def _process_audio_input( encoder_embeds = self.encoder(audio_input["input_features"]) # [bsz, , 4096] projected_embeds = self.projector(encoder_embeds) - # Apply mask on variable length audio features - masked_embeds = projected_embeds[audio_input["input_features_mask"]] + # Apply mask on variable length audio features. + with gpu_sync_allowed(): + masked_embeds = projected_embeds[audio_input["input_features_mask"]] # Split variable length features into a tuple return torch.split(masked_embeds, audio_input["audio_embed_sizes"]) diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py index 328ea9ce6326..e1219e3337dd 100644 --- a/vllm/model_executor/models/granitemoe.py +++ b/vllm/model_executor/models/granitemoe.py @@ -505,7 +505,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor( config.vocab_size, diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index a50a95a302e8..aa21c23f1e11 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -663,7 +663,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor( config.vocab_size, config.vocab_size, diff --git a/vllm/model_executor/models/granitemoeshared.py b/vllm/model_executor/models/granitemoeshared.py index 7abc682c58e5..603afc60d683 100644 --- a/vllm/model_executor/models/granitemoeshared.py +++ b/vllm/model_executor/models/granitemoeshared.py @@ -281,7 +281,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor( config.vocab_size, diff --git a/vllm/model_executor/models/hrm_text.py b/vllm/model_executor/models/hrm_text.py index a7546b0cc444..dd7ba93e7b8a 100644 --- a/vllm/model_executor/models/hrm_text.py +++ b/vllm/model_executor/models/hrm_text.py @@ -487,15 +487,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/hunyuan_v1.py b/vllm/model_executor/models/hunyuan_v1.py index ef7c74851960..b90d463627d9 100644 --- a/vllm/model_executor/models/hunyuan_v1.py +++ b/vllm/model_executor/models/hunyuan_v1.py @@ -927,7 +927,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/hy_v3.py b/vllm/model_executor/models/hy_v3.py index a4f4a558ec1c..0248ac3d9c02 100644 --- a/vllm/model_executor/models/hy_v3.py +++ b/vllm/model_executor/models/hy_v3.py @@ -673,7 +673,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index b6f5531074a2..72a26993a537 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -39,6 +39,7 @@ RowParallelLinear, ) from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from .utils import AutoWeightsLoader, WeightsMapper from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model @@ -96,10 +97,14 @@ def forward( size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0 ) - for batch_idx, p_attn_mask in enumerate(patch_attention_mask): - if tgt_sizes is not None: - nb_patches_h = tgt_sizes[batch_idx][0] - nb_patches_w = tgt_sizes[batch_idx][1] + with gpu_sync_allowed(): + patch_attention_mask_cpu = patch_attention_mask.cpu() + tgt_sizes_cpu = tgt_sizes.cpu() if tgt_sizes is not None else None + + for batch_idx, p_attn_mask in enumerate(patch_attention_mask_cpu): + if tgt_sizes_cpu is not None: + nb_patches_h = tgt_sizes_cpu[batch_idx][0] + nb_patches_w = tgt_sizes_cpu[batch_idx][1] else: nb_patches_h = p_attn_mask[:, 0].sum() nb_patches_w = p_attn_mask[0].sum() @@ -114,8 +119,10 @@ def forward( pos_ids = ( bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w ).flatten() - position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids - position_ids = position_ids.to(self.position_embedding.weight.device) + position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids + position_ids = position_ids.to( + self.position_embedding.weight.device, non_blocking=True + ) embeddings += self.position_embedding(position_ids) return embeddings diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index 732f5c021af1..7b3e552f0cdd 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -51,6 +51,7 @@ PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .idefics2_vision_model import ( @@ -443,11 +444,12 @@ def image_pixels_to_features( real_images_inds = (pixel_values == 0.0).sum( dim=(-1, -2, -3) ) != nb_values_per_image - pixel_values = pixel_values[real_images_inds].contiguous() + with gpu_sync_allowed(): + pixel_values = pixel_values[real_images_inds].contiguous() - # Handle the vision attention mask - # Remove padding images from the mask - pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() + # Handle the vision attention mask + # Remove padding images from the mask + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() patch_size = self.config.vision_config.patch_size patches_subgrid = pixel_attention_mask.unfold( @@ -539,7 +541,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.text_config.tie_word_embeddings: - self.lm_head.weight = self.model.text_model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.text_model.embed_tokens) self.logits_processor = LogitsProcessor(config.text_config.vocab_size) def _parse_and_validate_image_input(self, **kwargs: object) -> ImageInputs | None: diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 46d0ed30df45..555051237336 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -48,7 +48,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc from vllm.model_executor.models.interfaces_base import VllmModel from vllm.model_executor.models.utils import WeightsMapper - from vllm.multimodal.inputs import MultiModalFeatureSpec + from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem from vllm.multimodal.registry import _ProcessorFactories from vllm.sequence import IntermediateTensors from vllm.tasks import ScoreType @@ -75,6 +75,12 @@ | tuple[tuple[int, int, int]] | tuple[tuple[int, int], tuple[int, int]] | tuple[tuple[int, int], tuple[int, int, int]] + | tuple[ + tuple[int, int], + tuple[int, int, int], + tuple[int, int, int], + tuple[int, int, int], + ] ) @@ -385,6 +391,28 @@ def get_num_mm_connector_tokens(self, num_vision_tokens: int) -> int: """ ... + def get_mm_lora_token_counts( + self, + *, + modality: str, + mm_kwargs: "MultiModalKwargsItem | None", + num_mm_embeds: int, + ) -> tuple[int, int | None]: + """ + Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings. + + MM LoRA uses these counts to build adapter mappings for the tower and + connector forwards. Models with multiple modalities can override this + when each modality has different encoder padding or pooling behavior. + """ + del modality, mm_kwargs + num_encoder_tokens = self.get_num_mm_encoder_tokens(num_mm_embeds) + num_connector_tokens = self.get_num_mm_connector_tokens(num_encoder_tokens) + return ( + num_encoder_tokens, + num_connector_tokens if isinstance(num_connector_tokens, int) else None, + ) + @overload def embed_input_ids(self, input_ids: Tensor) -> Tensor: ... @@ -1077,8 +1105,8 @@ def supports_mamba_prefix_caching( @runtime_checkable class SupportsReplaySSM(Protocol): - """The interface for models whose Mamba2 layers support ReplaySSM cached - standard decode. + """The interface for models whose recurrent layers support ReplaySSM + cached decode. This is currently experimental. """ @@ -1545,13 +1573,15 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: parent_ref = self.get_language_model() elif hasattr(self, "language_model"): parent_ref = self.language_model - assert hasattr(parent_ref, "model"), ( - "Model instance must have 'model' attribute to set number of layers" - ) - assert isinstance(parent_ref.model, EagleModelMixin), ( + # A multimodal model that builds its decoder inside + # `_mark_language_model` has get_language_model() return the inner + # decoder itself, which IS the EagleModelMixin and has no further + # `.model`. Unwrap only when there is something to unwrap. + holder = getattr(parent_ref, "model", parent_ref) + assert isinstance(holder, EagleModelMixin), ( "Model instance must inherit from EagleModelMixin to set auxiliary layers" ) - parent_ref.model._set_aux_hidden_state_layers(layers) + holder._set_aux_hidden_state_layers(layers) def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: """ @@ -1568,13 +1598,12 @@ def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]: parent_ref = self.get_language_model() elif hasattr(self, "language_model"): parent_ref = self.language_model - assert hasattr(parent_ref, "model"), ( - "Model instance must have 'model' attribute to get number of layers" - ) - assert hasattr(parent_ref.model, "layers"), ( + # Same unwrap-only-if-needed rule as set_aux_hidden_state_layers. + holder = getattr(parent_ref, "model", parent_ref) + assert hasattr(holder, "layers"), ( "Model instance must have 'layers' attribute to get number of layers" ) - num_layers = len(parent_ref.model.layers) + num_layers = len(holder.layers) return (2, num_layers // 2, num_layers - 3) diff --git a/vllm/model_executor/models/internlm2.py b/vllm/model_executor/models/internlm2.py index 81487f9cad5b..ffc750c7c565 100644 --- a/vllm/model_executor/models/internlm2.py +++ b/vllm/model_executor/models/internlm2.py @@ -352,7 +352,7 @@ def __init__( prefix=maybe_prefix(prefix, "output"), ) if self.config.tie_word_embeddings: - self.output.weight = self.model.tok_embeddings.weight + self.output = self.output.tie_weights(self.model.tok_embeddings) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/interns1.py b/vllm/model_executor/models/interns1.py index a447f3e2ee8e..6dbc358e089e 100644 --- a/vllm/model_executor/models/interns1.py +++ b/vllm/model_executor/models/interns1.py @@ -424,7 +424,7 @@ def _get_mm_fields_config( pixel_values=MultiModalFieldConfig.flat_from_sizes( "image", image_num_patches ), - image_num_patches=MultiModalFieldConfig.batched("image"), + image_num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True), image_embeds=MultiModalFieldConfig.batched("image"), image_token_id=MultiModalFieldConfig.shared( "image", num_images, keep_on_cpu=True @@ -432,7 +432,7 @@ def _get_mm_fields_config( pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( "video", video_num_patches ), - video_num_patches=MultiModalFieldConfig.batched("video"), + video_num_patches=MultiModalFieldConfig.batched("video", keep_on_cpu=True), video_token_id=MultiModalFieldConfig.shared( "video", num_videos, keep_on_cpu=True ), diff --git a/vllm/model_executor/models/interns1_pro.py b/vllm/model_executor/models/interns1_pro.py index 378177a5575d..e85ad87cdc7c 100644 --- a/vllm/model_executor/models/interns1_pro.py +++ b/vllm/model_executor/models/interns1_pro.py @@ -481,7 +481,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(self.config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/interns2_mobius.py b/vllm/model_executor/models/interns2_mobius.py index e13de751e63b..b8111412f0ba 100644 --- a/vllm/model_executor/models/interns2_mobius.py +++ b/vllm/model_executor/models/interns2_mobius.py @@ -43,7 +43,7 @@ Qwen3_5RMSNorm, ) from .qwen3_5_mtp import Qwen3_5MoeMTP -from .qwen3_next import Qwen3NextAttention, _all_gather_hidden_and_residual +from .qwen3_next import Qwen3NextAttention from .qwen3_vl import ( Qwen3_VisionTransformer, Qwen3VLDummyInputsBuilder, @@ -233,11 +233,6 @@ def forward( meta_mlp: nn.ModuleList, ) -> tuple[torch.Tensor, torch.Tensor]: full_num_tokens = positions.shape[-1] - input_is_sequence_parallel = ( - self.use_attn_reduce_scatter_for_moe - and residual is not None - and hidden_states.shape[0] != full_num_tokens - ) if residual is None: residual = hidden_states @@ -245,7 +240,7 @@ def forward( else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - if input_is_sequence_parallel: + if self.use_attn_reduce_scatter_for_moe: hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) hidden_states = hidden_states[:full_num_tokens] @@ -266,8 +261,6 @@ def forward( sp_pad = (-hidden_states.shape[0]) % tp_world_size hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad)) hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) - if not input_is_sequence_parallel: - residual = sequence_parallel_chunk(residual) hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) block_idx = self.layer_idx % self.num_blocks @@ -343,6 +336,10 @@ def get_layer(prefix: str) -> InternS2MobiusDecoderLayer: def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + @property + def use_sequence_parallel(self) -> bool: + return self.layers[self.start_layer].use_attn_reduce_scatter_for_moe + def forward( self, input_ids: torch.Tensor | None, @@ -363,36 +360,21 @@ def forward( residual = intermediate_tensors["residual"] full_num_tokens = positions.shape[-1] + if self.use_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + assert residual is None + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): - if ( - hidden_states.shape[0] != full_num_tokens - and not layer.use_attn_reduce_scatter_for_moe - ): - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - full_num_tokens, - self.config.hidden_size, - ) hidden_states, residual = layer( positions=positions, hidden_states=hidden_states, residual=residual, meta_mlp=self.meta_mlp, ) - if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[ - 0 - ] != full_num_tokens: - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - full_num_tokens, - self.config.hidden_size, - ) self._maybe_add_hidden_state( aux_hidden_states, layer_idx + 1, @@ -404,14 +386,19 @@ def forward( return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) - if hidden_states.shape[0] != full_num_tokens: - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - full_num_tokens, - self.config.hidden_size, - ) hidden_states, _ = self.norm(hidden_states, residual) + if self.use_sequence_parallel: + if aux_hidden_states: + hidden_size = hidden_states.shape[-1] + hidden_states = torch.cat([hidden_states, *aux_hidden_states], dim=-1) + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + hidden_states, *aux_hidden_states = hidden_states.split( + hidden_size, dim=-1 + ) + else: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] if aux_hidden_states: return hidden_states, aux_hidden_states return hidden_states @@ -443,15 +430,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=self.quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/iquest_loopcoder.py b/vllm/model_executor/models/iquest_loopcoder.py index 3755cba5d1ae..fdb90b962d08 100644 --- a/vllm/model_executor/models/iquest_loopcoder.py +++ b/vllm/model_executor/models/iquest_loopcoder.py @@ -541,15 +541,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 5fae3031542a..0595df8e3a28 100644 --- a/vllm/model_executor/models/jina.py +++ b/vllm/model_executor/models/jina.py @@ -16,6 +16,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import PoolingTask from vllm.transformers_utils.repo_utils import get_hf_file_bytes +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.pool.metadata import PoolingMetadata from ..layers.pooler import DispatchPooler @@ -110,19 +111,24 @@ def forward( prompt_token_ids = pooling_metadata.get_prompt_token_ids() embeds_list = list[torch.Tensor | None]() - for data, token_ids in zip(pooled_data_lst, prompt_token_ids): - # for unfinished chunked prefill - if data is None: - embeds_list.append(None) - else: - docs_indexes = torch.where(torch.eq(token_ids, self.doc_token_id))[0] - query_indexes = torch.where(torch.eq(token_ids, self.query_token_id))[0] - - # The JinaForRanking model concatenates docs first, then query. - # Let's stay consistent with this novel design. - indexes = torch.cat([docs_indexes, query_indexes]) - embeds = self.projector(data[indexes]) - embeds_list.append(embeds) + # `torch.where` resolves the match count on the host. + with gpu_sync_allowed(): + for data, token_ids in zip(pooled_data_lst, prompt_token_ids): + # for unfinished chunked prefill + if data is None: + embeds_list.append(None) + else: + doc_match = torch.eq(token_ids, self.doc_token_id) + docs_indexes = torch.where(doc_match)[0] + query_indexes = torch.where( + torch.eq(token_ids, self.query_token_id) + )[0] + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + indexes = torch.cat([docs_indexes, query_indexes]) + embeds = self.projector(data[indexes]) + embeds_list.append(embeds) return embeds_list diff --git a/vllm/model_executor/models/keye.py b/vllm/model_executor/models/keye.py index 005af154c043..28ccddd47e9f 100644 --- a/vllm/model_executor/models/keye.py +++ b/vllm/model_executor/models/keye.py @@ -59,6 +59,7 @@ ) from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm.utils.torch_utils import async_tensor_h2d from .interfaces import ( MultiModalEmbeddings, @@ -608,7 +609,9 @@ def forward( [height_position_ids, width_position_ids], dim=-1, ) - max_grid_size = pids.max() + 1 + # The ids are built from the grids above, so `h`/`w` bound them + # and the table size is known on the host. + max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw) rope_emb_max_grid = self.rotary_pos_emb(max_grid_size) rope_emb = rope_emb_max_grid[pids].flatten(1) rope_emb = rope_emb.repeat(1, 2) @@ -694,19 +697,19 @@ def forward( last_hidden_state = self.post_layernorm(last_hidden_state) - sample_hidden_state = list() if cu_seqlens is None: raise ValueError( "cu_seqlens cannot be None for " "SiglipVisionTransformer output processing." ) - for i in range(cu_seqlens.shape[0] - 1): - start = cu_seqlens[i] - end = cu_seqlens[i + 1] - tensor = last_hidden_state[:, start:end, :].squeeze(0) - sample_hidden_state.append(tensor) - - return sample_hidden_state + # `cu_seqlens` is the running sum of the per-image `t * h * w`, so the + # split sizes are known on the host and slicing needs no device read. + split_sizes = [ + int(np.prod(thw)) for thw in self.encoder.flatten_list(image_grid_thw) + ] + return [ + tensor.squeeze(0) for tensor in last_hidden_state.split(split_sizes, dim=1) + ] class KeyeSiglipVisionModel(nn.Module): @@ -1307,13 +1310,17 @@ def _process_image_input(self, image_input: Any) -> tuple[torch.Tensor, ...]: ) else: pixel_values = image_input["pixel_values"].type(self.visual.dtype) + # These are all built on the host, so stage them over + # non-blocking rather than forcing a sync per transfer. siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to( - pixel_values.device + pixel_values.device, non_blocking=True + ) + cu_seqlens = async_tensor_h2d( + cu_seqlens, dtype=torch.int32, device=pixel_values.device ) - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to( - pixel_values.device + sample_indices = torch.concat(sample_indices, dim=0).to( + pixel_values.device, non_blocking=True ) - sample_indices = torch.concat(sample_indices, dim=0).to(pixel_values.device) image_embeds = self.visual( pixel_values=pixel_values, @@ -1357,14 +1364,15 @@ def _process_video_embeds( ) else: pixel_values_videos = pixel_values_videos.type(self.visual.dtype) + # These are host-built; stage them over non-blocking. siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to( - pixel_values_videos.device + pixel_values_videos.device, non_blocking=True ) - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to( - pixel_values_videos.device + cu_seqlens = async_tensor_h2d( + cu_seqlens, dtype=torch.int32, device=pixel_values_videos.device ) sample_indices = torch.concat(sample_indices, dim=0).to( - pixel_values_videos.device + pixel_values_videos.device, non_blocking=True ) video_embeds = self.visual( diff --git a/vllm/model_executor/models/mamba.py b/vllm/model_executor/models/mamba.py index 6a77a58abf4d..72893c574f02 100644 --- a/vllm/model_executor/models/mamba.py +++ b/vllm/model_executor/models/mamba.py @@ -188,14 +188,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "backbone") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.backbone.embeddings - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.backbone.embeddings) self.logits_processor = LogitsProcessor(config.vocab_size) diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py index c20fa00e3d6f..2eb35687d737 100644 --- a/vllm/model_executor/models/mellum.py +++ b/vllm/model_executor/models/mellum.py @@ -221,7 +221,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/mimo.py b/vllm/model_executor/models/mimo.py index e4247fa8d8df..56250ba2f4e8 100644 --- a/vllm/model_executor/models/mimo.py +++ b/vllm/model_executor/models/mimo.py @@ -101,15 +101,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index 4a3ad03ed784..d9358fed425f 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -44,7 +44,11 @@ from typing_extensions import TypeVar from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import ( + BaseDummyOptions, + ImageDummyOptions, + VideoDummyOptions, +) from vllm.inputs import ModalityData, MultiModalDataDict from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.resampler import ( @@ -94,6 +98,7 @@ ) from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.collection_utils import flatten_2d_lists +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import set_default_torch_dtype @@ -811,6 +816,18 @@ def get_dummy_mm_data( image_overrides = mm_options.get("image") video_overrides = mm_options.get("video") + # Convert video overrides to image overrides for per-frame image generation, + # and apply num_frames override to num_video_frames. + video_frame_overrides: ImageDummyOptions | None = None + if isinstance(video_overrides, VideoDummyOptions): + if video_overrides.num_frames: + num_video_frames = min(num_video_frames, video_overrides.num_frames) + if video_overrides.width or video_overrides.height: + video_frame_overrides = ImageDummyOptions( + width=video_overrides.width, + height=video_overrides.height, + ) + return { "image": self._get_dummy_images( width=image_width, @@ -823,7 +840,7 @@ def get_dummy_mm_data( width=video_width, height=video_height, num_images=num_video_frames, - overrides=video_overrides, + overrides=video_frame_overrides, ) ] * num_videos, @@ -1532,7 +1549,8 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) @@ -1624,7 +1642,9 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + # Needed as a Python int to size the mask below. + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) @@ -1721,7 +1741,8 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) @@ -1823,7 +1844,8 @@ def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tens all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) diff --git a/vllm/model_executor/models/mixtral.py b/vllm/model_executor/models/mixtral.py index f9b452f47f7e..d154208606c7 100644 --- a/vllm/model_executor/models/mixtral.py +++ b/vllm/model_executor/models/mixtral.py @@ -403,7 +403,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index cf4550a10765..bdc46d83b934 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -1340,15 +1340,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.img_patch_id = None + self.lm_head = ParallelLMHead( + config.embedding_size or config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.config.weight_tying: - self.lm_head = self.model.transformer.wte - else: - self.lm_head = ParallelLMHead( - config.embedding_size or config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.transformer.wte) self.logits_processor = LogitsProcessor( config.embedding_size or config.vocab_size diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py index 56204dd3c61e..29400f597b05 100644 --- a/vllm/model_executor/models/moonvit.py +++ b/vllm/model_executor/models/moonvit.py @@ -66,6 +66,7 @@ from vllm.model_executor.models.vision import is_vit_use_data_parallel from vllm.platforms import current_platform from vllm.transformers_utils.configs.moonvit import MoonViTConfig +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d @@ -811,9 +812,14 @@ def forward( merge_kernel_size=self.merge_kernel_size, ) - hidden_states = self.patch_embed(pixel_values, grid_hw) - hidden_states = self.encoder(hidden_states, grid_hw) - hidden_states = patch_merger( - hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size - ) + # Legacy path: patch_embed, encoder and patch_merger each iterate the + # per-image grids in Python, so all three read `grid_hw` back to the + # host. The `encoder_metadata` path above precomputes that outside the + # graph and does not sync. + with gpu_sync_allowed(): + hidden_states = self.patch_embed(pixel_values, grid_hw) + hidden_states = self.encoder(hidden_states, grid_hw) + hidden_states = patch_merger( + hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size + ) return hidden_states diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index bef8057d2fb5..bfcebe7f9210 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -55,6 +55,7 @@ ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.repo_utils import get_hf_file_to_dict +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -835,15 +836,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: from .utils import PPMissingLayer @@ -1668,15 +1668,19 @@ def _process_audio_input( """ audio_data = audio_input["audio_data"] audio_data_seqlens = audio_input["audio_data_seqlens"] - last_hidden_state, deepstack = self.audio_encoder( - audio_data.to(self.audio_encoder.dtype), - feature_lens=audio_data_seqlens, - output_deepstack_hidden_states=len(self.deepstack_audio_merger_list) > 0, - ) - audio_embeds = self.audio_adapter(last_hidden_state) - audio_lengths = MossAudioEncoder._compute_downsampled_length( - audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long) - ).tolist() + # The encoder chunks the input by per-audio feature lengths, which + # needs Python ints for `split`/`pad_sequence`. + want_deepstack = len(self.deepstack_audio_merger_list) > 0 + with gpu_sync_allowed(): + last_hidden_state, deepstack = self.audio_encoder( + audio_data.to(self.audio_encoder.dtype), + feature_lens=audio_data_seqlens, + output_deepstack_hidden_states=want_deepstack, + ) + audio_embeds = self.audio_adapter(last_hidden_state) + audio_lengths = MossAudioEncoder._compute_downsampled_length( + audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long) + ).tolist() main_embeddings = tuple(audio_embeds.squeeze(0).split(audio_lengths, dim=0)) deepstack_embeddings: list[tuple[torch.Tensor, ...]] = [] diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py new file mode 100644 index 000000000000..74e000131822 --- /dev/null +++ b/vllm/model_executor/models/muse_glimmer.py @@ -0,0 +1,1649 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only MuseGlimmer multimodal model for vLLM. + +Native port of the MuseGlimmer text decoder (``MuseGlimmerForCausalLM``). The text +stack is a Gemma2 derivative with the +following MuseGlimmer-specific deltas, each of which is handled explicitly here: + + * SiLU-gated MLP (``hidden_activation="silu"``), not Gemma's gelu-tanh. + * Scaleless RMSNorm on the token embeddings (no sqrt(hidden) scaling). + * Per-layer sandwich RMSNorms with a baked ``+1`` weight offset + (``x * (1 + w)``), matching Gemma, but with distinct eps for the + pre/post norms (``rms_norm_eps`` vs ``post_norm_eps``). + * QK-norm (weightless, fp32) applied *before* RoPE, followed by a query + pre-scale of ``qk_scale_factor / sqrt(head_dim)``. + * A per-head sigmoid attention output gate. + * iRoPE layout: NoPE layers use full attention, RoPE layers use sliding + window attention. RoPE is applied NEOX-style (``is_neox_style=True``): + the HF converter (``convert_muse_glimmer_weights_to_hf.py``, 20260806+) permutes + q/k into the half-split (NEOX) layout via ``_permute_for_rope`` so they + pair with ``rotate_half`` — matching the reference's interleaved rotation + on the *native* (unpermuted) weights. Serving the permuted HF weights with + ``is_neox_style=False`` scrambles q/k and causes token-repetition collapse. + * Final logits are pre-scaled by ``output_multiplier`` and then tanh + soft-capped at ``final_logit_softcapping``. + * Untied lm_head. + +The vision path supports variable-resolution images and temporally patched +videos. It mirrors the checkpoint's native vision encoder, including sparse +block attention, 2-D RoPE, pixel-shuffle downsampling, and the two-layer +adapter/projection stack. +""" + +import math +from collections.abc import Iterable, Mapping, Sequence +from itertools import islice +from typing import Annotated, Literal + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from torch import nn +from transformers import BatchFeature + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.distributed import ( + divide, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention, MMEncoderAttention +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.renderers import TokenizeParams +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.processors.muse_glimmer import MuseGlimmerProcessor +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from .interfaces import ( + EagleModelMixin, + MultiModalEmbeddings, + SupportsEagle3, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from .vision import get_fp8_padded_hidden_size, is_vit_use_data_parallel + +logger = init_logger(__name__) + + +def _text_config(config): + """MuseGlimmer checkpoints may nest the text config under ``text_config`` + (multimodal ``MuseGlimmerConfig``) or expose it directly + (``MuseGlimmerTextConfig``).""" + return getattr(config, "text_config", config) + + +def _vision_config(config): + return getattr(config, "vision_config", config) + + +def _muse_glimmer_has_vision(config) -> bool: + has_vision = getattr(config, "has_vision", None) + if has_vision is not None: + return bool(has_vision) + return hasattr(config, "vision_config") + + +IMAGE_TOKEN = "<|patch|>" +IMAGE_PROCESSOR_TOKEN = "<|image|>" +VIDEO_TOKEN = "<|video|>" + + +class MuseGlimmerImagePixelInputs(TensorSchema): + """Batched variable-resolution image inputs.""" + + type: Literal["image_pixels"] + pixel_values: Annotated[ + torch.Tensor | list[torch.Tensor], + TensorShape("bn", 3, "h", "w", dynamic_dims={"h", "w"}), + ] + feature_sizes: Annotated[torch.Tensor, TensorShape("bn")] + + +class MuseGlimmerVideoPixelInputs(TensorSchema): + """Batched variable-length, variable-resolution video inputs.""" + + type: Literal["video_pixels"] + pixel_values: Annotated[ + torch.Tensor | list[torch.Tensor], + TensorShape( + "bn", + "ng", + "c", + "h", + "w", + dynamic_dims={"ng", "h", "w"}, + ), + ] + feature_sizes: Annotated[torch.Tensor, TensorShape("bn")] + + +class MuseGlimmerProcessingInfo(BaseProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + if not _muse_glimmer_has_vision(self.get_hf_config()): + return {} + return {"image": None, "video": None} + + def get_default_tok_params(self) -> TokenizeParams: + return super().get_default_tok_params().with_kwargs(add_special_tokens=False) + + def get_hf_processor(self, **kwargs: object) -> MuseGlimmerProcessor: + return self.ctx.init_processor( + MuseGlimmerProcessor, + tokenizer=self.get_tokenizer(), + **kwargs, + ) + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + processor = self.get_hf_processor() + image_processor = processor.image_processor + video_processor = processor.video_processor + num_frames = self.get_num_frames_with_most_features(seq_len, mm_counts) + return { + "image": int(image_processor.max_image_tokens), + "video": ( + num_frames + // int(video_processor.patch_temporal) + * int(video_processor.max_video_frame_tokens) + ), + } + + def get_image_size_with_most_features(self) -> ImageSize: + image_processor = self.get_hf_processor().image_processor + grid_size = math.isqrt(int(image_processor.max_image_tokens)) + side = ( + int(image_processor.patch_size) + * int(image_processor.downsample_factor) + * grid_size + ) + return ImageSize(width=side, height=side) + + def get_num_frames_with_most_features( + self, seq_len: int, mm_counts: Mapping[str, int] + ) -> int: + video_processor = self.get_hf_processor().video_processor + num_videos = max(mm_counts.get("video", 0), 1) + groups = max( + 1, + seq_len // num_videos // int(video_processor.max_video_frame_tokens), + ) + patch_temporal = int(video_processor.patch_temporal) + max_groups = max( + 1, + int(video_processor.video_num_frames) // patch_temporal, + ) + return min(groups, max_groups) * patch_temporal + + +class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return IMAGE_TOKEN * mm_counts.get("image", 0) + VIDEO_TOKEN * mm_counts.get( + "video", 0 + ) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + processor = self.info.get_hf_processor() + video_processor = processor.video_processor + image_width, image_height = self.info.get_image_size_with_most_features() + video_grid = math.isqrt(int(video_processor.max_video_frame_tokens)) + video_size = ( + int(video_processor.patch_size) + * int(video_processor.downsample_factor) + * video_grid + ) + return { + "image": self._get_dummy_images( + width=image_width, + height=image_height, + num_images=mm_counts.get("image", 0), + overrides=mm_options.get("image"), + ), + "video": self._get_dummy_videos( + width=video_size, + height=video_size, + num_frames=self.info.get_num_frames_with_most_features( + seq_len, mm_counts + ), + num_videos=mm_counts.get("video", 0), + overrides=mm_options.get("video"), + ), + } + + +class MuseGlimmerMultiModalProcessor( + BaseMultiModalProcessor[MuseGlimmerProcessingInfo] +): + def _apply_hf_processor_text_only( + self, + prompt_text: str, + tokenization_kwargs: Mapping[str, object], + ) -> list[int]: + tokenizer = self.info.get_tokenizer() + return tokenizer.encode( + prompt_text, + **{"add_special_tokens": False, **tokenization_kwargs}, + ) + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + processor = self.info.get_hf_processor(**mm_kwargs) + tokenizer = processor.tokenizer + config = self.info.get_hf_config() + images = mm_data.get("images", ()) + videos = mm_data.get("videos", ()) + if not isinstance(images, Sequence) or not isinstance(videos, Sequence): + raise TypeError("MuseGlimmer multi-modal data must be a sequence") + + prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) + image_sentinel_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN) + video_sentinel_id = tokenizer.convert_tokens_to_ids(VIDEO_TOKEN) + if prompt_ids.count(image_sentinel_id) != len(images): + raise ValueError("The number of image placeholders does not match images") + if prompt_ids.count(video_sentinel_id) != len(videos): + raise ValueError("The number of video placeholders does not match videos") + + image_pixels: list[torch.Tensor] = [] + image_sizes: list[int] = [] + image_prompt_blocks: list[list[int]] = [] + for image in images: + output = processor( + text=IMAGE_PROCESSOR_TOKEN, + images=[image], + return_tensors=None, + ) + block_ids = output["input_ids"][0] + if isinstance(block_ids, torch.Tensor): + block_ids = block_ids.tolist() + pixels = output["pixel_values"] + if isinstance(pixels, torch.Tensor): + pixels = [pixels] if pixels.ndim == 3 else list(pixels) + if len(pixels) != 1: + raise ValueError( + "MuseGlimmer HF processor must return one tensor per image" + ) + image_pixels.append(pixels[0]) + image_sizes.append(block_ids.count(config.image_token_id)) + image_prompt_blocks.append(block_ids) + + video_pixels: list[torch.Tensor] = [] + video_sizes: list[int] = [] + video_prompt_blocks: list[list[int]] = [] + for video in videos: + if not isinstance(video, np.ndarray): + raise TypeError("MuseGlimmer video input must be a NumPy array") + frames = [Image.fromarray(frame).convert("RGB") for frame in video] + output = processor( + text=VIDEO_TOKEN, + videos=[frames], + return_tensors=None, + ) + block_ids = output["input_ids"][0] + if isinstance(block_ids, torch.Tensor): + block_ids = block_ids.tolist() + pixels = output["pixel_values"] + if isinstance(pixels, torch.Tensor): + pixels = list(pixels) + video_pixels.append(torch.stack(pixels)) + video_sizes.append(block_ids.count(config.video_token_id)) + video_prompt_blocks.append(block_ids) + + combined_prompt_ids: list[int] = [] + image_idx = video_idx = 0 + for token_id in prompt_ids: + if token_id == image_sentinel_id: + combined_prompt_ids.extend(image_prompt_blocks[image_idx]) + image_idx += 1 + elif token_id == video_sentinel_id: + combined_prompt_ids.extend(video_prompt_blocks[video_idx]) + video_idx += 1 + else: + combined_prompt_ids.append(token_id) + + data: dict[str, object] = { + "input_ids": [combined_prompt_ids], + } + if image_pixels: + data.update( + image_pixel_values=image_pixels, + image_feature_sizes=torch.tensor(image_sizes), + ) + if video_pixels: + data.update( + video_pixel_values=video_pixels, + video_feature_sizes=torch.tensor(video_sizes), + ) + del tok_kwargs + return BatchFeature(data=data, tensor_type=None) + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields = {} + if "image_pixel_values" in hf_inputs: + fields.update( + image_pixel_values=MultiModalFieldConfig.batched("image"), + image_feature_sizes=MultiModalFieldConfig.batched("image"), + ) + if "video_pixel_values" in hf_inputs: + fields.update( + video_pixel_values=MultiModalFieldConfig.batched("video"), + video_feature_sizes=MultiModalFieldConfig.batched("video"), + ) + return fields + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + config = self.info.get_hf_config() + tokenizer = self.info.get_tokenizer() + vocab = tokenizer.get_vocab() + + placeholder = { + "image": vocab[IMAGE_TOKEN], + "video": vocab[VIDEO_TOKEN], + } + image_start_id = vocab["<|image_start|>"] + image_end_id = vocab["<|image_end|>"] + video_start_id = vocab["<|vid_start|>"] + video_end_id = vocab["<|vid_end|>"] + video_separator_id = vocab["<|vid_frame_separator|>"] + + def image_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]: + out_item = out_mm_kwargs["image"][item_idx] + num_tokens = int(out_item["image_feature_sizes"].data) + replacement = ( + [image_start_id] + [config.image_token_id] * num_tokens + [image_end_id] + ) + return PromptUpdateDetails.select_token_id( + replacement, config.image_token_id + ) + + def video_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]: + out_item = out_mm_kwargs["video"][item_idx] + pixel_values = out_item["video_pixel_values"].data + num_groups = len(pixel_values) + if num_groups == 0: + raise ValueError( + "MuseGlimmer video must contain at least one frame group" + ) + + num_tokens = int(out_item["video_feature_sizes"].data) + tokens_per_group, remainder = divmod(num_tokens, num_groups) + if remainder: + raise ValueError( + "MuseGlimmer video feature size must be divisible by frame groups" + ) + + video_processor = processor.video_processor + patch_temporal = int(video_processor.patch_temporal) + sampling_fps = float(video_processor.video_sampling_fps) + if sampling_fps <= 0: + raise ValueError("MuseGlimmer video_sampling_fps must be positive") + + replacement = [video_start_id] + for group_idx in range(num_groups): + timestamp = group_idx * patch_temporal / sampling_fps + replacement.extend( + tokenizer.encode( + f"Time: {timestamp:.1f}s", add_special_tokens=False + ) + ) + replacement.extend([config.video_token_id] * tokens_per_group) + replacement.append( + video_separator_id if group_idx < num_groups - 1 else video_end_id + ) + return PromptUpdateDetails.select_token_id( + replacement, config.video_token_id + ) + + return [ + PromptReplacement( + modality="image", + target=[placeholder["image"]], + replacement=image_replacement, + ), + PromptReplacement( + modality="video", + target=[placeholder["video"]], + replacement=video_replacement, + ), + ] + + +def _muse_glimmer_use_qk_norm(config) -> bool: + """Whether QK-norm is applied. MuseGlimmer ALWAYS applies QK-norm; the modular HF + ``text_config`` schema simply omits ``use_qk_norm`` (reads as ``None``). + Treat a missing/None flag as True — only an explicit ``False`` disables it.""" + val = getattr(config, "use_qk_norm", None) + return True if val is None else bool(val) + + +def _muse_glimmer_use_attn_output_gate(config) -> bool: + """Whether the per-head sigmoid attention output gate is applied. MuseGlimmer ALWAYS + applies it; the modular HF ``text_config`` omits ``use_attn_output_gate`` + (reads as ``None``). Missing/None -> True; only explicit ``False`` disables.""" + val = getattr(config, "use_attn_output_gate", None) + return True if val is None else bool(val) + + +def _muse_glimmer_query_prescale(config) -> float: + """Post-QK-norm query pre-scale (``scale_query_by``), normalized across the + two config schemas so the net query scaling matches the native reference. + + HF native modeling computes ``scale_query_by = qk_scale_factor / sqrt(head_dim)`` + where the NATIVE ``qk_scale_factor`` is the raw ``params.json`` value + (~43.784). The modular HF ``text_config`` PRE-FOLDS the ``1/sqrt(head_dim)`` + factor and ships ``qk_scale_factor = 43.784 / sqrt(128) = 3.87`` already, + expecting it applied directly. Both must yield the SAME ``scale_query_by`` + (~3.87), then softmax uses ``scaling = head_dim**-0.5``. + + Precedence: + 1. explicit ``scale_query_by`` (already the final factor) -> use as-is. + 2. else derive from ``qk_scale_factor``: + - if it is already the folded value (``~= qk_scale_factor/sqrt(hd)`` is + NOT what we want; detect the native form and divide) — we decide by + magnitude: the native raw value is ``folded * sqrt(head_dim)``. If + ``qk_scale_factor`` is close to ``folded_expected * sqrt(hd)`` we treat + it as native and divide; otherwise it is already folded, use directly. + """ + head_dim = config.head_dim + sqrt_hd = head_dim**0.5 + + explicit = getattr(config, "scale_query_by", None) + if explicit is not None: + return float(explicit) + + qk_scale = getattr(config, "qk_scale_factor", None) + if qk_scale is None: + # No scale info at all: fall back to the plain 1/sqrt(head_dim) identity + # (net query scaling then just the softmax scaling); should not happen + # for real MuseGlimmer checkpoints, which always carry qk_scale_factor. + return 1.0 + + qk_scale = float(qk_scale) + # Disambiguate native (raw, ~43.78) vs modular (folded, ~3.87). The native + # form, when divided by sqrt(head_dim), yields the folded target; the folded + # form is already the target. Native values are ~sqrt(head_dim)x larger than + # folded. Use a threshold at sqrt(head_dim) (with margin): if qk_scale is + # comparable to or larger than sqrt(head_dim), it is the native raw value and + # must be divided; otherwise it is already folded and used directly. + # head_dim=128 -> sqrt=11.31; native 43.78 > 11.31 (divide -> 3.87), + # folded 3.87 < 11.31 (use as-is). + if qk_scale >= sqrt_hd: + return qk_scale / sqrt_hd + return qk_scale + + +class MuseGlimmerRMSNorm(nn.Module): + """RMSNorm mirroring HF MuseGlimmer exactly (fp32 compute, cast at the end). + + ``normed = _norm(x.float()) * (w.float() + weight_offset)`` cast back to the + input dtype. When ``with_scale`` is False the layer is weightless (used for + QK-norm and the token-embedding norm). + """ + + def __init__( + self, + dim: int | None = None, + eps: float = 1e-6, + with_scale: bool = True, + weight_offset: int = 0, + ) -> None: + super().__init__() + self.eps = eps + self.with_scale = with_scale + self.weight_offset = weight_offset + if with_scale: + assert dim is not None + self.weight = nn.Parameter(torch.zeros(dim)) + else: + self.register_parameter("weight", None) + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + out = self._norm(hidden_states.float()) + if self.with_scale: + out = out * (self.weight.float() + self.weight_offset) + return out.type_as(hidden_states) + + +class MuseGlimmerVisionAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + prefix: str = "", + ) -> None: + super().__init__() + if hidden_size % num_heads: + raise ValueError("MuseGlimmer vision hidden size must divide num heads") + + self.total_num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.use_data_parallel = is_vit_use_data_parallel(num_heads) + self.tp_size = ( + 1 if self.use_data_parallel else get_tensor_model_parallel_world_size() + ) + self.num_heads = divide(num_heads, self.tp_size) + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=num_heads, + total_num_kv_heads=num_heads, + bias=True, + prefix=f"{prefix}.qkv_proj", + disable_tp=self.use_data_parallel, + ) + self.o_proj = RowParallelLinear( + input_size=hidden_size, + output_size=hidden_size, + bias=True, + prefix=f"{prefix}.o_proj", + disable_tp=self.use_data_parallel, + ) + self.attn = MMEncoderAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + scale=self.head_dim**-0.5, + prefix=f"{prefix}.attn", + ) + self.apply_rotary_emb = ApplyRotaryEmb( + enforce_enable=True, + is_neox_style=True, + enable_fp32_compute=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + sequence_lengths: torch.Tensor | None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + qkv = qkv.view( + hidden_states.shape[0], + 3, + self.num_heads, + self.head_dim, + ) + query, key, value = qkv.unbind(1) + query, key = self.apply_rotary_emb( + torch.stack([query, key]).contiguous(), + rotary_pos_emb_cos, + rotary_pos_emb_sin, + ).unbind(0) + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + + output = self.attn( + query=query, + key=key, + value=value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + output = output.reshape(hidden_states.shape[0], -1) + output, _ = self.o_proj(output) + return output + + +class MuseGlimmerVisionMLP(nn.Module): + def __init__(self, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.c_fc = nn.Linear(hidden_size, intermediate_size, bias=True) + self.c_proj = nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.c_proj(F.gelu(self.c_fc(hidden_states))) + + +class MuseGlimmerVisionBlock(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + intermediate_size: int, + prefix: str = "", + ) -> None: + super().__init__() + self.ln_1 = nn.LayerNorm(hidden_size) + self.attn = MuseGlimmerVisionAttention( + hidden_size, + num_heads, + prefix=f"{prefix}.attn", + ) + self.ln_2 = nn.LayerNorm(hidden_size) + self.mlp = MuseGlimmerVisionMLP(hidden_size, intermediate_size) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: torch.Tensor, + sequence_lengths: torch.Tensor | None, + ) -> torch.Tensor: + batch_size, seq_len, hidden_size = hidden_states.shape + flattened = hidden_states.view(batch_size * seq_len, hidden_size) + flattened = flattened + self.attn( + self.ln_1(flattened), + rotary_pos_emb_cos, + rotary_pos_emb_sin, + cu_seqlens, + max_seqlen, + sequence_lengths, + ) + flattened = flattened + self.mlp(self.ln_2(flattened)) + return flattened.view(batch_size, seq_len, hidden_size) + + +class MuseGlimmerVisionEncoder(nn.Module): + def __init__(self, config, prefix: str = "") -> None: + super().__init__() + hidden_size = config.hidden_size + patch_dim = config.patch_temporal * 3 * config.patch_size**2 + self.hidden_size = hidden_size + self.patch_size = config.patch_size + self.patch_temporal = config.patch_temporal + self.merge_kernel_size = config.merge_kernel_size + self.pos_emb_height = config.pos_emb_height + self.pos_emb_width = config.pos_emb_width + self.head_dim = hidden_size // config.num_attention_heads + self.layer_types = list(config.layer_types) + if len(self.layer_types) != config.num_hidden_layers: + raise ValueError( + "MuseGlimmer vision layer_types must match num_hidden_layers" + ) + if self.head_dim % 4: + raise ValueError("MuseGlimmer vision head dimension must be divisible by 4") + + self.conv1_linear = nn.Linear(patch_dim, hidden_size, bias=False) + self.positional_embedding_vlm = nn.Parameter( + torch.zeros(config.pos_emb_height * config.pos_emb_width, hidden_size) + ) + self.ln_pre = nn.LayerNorm(hidden_size) + self.transformer = nn.ModuleList( + [ + MuseGlimmerVisionBlock( + hidden_size, + config.num_attention_heads, + config.intermediate_size, + prefix=f"{prefix}.transformer.{layer_idx}", + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + vision_attention = self.transformer[0].attn + self.tp_size = vision_attention.tp_size + self.attn_backend = vision_attention.attn.attn_backend + self.fp8_padded_hidden_size = get_fp8_padded_hidden_size( + config.num_attention_heads, self.head_dim + ) + self.ln_post = nn.LayerNorm(hidden_size) + + expected_output = hidden_size * self.merge_kernel_size**2 + if config.output_dim != expected_output: + raise ValueError( + f"MuseGlimmer vision output_dim={config.output_dim} does not match " + f"pixel-shuffle output {expected_output}" + ) + + def _make_2d_rope( + self, grid_height: int, grid_width: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor]: + spatial_dim = self.head_dim // 2 + inv_freq = 1.0 / ( + 10000.0 + ** ( + torch.arange(0, spatial_dim, 2, dtype=torch.float32, device=device) + / spatial_dim + ) + ) + height = torch.arange(1, grid_height + 1, dtype=torch.float32, device=device) + width = torch.arange(1, grid_width + 1, dtype=torch.float32, device=device) + height = height.unsqueeze(1).expand(-1, grid_width).reshape(-1) + width = width.unsqueeze(0).expand(grid_height, -1).reshape(-1) + freq_w = torch.outer(width, inv_freq) + freq_h = torch.outer(height, inv_freq) + freqs = torch.cat([freq_w, freq_h], dim=-1) + return torch.cos(freqs), torch.sin(freqs) + + def _get_pos_emb( + self, grid_height: int, grid_width: int, device: torch.device + ) -> torch.Tensor: + h_grid = ( + torch.arange(grid_height, device=device, dtype=torch.float32) + 0.5 + ) * (self.pos_emb_height / grid_height) - 0.5 + w_grid = ( + torch.arange(grid_width, device=device, dtype=torch.float32) + 0.5 + ) * (self.pos_emb_width / grid_width) - 0.5 + h_floor = torch.floor(h_grid).long() + w_floor = torch.floor(w_grid).long() + h_ceil = h_floor + 1 + w_ceil = w_floor + 1 + h_frac = h_grid - h_floor.float() + w_frac = w_grid - w_floor.float() + + h_floor_valid = (h_floor >= 0) & (h_floor < self.pos_emb_height) + h_ceil_valid = (h_ceil >= 0) & (h_ceil < self.pos_emb_height) + w_floor_valid = (w_floor >= 0) & (w_floor < self.pos_emb_width) + w_ceil_valid = (w_ceil >= 0) & (w_ceil < self.pos_emb_width) + h_floor = h_floor.clamp(0, self.pos_emb_height - 1) + h_ceil = h_ceil.clamp(0, self.pos_emb_height - 1) + w_floor = w_floor.clamp(0, self.pos_emb_width - 1) + w_ceil = w_ceil.clamp(0, self.pos_emb_width - 1) + + h_floor_offset = h_floor * self.pos_emb_width + h_ceil_offset = h_ceil * self.pos_emb_width + indices = torch.stack( + [ + (h_floor_offset[:, None] + w_floor[None, :]).flatten(), + (h_floor_offset[:, None] + w_ceil[None, :]).flatten(), + (h_ceil_offset[:, None] + w_floor[None, :]).flatten(), + (h_ceil_offset[:, None] + w_ceil[None, :]).flatten(), + ] + ) + weights = torch.stack( + [ + ( + (1 - h_frac)[:, None] + * (1 - w_frac)[None, :] + * (h_floor_valid[:, None] & w_floor_valid[None, :]) + ).flatten(), + ( + (1 - h_frac)[:, None] + * w_frac[None, :] + * (h_floor_valid[:, None] & w_ceil_valid[None, :]) + ).flatten(), + ( + h_frac[:, None] + * (1 - w_frac)[None, :] + * (h_ceil_valid[:, None] & w_floor_valid[None, :]) + ).flatten(), + ( + h_frac[:, None] + * w_frac[None, :] + * (h_ceil_valid[:, None] & w_ceil_valid[None, :]) + ).flatten(), + ] + ) + return (self.positional_embedding_vlm[indices] * weights[..., None]).sum(0) + + def _pixel_shuffle_downsample( + self, hidden_states: torch.Tensor, grid_height: int, grid_width: int + ) -> torch.Tensor: + factor = self.merge_kernel_size + output_tokens = (grid_height // factor) * (grid_width // factor) + permutation = torch.arange( + grid_height * grid_width, device=hidden_states.device + ) + permutation = permutation.view( + grid_height // factor, factor, grid_width // factor, factor + ) + permutation = permutation.permute(0, 2, 1, 3).reshape(-1) + hidden_states = hidden_states.squeeze(0)[permutation] + hidden_size = hidden_states.shape[-1] + hidden_states = ( + hidden_states.view(output_tokens, factor * factor, hidden_size) + .permute(0, 2, 1) + .contiguous() + .view(output_tokens, hidden_size * factor * factor) + ) + return hidden_states.unsqueeze(0) + + def _get_sparse_permutation( + self, grid_height: int, grid_width: int, device: torch.device + ) -> tuple[torch.Tensor, list[int]]: + block_height = self.pos_emb_height + block_width = self.pos_emb_width + padded_height = math.ceil(grid_height / block_height) * block_height + padded_width = math.ceil(grid_width / block_width) * block_width + indices = torch.arange(grid_height * grid_width, device=device).view( + grid_height, grid_width + ) + indices = F.pad( + indices, + (0, padded_width - grid_width, 0, padded_height - grid_height), + value=-1, + ).flatten() + indices = indices.view( + padded_height // block_height, + block_height, + padded_width // block_width, + block_width, + ) + indices = indices.permute(0, 2, 1, 3).reshape(-1) + valid = (indices != -1).view(-1, block_height * block_width) + return indices[indices != -1], valid.sum(dim=1).tolist() + + def _get_attention_metadata( + self, + seq_lens: Sequence[int], + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + cu_seqlens_np = np.concatenate( + [ + np.zeros(1, dtype=np.int32), + np.asarray(seq_lens, dtype=np.int32).cumsum(dtype=np.int32), + ] + ) + sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens( + self.attn_backend, + cu_seqlens_np, + device, + ) + max_seqlen = torch.tensor( + MMEncoderAttention.compute_max_seqlen( + self.attn_backend, + cu_seqlens_np, + ), + dtype=torch.int32, + ) + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens( + self.attn_backend, + cu_seqlens_np, + self.hidden_size, + self.tp_size, + device, + fp8_padded_hidden_size=self.fp8_padded_hidden_size, + ) + return cu_seqlens, max_seqlen, sequence_lengths + + def _patchify(self, pixels: torch.Tensor) -> torch.Tensor: + patch_size = self.patch_size + _, channels, height, width = pixels.shape + grid_height = height // patch_size + grid_width = width // patch_size + if channels == 3: + patches = pixels.unfold(2, patch_size, patch_size).unfold( + 3, patch_size, patch_size + ) + patches = patches.contiguous().view( + 1, channels, grid_height, grid_width, patch_size, patch_size + ) + patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous() + patches = patches.unsqueeze(3).expand( + -1, -1, -1, self.patch_temporal, -1, -1, -1 + ) + elif channels == self.patch_temporal * 3: + frame_patches = [] + for frame_idx in range(self.patch_temporal): + frame = pixels[:, frame_idx * 3 : (frame_idx + 1) * 3] + frame = frame.unfold(2, patch_size, patch_size).unfold( + 3, patch_size, patch_size + ) + frame = frame.contiguous().view( + 1, 3, grid_height, grid_width, patch_size, patch_size + ) + frame_patches.append(frame.permute(0, 2, 3, 1, 4, 5).contiguous()) + patches = torch.stack(frame_patches, dim=3) + else: + raise ValueError( + f"MuseGlimmer vision input has {channels} channels; expected 3 or " + f"{self.patch_temporal * 3}" + ) + return patches.reshape(1, grid_height * grid_width, -1) + + def forward(self, pixel_values: Sequence[torch.Tensor]) -> torch.Tensor: + device = self.conv1_linear.weight.device + dtype = self.conv1_linear.weight.dtype + has_sparse_layers = any( + layer_type != "full_attention" for layer_type in self.layer_types + ) + + all_hidden_states = [] + all_rotary_pos_emb_cos = [] + all_rotary_pos_emb_sin = [] + sparse_seq_lens: list[int] = [] + global_seq_lens: list[int] = [] + metadata = [] + for pixels in pixel_values: + pixels = pixels.to(device=device, dtype=dtype) + if pixels.ndim == 3: + pixels = pixels.unsqueeze(0) + if pixels.ndim != 4 or pixels.shape[0] != 1: + raise ValueError( + "Each MuseGlimmer vision input must have shape [C, H, W]" + ) + spatial_stride = self.patch_size * self.merge_kernel_size + if pixels.shape[-2] % spatial_stride or pixels.shape[-1] % spatial_stride: + raise ValueError( + "MuseGlimmer vision input dimensions must divide the " + "merged patch stride" + ) + grid_height = pixels.shape[-2] // self.patch_size + grid_width = pixels.shape[-1] // self.patch_size + num_tokens = grid_height * grid_width + hidden_states = self.conv1_linear(self._patchify(pixels)) + hidden_states = hidden_states + self._get_pos_emb( + grid_height, grid_width, device + ).unsqueeze(0).to(dtype) + hidden_states = self.ln_pre(hidden_states.view(-1, self.hidden_size)).view( + 1, -1, self.hidden_size + ) + rotary_pos_emb_cos, rotary_pos_emb_sin = self._make_2d_rope( + grid_height, grid_width, device + ) + + permutation = None + if has_sparse_layers: + permutation, seq_lens = self._get_sparse_permutation( + grid_height, grid_width, device + ) + hidden_states = hidden_states[:, permutation] + rotary_pos_emb_cos = rotary_pos_emb_cos[permutation] + rotary_pos_emb_sin = rotary_pos_emb_sin[permutation] + sparse_seq_lens.extend(seq_lens) + + all_hidden_states.append(hidden_states.squeeze(0)) + all_rotary_pos_emb_cos.append(rotary_pos_emb_cos) + all_rotary_pos_emb_sin.append(rotary_pos_emb_sin) + global_seq_lens.append(num_tokens) + metadata.append((grid_height, grid_width, num_tokens, permutation)) + + hidden_states = torch.cat(all_hidden_states).unsqueeze(0) + rotary_pos_emb_cos = torch.cat(all_rotary_pos_emb_cos) + rotary_pos_emb_sin = torch.cat(all_rotary_pos_emb_sin) + global_attention_metadata = self._get_attention_metadata( + global_seq_lens, + device, + ) + sparse_attention_metadata = ( + self._get_attention_metadata(sparse_seq_lens, device) + if sparse_seq_lens + else None + ) + for layer_type, block in zip(self.layer_types, self.transformer): + attention_metadata = ( + global_attention_metadata + if layer_type == "full_attention" + else sparse_attention_metadata + ) + if attention_metadata is None: + raise ValueError("MuseGlimmer sparse attention metadata is missing") + hidden_states = block( + hidden_states, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + *attention_metadata, + ) + + features = [] + offset = 0 + for grid_height, grid_width, num_tokens, permutation in metadata: + item = hidden_states[:, offset : offset + num_tokens] + offset += num_tokens + if permutation is not None: + inverse = torch.empty_like(permutation) + inverse[permutation] = torch.arange(len(permutation), device=device) + item = item[:, inverse] + item = self.ln_post(item.view(-1, self.hidden_size)).view( + 1, -1, self.hidden_size + ) + features.append( + self._pixel_shuffle_downsample(item, grid_height, grid_width).squeeze(0) + ) + return torch.cat(features) + + +class MuseGlimmerVisionAdapter(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.c_fc = nn.Linear(config.output_dim, config.adapter_dim, bias=False) + self.c_proj = nn.Linear(config.adapter_dim, config.adapter_dim, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.gelu(self.c_proj(F.gelu(self.c_fc(hidden_states)))) + + +class MuseGlimmerMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_activation: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_activation != "silu": + raise ValueError( + f"MuseGlimmer uses `silu` as the hidden activation; " + f"got `{hidden_activation}`." + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MuseGlimmerAttention(nn.Module): + def __init__( + self, + config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.layer_idx = extract_layer_index(prefix) + + tp_size = get_tensor_model_parallel_world_size() + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + # MuseGlimmer overrides Gemma2's query_pre_attn_scalar scaling with the + # standard 1/sqrt(head_dim). The query is *additionally* pre-scaled by + # scale_query_by after QK-norm (see forward). + self.scaling = self.head_dim**-0.5 + + # iRoPE: NoPE layers (no_rope_layers[i] == 0) run full attention; RoPE + # layers run sliding-window attention. + self.use_rope = config.no_rope_layers[self.layer_idx] == 1 + + self.use_qk_norm = _muse_glimmer_use_qk_norm(config) + if self.use_qk_norm: + # Weightless, computed in fp32, applied per head over head_dim. + self.qk_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False) + # Post-QK-norm query pre-scale, normalized across the native (raw + # ~43.78) and modular (pre-folded ~3.87) config schemas. + self.scale_query_by = _muse_glimmer_query_prescale(config) + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.use_output_gate = _muse_glimmer_use_attn_output_gate(config) + if self.use_output_gate: + self.output_gate_proj = ColumnParallelLinear( + self.hidden_size, + self.total_num_heads * self.head_dim, + bias=False, + gather_output=False, + quant_config=quant_config, + prefix=f"{prefix}.output_gate_proj", + ) + + self.rotary_emb = ( + get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters, + # HF converter permutes q/k to NEOX (half-split) + # layout via _permute_for_rope + is_neox_style=True, + ) + if self.use_rope + else None + ) + + # Full attention on NoPE layers, sliding window otherwise. + sliding_window = None if not self.use_rope else config.sliding_window + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + logits_soft_cap=None, # MuseGlimmer sets attn_logit_softcapping = None + per_layer_sliding_window=sliding_window, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.use_qk_norm: + # QK-norm over head_dim, fp32, applied BEFORE RoPE; then pre-scale q. + q = q.reshape(-1, self.head_dim) + q = self.qk_norm(q).reshape(-1, self.q_size) * self.scale_query_by + k = k.reshape(-1, self.head_dim) + k = self.qk_norm(k).reshape(-1, self.kv_size) + q = q.to(v.dtype) + k = k.to(v.dtype) + + if self.rotary_emb is not None: + q, k = self.rotary_emb(positions, q, k) + + attn_output = self.attn(q, k, v) + + if self.use_output_gate: + # Gate reads the layer input hidden states (not the attn output). + gate, _ = self.output_gate_proj(hidden_states) + attn_output = torch.sigmoid(gate) * attn_output + + output, _ = self.o_proj(attn_output) + return output + + +class MuseGlimmerDecoderLayer(nn.Module): + def __init__( + self, + config, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = MuseGlimmerAttention( + config=config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp = MuseGlimmerMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_activation=config.hidden_activation, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + # Sandwich norms with baked +1 offset. Pre-norms use rms_norm_eps; the + # post-norms use the (typically smaller) post_norm_eps. + self.input_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.rms_norm_eps, weight_offset=1 + ) + self.post_attention_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.post_norm_eps, weight_offset=1 + ) + self.pre_feedforward_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.rms_norm_eps, weight_offset=1 + ) + self.post_feedforward_layernorm = MuseGlimmerRMSNorm( + config.hidden_size, eps=config.post_norm_eps, weight_offset=1 + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Gemma2-style sandwich, replicated explicitly (matches HF MuseGlimmer). + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + return hidden_states, residual + + +@support_torch_compile +class MuseGlimmerModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = _text_config(vllm_config.model_config.hf_config) + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + # MuseGlimmer normalizes token embeddings with a weightless RMSNorm instead of + # Gemma's sqrt(hidden_size) multiplier. + self.embed_norm = MuseGlimmerRMSNorm(eps=config.rms_norm_eps, with_scale=False) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: MuseGlimmerDecoderLayer( + config, cache_config, quant_config, prefix=prefix + ), + prefix=f"{prefix}.layers", + ) + # Final norm: weight-as-scale, no offset. + self.norm = MuseGlimmerRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_norm(self.embed_tokens(input_ids)) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + aux_hidden_states = self._maybe_add_hidden_state( + [], self.start_layer, hidden_states, None + ) + for layer_idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + hidden_states, residual = layer(positions, hidden_states, residual) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, None + ) + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states = self.norm(hidden_states) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + +@MULTIMODAL_REGISTRY.register_processor( + MuseGlimmerMultiModalProcessor, + info=MuseGlimmerProcessingInfo, + dummy_inputs=MuseGlimmerDummyInputsBuilder, +) +class MuseGlimmerForCausalLM( + nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + # Weight-name normalization. Two checkpoint conventions are supported: + # + # * HF MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``): the + # multimodal ``MuseGlimmerConfig`` prefixes the language model with + # ``model.language_model.``; the per-layer sandwich norms are already + # named ``input_layernorm`` / ``post_attention_layernorm`` / + # ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm``. + # + # * Legacy HF export (an earlier checkpoint convention): uses ``model.`` + # and a different sandwich-norm naming where ``post_attn_norm`` is the + # true post-attention norm and ``post_attention_layernorm`` is actually + # the pre-feedforward norm. We remap those to MuseGlimmer's names. + # + # CONVENTION DISAMBIGUATION (critical): the two checkpoint families use + # DIFFERENT sandwich-norm names, and they must not be conflated: + # + # * Canonical MuseGlimmer export (current + # ``convert_muse_glimmer_weights_to_hf.py`` — what partners ship): + # keys are ``model.language_model.layers.N.*`` and + # the norms are ALREADY named ``input_layernorm`` / + # ``post_attention_layernorm`` / ``pre_feedforward_layernorm`` / + # ``post_feedforward_layernorm``. No norm rename needed — pass through. + # + # * Legacy HF export (an earlier checkpoint convention): keys are + # ``model.layers.N.*`` and the sandwich norms are + # named ``input_layernorm`` / ``post_attention_layernorm`` (this one is + # actually the PRE-feedforward norm) / ``post_attn_norm`` (the true + # post-attention norm) / ``post_ffn_norm``. These must be remapped. + # + # The unambiguous discriminator is the PREFIX: legacy keys start with + # ``model.layers.`` while canonical keys start with + # ``model.language_model.layers.``. ``orig_to_new_regex`` runs BEFORE the + # prefix strip (see WeightsMapper._map_name_with_shard), so we anchor the + # legacy renames on ``^model\.layers\.`` — they fire ONLY on legacy keys and + # leave canonical/partner checkpoints untouched. Rule order within the regex + # dict matters: the ``post_attention_layernorm`` -> ``pre_feedforward_...`` + # rule must precede the ``post_attn_norm`` -> ``post_attention_layernorm`` + # rule so the latter's output is not re-captured by the former (regex rules + # apply as a single forward pass). + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + "model.vision_tower.patch_embedder.position_embedding_table.weight": ( + "model.vision_tower.positional_embedding_vlm" + ), + "model.vision_tower.layers.": "model.vision_tower.transformer.", + ".norm1.": ".ln_1.", + ".norm2.": ".ln_2.", + ".attn.proj.": ".attn.o_proj.", + ".mlp.fc1.": ".mlp.c_fc.", + ".mlp.fc2.": ".mlp.c_proj.", + ".self_attn.gate_proj": ".self_attn.output_gate_proj", + }, + orig_to_new_prefix={ + "model.rotary_emb.": None, + "model.language_model.": "model.", + "language_model.": "model.", + "model.vision_tower.patch_embedder.patch_embedding.": ( + "model.vision_tower.conv1_linear." + ), + "model.vision_tower.": "vision_encoder.", + "vision_tower.": "vision_encoder.", + "model.vision_encoder.": "vision_encoder.", + "model.vision_adapter.fc1.": "model.vision_adapter.c_fc.", + "model.vision_adapter.fc2.": "model.vision_adapter.c_proj.", + "model.vision_adapter.": "vision_adapter.", + "model.vision_projection.": "vision_projection.", + "model.perception_emb_norm.": "perception_emb_norm.", + }, + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + }, + ) + + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return IMAGE_TOKEN + if modality.startswith("video"): + return VIDEO_TOKEN + raise ValueError(f"Unsupported MuseGlimmer modality: {modality}") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + text_config = _text_config(config) + vision_config = _vision_config(config) + quant_config = vllm_config.quant_config + self.config = config + self.text_config = text_config + self.quant_config = quant_config + self.has_vision = _muse_glimmer_has_vision(config) + + with self._mark_language_model(vllm_config): + self.model = MuseGlimmerModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + if self.has_vision: + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.vision_encoder = MuseGlimmerVisionEncoder( + vision_config, + prefix=maybe_prefix(prefix, "vision_encoder"), + ) + self.vision_adapter = MuseGlimmerVisionAdapter(vision_config) + self.vision_projection = nn.Linear( + vision_config.adapter_dim, + text_config.hidden_size, + bias=False, + ) + self.perception_emb_norm = ( + MuseGlimmerRMSNorm(eps=text_config.rms_norm_eps, with_scale=False) + if text_config.normalize_tok_embeddings + else nn.Identity() + ) + else: + self.vision_encoder = None + self.vision_adapter = None + self.vision_projection = None + self.perception_emb_norm = None + + self.lm_head = ParallelLMHead( + text_config.vocab_size, + text_config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.output_multiplier = text_config.output_multiplier + self.final_logit_softcapping = text_config.final_logit_softcapping + self.logits_processor = LogitsProcessor(text_config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + image_token_id = getattr( + config, "image_token_id", getattr(config, "patch_token_id", 200092) + ) + video_token_id = getattr(config, "video_token_id", 200091) + self.configure_mm_token_handling( + text_config.vocab_size, [image_token_id, video_token_id] + ) + + def _parse_and_validate_image_input( + self, + **kwargs: object, + ) -> MuseGlimmerImagePixelInputs | None: + pixel_values = kwargs.pop("image_pixel_values", None) + feature_sizes = kwargs.pop("image_feature_sizes", None) + if pixel_values is None and feature_sizes is None: + return None + if pixel_values is None or feature_sizes is None: + raise ValueError( + "MuseGlimmer image_pixel_values and image_feature_sizes " + "must be provided together" + ) + if not isinstance(feature_sizes, torch.Tensor): + raise ValueError("MuseGlimmer image_feature_sizes must be a tensor") + if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + return MuseGlimmerImagePixelInputs( + type="image_pixels", + pixel_values=pixel_values, + feature_sizes=feature_sizes.reshape(-1), + ) + + def _parse_and_validate_video_input( + self, + **kwargs: object, + ) -> MuseGlimmerVideoPixelInputs | None: + pixel_values = kwargs.pop("video_pixel_values", None) + feature_sizes = kwargs.pop("video_feature_sizes", None) + if pixel_values is None and feature_sizes is None: + return None + if pixel_values is None or feature_sizes is None: + raise ValueError( + "MuseGlimmer video_pixel_values and video_feature_sizes " + "must be provided together" + ) + if not isinstance(feature_sizes, torch.Tensor): + raise ValueError("MuseGlimmer video_feature_sizes must be a tensor") + if isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 4: + pixel_values = pixel_values.unsqueeze(0) + patch_temporal = int(_vision_config(self.config).patch_temporal) + return MuseGlimmerVideoPixelInputs( + type="video_pixels", + pixel_values=pixel_values, + feature_sizes=feature_sizes.reshape(-1), + resolve_bindings={"c": patch_temporal * 3}, + ) + + def _encode_pixel_groups( + self, + pixel_groups: Sequence[torch.Tensor], + feature_sizes: Sequence[int], + ) -> tuple[torch.Tensor, ...]: + if ( + self.vision_encoder is None + or self.vision_adapter is None + or self.vision_projection is None + or self.perception_emb_norm is None + ): + raise ValueError("This MuseGlimmer checkpoint has no vision tower") + features = self.vision_encoder(pixel_groups) + features = self.vision_adapter(features) + features = self.vision_projection(features) + features = self.perception_emb_norm(features) + if features.shape[0] != sum(feature_sizes): + raise ValueError( + f"MuseGlimmer produced {features.shape[0]} vision features for " + f"{sum(feature_sizes)} placeholder tokens" + ) + return features.split(list(feature_sizes)) + + def _process_image_input( + self, + image_input: MuseGlimmerImagePixelInputs, + ) -> tuple[torch.Tensor, ...]: + images = list(image_input["pixel_values"]) + sizes = [int(size) for size in image_input["feature_sizes"].tolist()] + if len(images) != len(sizes): + raise ValueError("MuseGlimmer image batch metadata does not match pixels") + return self._encode_pixel_groups(images, sizes) + + def _process_video_input( + self, + video_input: MuseGlimmerVideoPixelInputs, + ) -> tuple[torch.Tensor, ...]: + videos = list(video_input["pixel_values"]) + sizes = [int(size) for size in video_input["feature_sizes"].tolist()] + if len(videos) != len(sizes): + raise ValueError("MuseGlimmer video batch metadata does not match pixels") + groups = [group for video in videos for group in video] + return self._encode_pixel_groups(groups, sizes) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + video_input = self._parse_and_validate_video_input(**kwargs) + + embeddings: MultiModalEmbeddings = [] + for key in kwargs: + if key == "image_pixel_values" and image_input is not None: + embeddings.extend(self._process_image_input(image_input)) + elif key == "video_pixel_values" and video_input is not None: + embeddings.extend(self._process_video_input(video_input)) + return embeddings + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + logits = logits * self.output_multiplier + if self.final_logit_softcapping is not None: + cap = self.final_logit_softcapping + logits = cap * torch.tanh(logits / cap) + return logits + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/nemotron.py b/vllm/model_executor/models/nemotron.py index 6f0b61205b3c..97199c184111 100644 --- a/vllm/model_executor/models/nemotron.py +++ b/vllm/model_executor/models/nemotron.py @@ -404,7 +404,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/olmo3.py b/vllm/model_executor/models/olmo3.py index 922834a8ee68..9ac4e61cf47d 100644 --- a/vllm/model_executor/models/olmo3.py +++ b/vllm/model_executor/models/olmo3.py @@ -369,15 +369,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = Olmo3Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=vllm_config.quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/olmo_hybrid.py b/vllm/model_executor/models/olmo_hybrid.py index 51bc410363ce..122a1c1c5b25 100644 --- a/vllm/model_executor/models/olmo_hybrid.py +++ b/vllm/model_executor/models/olmo_hybrid.py @@ -397,15 +397,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=vllm_config.quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( diff --git a/vllm/model_executor/models/openpangu.py b/vllm/model_executor/models/openpangu.py index 3557bdd0a56c..4a2030cfddc8 100644 --- a/vllm/model_executor/models/openpangu.py +++ b/vllm/model_executor/models/openpangu.py @@ -1122,7 +1122,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) diff --git a/vllm/model_executor/models/opt.py b/vllm/model_executor/models/opt.py index 32bb532f5c5b..d78e51e3dde5 100644 --- a/vllm/model_executor/models/opt.py +++ b/vllm/model_executor/models/opt.py @@ -349,14 +349,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model = OPTModel( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.word_embed_proj_dim, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.config.tie_word_embeddings: - self.lm_head = self.model.decoder.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.word_embed_proj_dim, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.decoder.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/orion.py b/vllm/model_executor/models/orion.py index 0871c347ac5c..2b58ee1e015c 100644 --- a/vllm/model_executor/models/orion.py +++ b/vllm/model_executor/models/orion.py @@ -305,7 +305,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/paddleocr_vl.py b/vllm/model_executor/models/paddleocr_vl.py index 0dae22115d43..bd635aaada2d 100644 --- a/vllm/model_executor/models/paddleocr_vl.py +++ b/vllm/model_executor/models/paddleocr_vl.py @@ -68,6 +68,7 @@ ) from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum from .ernie45 import Ernie4_5ForCausalLM @@ -800,16 +801,18 @@ def forward( [height_position_ids, width_position_ids], dim=-1, ) - max_grid_size = pids.max() + 1 + # The ids are built from the grids above, so `h`/`w` bound them + # and the table size is known on the host. + max_grid_size = max(max(h, w) for _, h, w in flatten_image_grid_thw) rope_emb_max_grid = self.rotary_pos_emb(max_grid_size) rotary_pos_emb = rope_emb_max_grid[pids].flatten(1) if cu_seqlens is None: raise ValueError("cu_seqlens cannot be None for SiglipEncoder.") if not isinstance(cu_seqlens, torch.Tensor): - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + cu_seqlens = async_tensor_h2d(cu_seqlens, dtype=torch.int32, device=device) else: - cu_seqlens = cu_seqlens.to(device=device) + cu_seqlens = cu_seqlens.to(device=device, non_blocking=True) max_seqlen = None if self.attn_backend in { @@ -1131,10 +1134,13 @@ def encode_image( siglip_position_ids.append(image_position_ids) cu_seqlens.append(cu_seqlens[-1] + numel) + # Both are built on the host; stage them over non-blocking. siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to( - pixel_values.device + pixel_values.device, non_blocking=True + ) + cu_seqlens = async_tensor_h2d( + cu_seqlens, dtype=torch.int32, device=pixel_values.device ) - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(pixel_values.device) vision_outputs = self.visual( pixel_values=pixel_values, diff --git a/vllm/model_executor/models/param2moe.py b/vllm/model_executor/models/param2moe.py index 7830272c324e..d865301a211f 100644 --- a/vllm/model_executor/models/param2moe.py +++ b/vllm/model_executor/models/param2moe.py @@ -649,15 +649,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index 8ea7b58eaf97..9a8da5ed3215 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -49,6 +49,7 @@ ResolvedPromptUpdate, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .idefics2_vision_model import Idefics2VisionTransformer @@ -396,12 +397,23 @@ def forward( w * base_feat_width // base_feat_height_reduction, ) ) - useful_height = int(reshaped_image_attention_mask[0, :, 0].sum().item()) - useful_width = int(reshaped_image_attention_mask[0, 0, :].sum().item()) + # The mask stays on device for the encoder above, so these + # per-image counts have to come back to the host to drive the + # slicing and the Python-level length arithmetic. + with gpu_sync_allowed(): + useful_height = int( + reshaped_image_attention_mask[0, :, 0].sum().item() + ) + useful_width = int( + reshaped_image_attention_mask[0, 0, :].sum().item() + ) + mask_token_count = int( + image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item() + ) sub_img = sub_img[:, :useful_height, :useful_width] temp_sub_GN = self.sub_GN.repeat(1, useful_height, 1, 1) temp_len = ( - int(image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item()) + mask_token_count + (useful_height + 1) + base_feat_height // base_feat_height_reduction ) @@ -905,7 +917,7 @@ def _get_mm_fields_config( return dict( input_image_embeds=MultiModalFieldConfig.batched("image"), image_attention_mask=MultiModalFieldConfig.batched("image"), - image_sizes=MultiModalFieldConfig.batched("image"), + image_sizes=MultiModalFieldConfig.batched("image", keep_on_cpu=True), num_img_tokens=MultiModalFieldConfig.batched("image"), input_audio_embeds=MultiModalFieldConfig.batched("audio"), ) diff --git a/vllm/model_executor/models/qwen2.py b/vllm/model_executor/models/qwen2.py index f1d5f23a264c..3ec8ac931913 100644 --- a/vllm/model_executor/models/qwen2.py +++ b/vllm/model_executor/models/qwen2.py @@ -460,15 +460,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 0715235ebb15..96a58df1ddaa 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -89,6 +89,7 @@ PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -1036,12 +1037,14 @@ def _process_audio_input( self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths) ) - audio_outputs = self.audio_tower( - input_features.to(self.audio_tower.dtype), - feature_lens=audio_feature_lengths, - aftercnn_lens=audio_feat_lengths, - ) - return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist()) + with gpu_sync_allowed(): + audio_outputs = self.audio_tower( + input_features.to(self.audio_tower.dtype), + feature_lens=audio_feature_lengths, + aftercnn_lens=audio_feat_lengths, + ) + split_sizes = audio_output_lengths.tolist() + return audio_outputs.last_hidden_state.split(split_sizes) def _process_image_input( self, image_input: Qwen2_5_VLImageInputs @@ -1526,12 +1529,13 @@ def embed_input_ids( video_token_id = self.config.video_token_index audio_token_id = self.config.audio_token_index - input_ids_cpu = input_ids.cpu() - is_video = is_multimodal & (input_ids_cpu == video_token_id) - is_audio = is_multimodal & (input_ids_cpu == audio_token_id) + with gpu_sync_allowed(): + input_ids_cpu = input_ids.cpu() + is_video = is_multimodal & (input_ids_cpu == video_token_id) + is_audio = is_multimodal & (input_ids_cpu == audio_token_id) - num_video = is_video.sum().item() - num_audio = is_audio.sum().item() + num_video = is_video.sum().item() + num_audio = is_audio.sum().item() if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio): inputs_embeds = self._embed_text_input_ids( diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index aa105ed8e8f7..d7d0d1fa5a76 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -87,6 +87,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers +from ...utils.gpu_sync_debug import gpu_sync_allowed from .interfaces import ( MultiModalEmbeddings, SupportsEagle, @@ -1569,8 +1570,9 @@ def _postprocess_video_embeds_evs( video_second_per_grid=video_second_per_grid_t.item(), ).to(emb.device, non_blocking=True) - emb = emb[retention_mask] - positions = positions[retention_mask] + with gpu_sync_allowed(): + emb = emb[retention_mask] + positions = positions[retention_mask] emb = torch.cat([emb, positions], dim=1) video_embeds_out.append(emb) return tuple(video_embeds_out) @@ -1624,15 +1626,16 @@ def recompute_mrope_positions( mm[:, -4:].permute(1, 0).long() for mm in multimodal_embeddings ] - positions, mrope_positions_delta = recompute_mrope_positions( - input_ids_t, - mm_embeddings_pos, - mrope_positions, - num_computed_tokens, - vision_start_token_id, - image_token_id, - video_token_id, - ) + with gpu_sync_allowed(): + positions, mrope_positions_delta = recompute_mrope_positions( + input_ids_t, + mm_embeddings_pos, + mrope_positions, + num_computed_tokens, + vision_start_token_id, + image_token_id, + video_token_id, + ) return mm_embeddings_out, positions, mrope_positions_delta diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 9692d9615b51..115a7f7f79b6 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -60,6 +60,7 @@ PromptUpdate, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP @@ -444,12 +445,15 @@ def _process_audio_input( ) < audio_output_lengths ) - masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim) + with gpu_sync_allowed(): + masked_audio_features = audio_features[audio_features_mask].view( + -1, embed_dim + ) - # Split to tuple of embeddings for individual audio input. - return torch.split( - masked_audio_features, audio_output_lengths.flatten().tolist() - ) + # Split to tuple of embeddings for individual audio input. + return torch.split( + masked_audio_features, audio_output_lengths.flatten().tolist() + ) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: audio_input = self._parse_and_validate_audio_input(**kwargs) diff --git a/vllm/model_executor/models/qwen2_moe.py b/vllm/model_executor/models/qwen2_moe.py index 43274e4ba7e0..df1794d2893e 100644 --- a/vllm/model_executor/models/qwen2_moe.py +++ b/vllm/model_executor/models/qwen2_moe.py @@ -461,7 +461,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/qwen3.py b/vllm/model_executor/models/qwen3.py index cf136db5be48..7db73a749f75 100644 --- a/vllm/model_executor/models/qwen3.py +++ b/vllm/model_executor/models/qwen3.py @@ -295,15 +295,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index 6ebf26c8719a..28c6a1189937 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -334,15 +334,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=self.quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/qwen3_5_mtp.py b/vllm/model_executor/models/qwen3_5_mtp.py index 6421c8b60d2d..67f06f4f5f24 100644 --- a/vllm/model_executor/models/qwen3_5_mtp.py +++ b/vllm/model_executor/models/qwen3_5_mtp.py @@ -10,7 +10,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig, get_current_vllm_config -from vllm.distributed.parallel_state import get_pp_group +from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -26,9 +26,9 @@ ) from vllm.model_executor.models.qwen3_next import ( QwenNextMixtureOfExperts, - _all_gather_hidden_and_residual, _is_shared_expert_fse_compatible, ) +from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig @@ -161,6 +161,10 @@ def forward( current_step_idx = spec_step_idx % self.num_mtp_layers mtp_layer = self.layers[current_step_idx] + if mtp_layer.use_attn_reduce_scatter_for_moe: + assert hidden_states.shape[0] == positions.shape[-1] + hidden_states = sequence_parallel_chunk(hidden_states) + assert residual is None hidden_states, residual = mtp_layer( positions=positions, hidden_states=hidden_states, @@ -172,15 +176,10 @@ def forward( {"hidden_states": hidden_states, "residual": residual} ) - if mtp_layer.use_attn_reduce_scatter_for_moe: - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - positions.shape[-1], - self.config.hidden_size, - ) - hidden_states, _ = self.norm(hidden_states, residual) + if mtp_layer.use_attn_reduce_scatter_for_moe: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[-1]] return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -238,15 +237,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=self.quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index b9224adc6c8e..f5ec1a452629 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -92,6 +92,7 @@ from vllm.transformers_utils.processors.qwen3_asr import ( Qwen3ASRProcessor, ) +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) _ASR_TEXT_TAG = "" @@ -467,7 +468,9 @@ def _process_audio_input( feature_lens=audio_feature_lengths, aftercnn_lens=audio_output_lengths, ) - return audio_features.split(audio_output_lengths.tolist()) + with gpu_sync_allowed(): + split_sizes = audio_output_lengths.tolist() + return audio_features.split(split_sizes) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 90b2f48faa82..35c04331b001 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -72,6 +72,26 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool: ) +def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: + """The target's RoPE layout, from its first attention layer. + + A DFlash head must rotate Q/K the way the target it was distilled against + does, and a mismatch is silent — acceptance collapses but nothing errors and + the output stays correct. Draft checkpoints do not carry this, so take it + from the target. None if the target uses no RoPE. + """ + language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + for module in language_model.modules(): + style = getattr(module, "is_neox_style", None) + if isinstance(style, bool): + return style + return None + + def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: spec_config = vllm_config.speculative_config config = spec_config.draft_model_config.hf_config @@ -166,6 +186,7 @@ def __init__( add_swa_attention_sink_bias: bool = False, sliding_window: int | None = None, causal: bool = False, + is_neox_style: bool = True, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -209,6 +230,7 @@ def __init__( self.rotary_emb = get_rope( self.head_dim, max_position=max_position, + is_neox_style=is_neox_style, rope_parameters=rope_parameters, ) @@ -291,6 +313,13 @@ def __init__( # non-causal) from the draft config. sliding_window, causal = _resolve_layer_attention(config, layer_idx) + # RoPE layout, copied off the target at load time by the draft loader + # (see `dflash_target_rope_is_neox_style`). Checkpoints do not carry it: + # a head distilled from an interleaved-RoPE target must rotate the way + # that target does, or every drafted Q/K is wrong and acceptance + # collapses with no error raised. + is_neox_style = getattr(config, "is_neox_style", True) + self.self_attn = DFlashQwen3Attention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, @@ -301,6 +330,7 @@ def __init__( add_swa_attention_sink_bias=add_swa_attention_sink_bias, sliding_window=sliding_window, causal=causal, + is_neox_style=is_neox_style, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, @@ -345,7 +375,14 @@ def forward( @support_torch_compile class DFlashQwen3Model(nn.Module): hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={"midlayer.": "layers.0."}, + orig_to_new_substr={ + "midlayer.": "layers.0.", + # Muse-Glimmer-30B-assistant names the aux-hidden-state encoder + # `encoder.fc` / `encoder.output_norm_enc`; this head calls them + # `fc` / `hidden_norm`. Same tensors and shapes, different names. + "encoder.output_norm_enc.": "hidden_norm.", + "encoder.fc.": "fc.", + }, orig_to_new_stacked={ ".q_proj": (".qkv_proj", "q"), ".k_proj": (".qkv_proj", "k"), diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index c33166469a2b..5df982d741ea 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -575,7 +575,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ffeada5444b7..b241baca6cbc 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -78,6 +78,18 @@ KVCache = tuple[torch.Tensor, torch.Tensor] +def _should_use_sequence_parallel(vllm_config: VllmConfig) -> bool: + config = vllm_config.model_config.hf_text_config + parallel_config = vllm_config.parallel_config + return ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + and getattr(config, "num_experts", 0) > 0 + and not getattr(config, "mlp_only_layers", []) + and getattr(config, "decoder_sparse_step", 1) == 1 + ) + + def _is_shared_expert_fse_compatible(quant_config) -> bool: """Check if shared expert can be fused with routed experts. @@ -397,7 +409,6 @@ def __init__( model_config = vllm_config.model_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config - parallel_config = vllm_config.parallel_config self.layer_type = layer_type self.layer_idx = extract_layer_index(prefix) @@ -409,10 +420,8 @@ def __init__( config.num_experts > 0 and (self.layer_idx + 1) % config.decoder_sparse_step == 0 ) - self.use_attn_reduce_scatter_for_moe = ( - parallel_config.use_sequence_parallel_moe - and parallel_config.pipeline_parallel_size == 1 - and is_moe_layer + self.use_attn_reduce_scatter_for_moe = _should_use_sequence_parallel( + vllm_config ) if self.layer_type == "linear_attention": @@ -481,11 +490,6 @@ def forward( **kwargs: object, ): full_num_tokens = positions.shape[-1] - input_is_sequence_parallel = ( - self.use_attn_reduce_scatter_for_moe - and residual is not None - and hidden_states.shape[0] != full_num_tokens - ) if residual is None: residual = hidden_states @@ -493,7 +497,7 @@ def forward( else: hidden_states, residual = self.input_layernorm(hidden_states, residual) - if input_is_sequence_parallel: + if self.use_attn_reduce_scatter_for_moe: hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) hidden_states = hidden_states[:full_num_tokens] @@ -524,8 +528,6 @@ def forward( # pad if not divisible by world size hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad)) hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) - if not input_is_sequence_parallel: - residual = sequence_parallel_chunk(residual) # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) @@ -554,24 +556,6 @@ def forward( return hidden_states, residual -def _all_gather_hidden_and_residual( - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - full_num_tokens: int, - hidden_size: int, -) -> tuple[torch.Tensor, torch.Tensor | None]: - if residual is None: - hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) - hidden_states = hidden_states[:full_num_tokens] - return hidden_states, None - - combined_states = torch.cat([hidden_states, residual], dim=-1) - combined_states = tensor_model_parallel_all_gather(combined_states, 0) - combined_states = combined_states[:full_num_tokens] - hidden_states, residual = combined_states.split([hidden_size, hidden_size], dim=-1) - return hidden_states, residual - - @support_torch_compile class Qwen3NextModel(nn.Module, EagleModelMixin): hf_to_vllm_mapper = WeightsMapper( @@ -629,6 +613,10 @@ def get_layer(prefix: str): def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) + @property + def use_sequence_parallel(self) -> bool: + return self.layers[self.start_layer].use_attn_reduce_scatter_for_moe + def forward( self, input_ids: torch.Tensor | None, @@ -648,35 +636,20 @@ def forward( residual = intermediate_tensors["residual"] full_num_tokens = positions.shape[-1] + if self.use_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + assert residual is None + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): - if ( - hidden_states.shape[0] != full_num_tokens - and not layer.use_attn_reduce_scatter_for_moe - ): - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - full_num_tokens, - self.config.hidden_size, - ) hidden_states, residual = layer( positions=positions, hidden_states=hidden_states, residual=residual, ) - if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[ - 0 - ] != full_num_tokens: - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - full_num_tokens, - self.config.hidden_size, - ) self._maybe_add_hidden_state( aux_hidden_states, layer_idx + 1, hidden_states, residual ) @@ -685,14 +658,19 @@ def forward( return IntermediateTensors( {"hidden_states": hidden_states, "residual": residual} ) - if hidden_states.shape[0] != full_num_tokens: - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - full_num_tokens, - self.config.hidden_size, - ) hidden_states, _ = self.norm(hidden_states, residual) + if self.use_sequence_parallel: + if aux_hidden_states: + hidden_size = hidden_states.shape[-1] + hidden_states = torch.cat([hidden_states, *aux_hidden_states], dim=-1) + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + hidden_states, *aux_hidden_states = hidden_states.split( + hidden_size, dim=-1 + ) + else: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] if aux_hidden_states: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index c832e955b24a..35f550700a54 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -9,7 +9,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig -from vllm.distributed.parallel_state import get_pp_group +from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather from vllm.logger import init_logger from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -22,8 +22,8 @@ Qwen3NextModel, Qwen3NextRMSNorm, QwenNextMixtureOfExperts, - _all_gather_hidden_and_residual, ) +from vllm.model_executor.models.utils import sequence_parallel_chunk from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig @@ -131,6 +131,10 @@ def forward( current_step_idx = spec_step_idx % self.num_mtp_layers mtp_layer = self.layers[current_step_idx] + if mtp_layer.use_attn_reduce_scatter_for_moe: + assert hidden_states.shape[0] == positions.shape[-1] + hidden_states = sequence_parallel_chunk(hidden_states) + assert residual is None hidden_states, residual = mtp_layer( positions=positions, hidden_states=hidden_states, @@ -142,14 +146,10 @@ def forward( {"hidden_states": hidden_states, "residual": residual} ) - if mtp_layer.use_attn_reduce_scatter_for_moe: - hidden_states, residual = _all_gather_hidden_and_residual( - hidden_states, - residual, - positions.shape[-1], - self.config.hidden_size, - ) hidden_states, _ = self.norm(hidden_states, residual) + if mtp_layer.use_attn_reduce_scatter_for_moe: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[: positions.shape[-1]] return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index bb938167aec2..48a16e97d1b4 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -77,6 +77,7 @@ ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -434,11 +435,14 @@ def forward( feature_lens: torch.Tensor, aftercnn_lens: torch.Tensor, ): - # Compute chunk information + # Compute chunk information. chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long() + with gpu_sync_allowed(): + total_chunks = int(chunk_num.sum()) + chunk_lengths = async_tensor_h2d( - [self.n_window * 2] * chunk_num.sum(), + [self.n_window * 2] * total_chunks, dtype=torch.long, device=feature_lens.device, ) @@ -447,15 +451,18 @@ def forward( chunk_lengths[chunk_lengths == 0] = self.n_window * 2 # Split input features into chunks and pad - chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0) + with gpu_sync_allowed(): + chunk_lengths_list = chunk_lengths.tolist() + chunk_list = input_features.T.split(chunk_lengths_list, dim=0) padded_feature = nn.utils.rnn.pad_sequence( chunk_list, batch_first=True ).transpose(1, 2) # Compute feature lengths after CNN feature_lens_after_cnn = self._get_cnn_output_lengths(chunk_lengths) - # Vectorized mask creation: avoid creating many small tensors - max_len_after_cnn = feature_lens_after_cnn.max().item() + # Vectorized mask creation: avoid creating many small tensors. + with gpu_sync_allowed(): + max_len_after_cnn = feature_lens_after_cnn.max().item() indices = torch.arange(max_len_after_cnn, device=padded_feature.device) padded_mask_after_cnn = indices.unsqueeze(0) < feature_lens_after_cnn.unsqueeze( 1 @@ -494,16 +501,17 @@ def forward( ) padded_embed = padded_embed + positional_embedding - # Extract valid hidden states and compute cu_seqlens - hidden_states = padded_embed[padded_mask_after_cnn] + with gpu_sync_allowed(): + hidden_states = padded_embed[padded_mask_after_cnn] + # Use tolist() for efficient batch conversion from tensor to Python. + aftercnn_lens_list = aftercnn_lens.tolist() # Compute cumulative sequence lengths for chunked attention cu_chunk_lens = [0] window_aftercnn = padded_mask_after_cnn.shape[-1] * ( self.n_window_infer // (self.n_window * 2) ) - # Use tolist() for efficient batch conversion from tensor to Python - for cnn_len in aftercnn_lens.tolist(): + for cnn_len in aftercnn_lens_list: num_full_chunks = cnn_len // window_aftercnn remainder = cnn_len % window_aftercnn cu_chunk_lens.extend([window_aftercnn] * num_full_chunks) @@ -1111,7 +1119,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config.vocab_size, config.hidden_size, quant_config=quant_config ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index c34d95db3258..de341d302a35 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -115,6 +115,7 @@ from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers +from ...utils.gpu_sync_debug import gpu_sync_allowed from ...utils.torch_utils import async_tensor_h2d from .interfaces import ( MultiModalEmbeddings, @@ -1707,15 +1708,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix="lm_head", + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix="lm_head", - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() @@ -2369,19 +2369,17 @@ def _postprocess_video_embeds_evs( spatial_merge_size=self.visual.spatial_merge_size, q=self.video_pruning_rate, ) - # Apply retention mask. - emb = emb[retention_mask] - # Calculate the actual number of retained tokens per frame. - num_frames, rows, cols = ( - t, - h // merge_size, - w // merge_size, - ) - retention_mask_thw = retention_mask.reshape(num_frames, rows, cols) - num_tokens_per_frame = ( - retention_mask_thw.sum(dim=(1, 2)).long().tolist() - ) + with gpu_sync_allowed(): + # Apply retention mask. + emb = emb[retention_mask] + + # Calculate the actual number of retained tokens per frame. + num_frames, rows, cols = t, h // merge_size, w // merge_size + retention_mask_thw = retention_mask.reshape(num_frames, rows, cols) + num_tokens_per_frame = ( + retention_mask_thw.sum(dim=(1, 2)).long().tolist() + ) else: feature_size = emb.shape[0] // num_frames num_tokens_per_frame = [feature_size] * num_frames @@ -2549,10 +2547,14 @@ def _get_expanded_positions( .permute(1, 0) ) full_is_video_embed = unpruned_token_ids_tensor == embed_token_id - expanded_positions[is_video_embed, :3] = original_mrope[full_is_video_embed][ - retention_mask - ] - expanded_positions[~is_video_embed, :3] = original_mrope[~full_is_video_embed] + + with gpu_sync_allowed(): + expanded_positions[is_video_embed, :3] = original_mrope[ + full_is_video_embed + ][retention_mask] + expanded_positions[~is_video_embed, :3] = original_mrope[ + ~full_is_video_embed + ] expanded_positions[..., 3] = is_vision_start expanded_positions[..., 4] = is_video_embed @@ -2815,15 +2817,16 @@ def _recompute_mrope_positions( torch.empty(5, 0, device=device, dtype=torch.long) ) - positions, mrope_positions_delta = recompute_mrope_positions( - input_ids_t, - mm_embeddings_pos, - mrope_positions, - num_computed_tokens, - vision_start_token_id, - image_token_id, - video_token_id, - ) + with gpu_sync_allowed(): + positions, mrope_positions_delta = recompute_mrope_positions( + input_ids_t, + mm_embeddings_pos, + mrope_positions, + num_computed_tokens, + vision_start_token_id, + image_token_id, + video_token_id, + ) return mm_embeddings_out, positions, mrope_positions_delta diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index 1360edfdd7b3..c7cc440b8c2f 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -168,7 +168,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(self.config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 87545b092542..1e4bf6731c96 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -183,6 +183,7 @@ "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), "Olmo3ForCausalLM": ("olmo3", "Olmo3ForCausalLM"), + "MuseGlimmerForCausalLM": ("muse_glimmer", "MuseGlimmerForCausalLM"), "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"), "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), "OPTForCausalLM": ("opt", "OPTForCausalLM"), @@ -530,6 +531,7 @@ "NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"), "NVLM_D": ("nvlm_d", "NVLM_D_Model"), + "MuseGlimmerForConditionalGeneration": ("muse_glimmer", "MuseGlimmerForCausalLM"), "OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"), "OpenPanguVLForConditionalGeneration": ( "openpangu_vl", @@ -623,6 +625,14 @@ "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"), "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"), "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + # Muse Glimmer's DFlash draft head, reusing the generic qwen3_dflash + # implementation. EAGLEConfig rewrites a dflash draft's architecture to + # DFlash{arch} unless it already starts or ends with "DFlash" (see + # transformers_utils/configs/eagle.py), so the name the registry is asked + # for is DFlashMuseGlimmerAssistantModel -- same convention as + # DFlashLagunaForCausalLM. The bare name is kept as a defensive alias. + "MuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), + "DFlashMuseGlimmerAssistantModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"), "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "K3DSparkModel": ( diff --git a/vllm/model_executor/models/sarvam.py b/vllm/model_executor/models/sarvam.py index 2b2fbfe70897..04590a2a913a 100644 --- a/vllm/model_executor/models/sarvam.py +++ b/vllm/model_executor/models/sarvam.py @@ -615,15 +615,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.tie_word_embeddings = getattr(config, "tie_word_embeddings", False) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py index d2c767846d70..f3fa8b2d5de2 100644 --- a/vllm/model_executor/models/seed_oss.py +++ b/vllm/model_executor/models/seed_oss.py @@ -388,15 +388,14 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index 5808c9539bfc..6d924a286501 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -54,6 +54,7 @@ TimingContext, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant @@ -998,27 +999,31 @@ def _flip_sequences_by_position_ids( position_diffs = position_ids[1:] - position_ids[:-1] boundary_mask = position_diffs <= 0 - boundary_indices = torch.cat( - [ - torch.tensor([0], device=features.device), - torch.where(boundary_mask)[0] + 1, - torch.tensor([len(features)], device=features.device), - ] - ) + with gpu_sync_allowed(): + boundary_mid_cpu = torch.where(boundary_mask.cpu())[0] + 1 + + zero = torch.zeros(1, dtype=boundary_mid_cpu.dtype) + end = torch.full((1,), len(features), dtype=boundary_mid_cpu.dtype) + boundary_indices_cpu = torch.cat([zero, boundary_mid_cpu, end]) # For each sequence [start, end), position i flips to: start + end - 1 - i - lengths = boundary_indices[1:] - boundary_indices[:-1] - starts = boundary_indices[:-1] - ends = boundary_indices[1:] + lengths_cpu = boundary_indices_cpu[1:] - boundary_indices_cpu[:-1] + starts_cpu = boundary_indices_cpu[:-1] + ends_cpu = boundary_indices_cpu[1:] # Assign sequence ID to each element - sequence_ids = torch.arange( - len(lengths), device=features.device - ).repeat_interleave(lengths) + sequence_ids_cpu = torch.arange( + len(lengths_cpu), dtype=boundary_mid_cpu.dtype + ).repeat_interleave(lengths_cpu) # Calculate flipped indices for all positions at once - current_positions = torch.arange(len(features), device=features.device) - flip_indices = starts[sequence_ids] + ends[sequence_ids] - 1 - current_positions + current_positions_cpu = torch.arange( + len(features), dtype=boundary_mid_cpu.dtype + ) + flip_indices_cpu = ( + starts_cpu[sequence_ids_cpu] + ends_cpu[sequence_ids_cpu] + ) - (1 + current_positions_cpu) + flip_indices = flip_indices_cpu.to(features.device, non_blocking=True) return features[flip_indices] diff --git a/vllm/model_executor/models/solar.py b/vllm/model_executor/models/solar.py index 478a61da6754..430337c77d1c 100644 --- a/vllm/model_executor/models/solar.py +++ b/vllm/model_executor/models/solar.py @@ -386,7 +386,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=maybe_prefix(prefix, "lm_head"), ) if config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) logit_scale = getattr(config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( diff --git a/vllm/model_executor/models/stablelm.py b/vllm/model_executor/models/stablelm.py index 58758b11cdda..b18f4051cafe 100644 --- a/vllm/model_executor/models/stablelm.py +++ b/vllm/model_executor/models/stablelm.py @@ -294,7 +294,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.lm_head", ) if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight + self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens) self.logits_processor = LogitsProcessor(config.vocab_size) self.make_empty_intermediate_tensors = ( self.model.make_empty_intermediate_tensors diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index b49f7827dafe..a160d5bd17f7 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -61,6 +61,7 @@ TimingContext, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal from .interfaces_base import attn_type @@ -280,7 +281,9 @@ def forward( inputs_embeds: torch.Tensor | None = None, **kwargs: object, ): - model_output = self.inference_runner.forward(**kwargs) + # terratorch's forward has internal GPU syncs. + with gpu_sync_allowed(): + model_output = self.inference_runner.forward(**kwargs) return model_output.output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index b637df116493..a32e2b54ba42 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -59,9 +59,10 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self.lm_head = self.lm_head.tie_weights(module) break - logit_scale = getattr(self.text_config, "logit_scale", 1.0) self.logits_processor = LogitsProcessor( - self.text_config.vocab_size, scale=logit_scale + self.text_config.vocab_size, + scale=getattr(self.text_config, "logit_scale", 1.0), + soft_cap=getattr(self.text_config, "final_logit_softcapping", None), ) else: self.lm_head = PPMissingLayer() diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 84d6813c8e2c..f328a18a3ce0 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -54,6 +54,7 @@ TimingContext, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed if TYPE_CHECKING: from transformers import BatchFeature, PreTrainedModel @@ -726,12 +727,17 @@ def _process_audio_input(self, **kwargs) -> list[torch.Tensor] | None: kwargs.pop("token_type_ids", None) kwargs.pop("mm_token_type_ids", None) - audio_output = self.model.get_audio_features( - input_features, return_dict=True, **kwargs - ) + # HuggingFace's `get_audio_features` implementations branch on + # per-sample feature lengths internally. + with gpu_sync_allowed(): + audio_output = self.model.get_audio_features( + input_features, return_dict=True, **kwargs + ) audio_embeddings = audio_output.pooler_output - split_sizes = num_audio_tokens.flatten().tolist() + # Per-audio token counts are needed as Python ints to split. + with gpu_sync_allowed(): + split_sizes = num_audio_tokens.flatten().tolist() return self._split_embeddings(audio_embeddings, split_sizes) def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: @@ -749,7 +755,12 @@ def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: num_image_patches = kwargs.pop("num_image_patches") - vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) + # The underlying HuggingFace `get_image_features` implementations + # contain model-internal syncs (e.g. Idefics3 filters all-zero + # padding images via boolean-mask indexing, LlavaOnevision + # branches on per-sample batch counts). + with gpu_sync_allowed(): + vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) # Transformers `v5`, `self.get_image_features` returns a tuple # containing the features and optionally attentions/hidden_states diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index ed3b6f1762fd..551c2c5694b2 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -43,6 +43,7 @@ from vllm.renderers import TokenizeParams from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.ultravox import UltravoxConfig +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -244,7 +245,11 @@ def _get_mm_fields_config( # higher than the number of audio samples audio_features=MultiModalFieldConfig.flat_from_sizes("audio", num_chunks), audio_token_len=MultiModalFieldConfig.flat_from_sizes("audio", num_chunks), - audio_lens=MultiModalFieldConfig.flat_from_sizes("audio", num_chunks), + # Only ever used to derive the encoder attention metadata on the + # host, so keep it there. + audio_lens=MultiModalFieldConfig.flat_from_sizes( + "audio", num_chunks, keep_on_cpu=True + ), # num_chunks can convert audio_chunked to audio batch dimension audio_num_chunks=MultiModalFieldConfig.batched("audio"), audio_embeds=MultiModalFieldConfig.batched("audio"), @@ -837,15 +842,19 @@ def _process_audio_input( embeddings.shape[0], -1 ) mask = indices < audio_token_len[:, None] - # Apply mask and flatten - flattened_embeddings = embeddings[mask] - # Return one tensor per input audio - embed_lens = [ - chunk_lens.sum().item() - for chunk_lens in audio_token_len.split(audio_input["num_chunks"].tolist()) - ] - return flattened_embeddings.split(embed_lens) + with gpu_sync_allowed(): + # Apply mask and flatten + flattened_embeddings = embeddings[mask] + + # Return one tensor per input audio + embed_lens = [ + chunk_lens.sum().item() + for chunk_lens in audio_token_len.split( + audio_input["num_chunks"].tolist() + ) + ] + return flattened_embeddings.split(embed_lens) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: audio_input = self._parse_and_validate_audio_input(**kwargs) diff --git a/vllm/model_executor/warmup/b12x_warmup.py b/vllm/model_executor/warmup/b12x_warmup.py new file mode 100644 index 000000000000..5e0e4a0fa498 --- /dev/null +++ b/vllm/model_executor/warmup/b12x_warmup.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm B12X JIT kernels used by a loaded model.""" + +from collections import Counter +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import torch + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def _collect_warmup_units( + model: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, +) -> Iterable[B12xWarmupUnit]: + units: dict[object, B12xWarmupUnit] = {} + for layer in model.modules(): + provider = getattr(layer, "b12x_warmup_provider", None) + get_unit = getattr(provider, "get_b12x_warmup_unit", None) + if not callable(get_unit): + continue + unit = get_unit(layer, token_counts, output_dtype) + assert isinstance(unit, B12xWarmupUnit) + units.setdefault(unit.key, unit) + return units.values() + + +def _compile_warmup_units( + units: Iterable[B12xWarmupUnit], +) -> Counter[str]: + warmed: Counter[str] = Counter() + with torch.inference_mode(): + for unit in units: + unit.compile() + warmed[unit.name] += 1 + if warmed: + torch.accelerator.synchronize() + return warmed + + +def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None: + if not current_platform.is_cuda(): + return + if not current_platform.is_device_capability_family(120): + return + + output_dtype = getattr( + getattr(worker, "model_config", None), + "dtype", + torch.bfloat16, + ) + if output_dtype not in (torch.bfloat16, torch.float16): + output_dtype = torch.bfloat16 + token_counts = b12x_warmup_token_counts( + max_tokens=worker.scheduler_config.max_num_batched_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, + ) + units = _collect_warmup_units( + worker.get_model(), + token_counts, + output_dtype, + ) + for name, count in _compile_warmup_units(units).items(): + logger.info_once( + "Warmed up %d B12X %s linear GEMM signatures.", + count, + name, + ) diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index d8c51409502e..9c43bbfc33d8 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -5,16 +5,20 @@ from __future__ import annotations import ast +import builtins import inspect import itertools import operator import textwrap from abc import ABC, abstractmethod -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass from typing import Any, Generic, TypeVar __all__ = [ + "JitWarmupRegistry", "VllmJitKernel", "WarmupIntRange", "get_ast_full_name", @@ -136,8 +140,15 @@ class _CompileKeyDispatchTrace: field_exprs: tuple[tuple[str, ast.AST], ...] globals: Mapping[str, Any] input_names: frozenset[str] + # Named parameters are excluded from direct **kwargs forwarding. + named_parameters: frozenset[str] | None defaults: Mapping[str, Any] + def input_names_for(self, available_names: set[str]) -> frozenset[str]: + if self.named_parameters is None: + return self.input_names + return self.input_names | (available_names - self.named_parameters) + def compile_key( self, compile_key_type: type[CompileKeyT], @@ -146,12 +157,20 @@ def compile_key( dispatch_values = _eval_local_exprs( self.local_exprs, {**self.defaults, **kwargs}, self.globals ) - return compile_key_type( - **{ - field: _eval_dispatch_expr(expr, dispatch_values, self.globals) - for field, expr in self.field_exprs + named_parameters = self.named_parameters + # Materialize direct fields before evaluating named AST expressions. + fields: dict[str, Any] = {} + if named_parameters is not None: + fields = { + name: value + for name, value in kwargs.items() + if name not in named_parameters } - ) + for field, expr in self.field_exprs: + if field in fields: + raise TypeError(f"CompileKey field '{field}' is specified twice") + fields[field] = _eval_dispatch_expr(expr, dispatch_values, self.globals) + return compile_key_type(**fields) @dataclass(frozen=True) @@ -186,6 +205,10 @@ def matches(self, kwargs: Mapping[str, Any]) -> bool: ast.LtE: operator.le, ast.Gt: operator.gt, ast.GtE: operator.ge, + ast.In: lambda left, right: left in right, + ast.NotIn: lambda left, right: left not in right, + ast.Is: operator.is_, + ast.IsNot: operator.is_not, } @@ -200,8 +223,9 @@ def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError: return ValueError( f"{reason}: {_dispatch_expr_source(node)}. " "Supported dispatch expressions are names, constants, attributes, " - "tuple/list literals, conditional expressions, comparisons, boolean " - "operators, unary not/minus, arithmetic, and calls without **kwargs." + "subscriptions, tuple/list literals, conditional expressions, " + "comparisons, boolean operators, unary not/minus, arithmetic, and " + "calls without **kwargs." ) @@ -225,6 +249,8 @@ def visit_Name(self, node: ast.Name) -> Any: return self.values[node.id] if node.id in self.globals: return self.globals[node.id] + if hasattr(builtins, node.id): + return getattr(builtins, node.id) raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'") def visit_Constant(self, node: ast.Constant) -> Any: @@ -299,6 +325,9 @@ def visit_Call(self, node: ast.Call) -> Any: def visit_Attribute(self, node: ast.Attribute) -> Any: return getattr(self.visit(node.value), node.attr) + def visit_Subscript(self, node: ast.Subscript) -> Any: + return self.visit(node.value)[self.visit(node.slice)] + def get_ast_full_name(node: ast.AST) -> str | None: if isinstance(node, ast.Name): @@ -310,15 +339,23 @@ def get_ast_full_name(node: ast.AST) -> str | None: return None -def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef: +def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef | ast.Lambda: source_fn = getattr(fn, "fn", fn) source = textwrap.dedent(inspect.getsource(source_fn)) tree = ast.parse(source) function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)] - if len(function_defs) != 1: - name = getattr(source_fn, "__name__", type(source_fn).__name__) - raise ValueError(f"Expected one function in {name}, found {len(function_defs)}") - return function_defs[0] + if len(function_defs) == 1: + return function_defs[0] + + lambdas = [node for node in ast.walk(tree) if isinstance(node, ast.Lambda)] + if len(lambdas) == 1: + return lambdas[0] + + name = getattr(source_fn, "__name__", type(source_fn).__name__) + raise ValueError( + f"Expected one function or lambda in {name}, found " + f"{len(function_defs)} functions and {len(lambdas)} lambdas" + ) def _eval_dispatch_expr( @@ -349,8 +386,11 @@ def _collect_input_names( def _collect_expression_body( fn: Callable[..., Any], - function_def: ast.FunctionDef, + function_def: ast.FunctionDef | ast.Lambda, ) -> tuple[list[tuple[str, ast.AST]], ast.AST]: + if isinstance(function_def, ast.Lambda): + return [], function_def.body + local_exprs: list[tuple[str, ast.AST]] = [] for statement in function_def.body: if ( @@ -421,6 +461,10 @@ def _trace_compile_key_dispatch( source_fn = getattr(fn, "__func__", fn) globals_ = source_fn.__globals__ function_def = get_function_source_node(fn) + if isinstance(function_def, ast.Lambda): + raise _dispatch_expr_error( + function_def, "Dispatch must be a function definition" + ) local_exprs, return_expr = _collect_expression_body(fn, function_def) if not isinstance(return_expr, ast.Call): @@ -431,13 +475,37 @@ def _trace_compile_key_dispatch( field_exprs: list[tuple[str, ast.AST]] = [] defaults, candidate_names = _function_trace_inputs(fn) + # Fields captured by dispatch **kwargs are forwarded, not AST-evaluated. + signature = inspect.signature(fn) + variadic_keyword = next( + ( + name + for name, parameter in signature.parameters.items() + if parameter.kind is inspect.Parameter.VAR_KEYWORD + ), + None, + ) + candidate_names.discard(variadic_keyword) input_names: set[str] = set() local_names = {name for name, _ in local_exprs} for _, expr in local_exprs: input_names.update(_collect_input_names(expr, candidate_names)) + named_parameters: frozenset[str] | None = None for keyword in return_expr.keywords: + # CompileKey may unpack only that **kwargs parameter, once. if keyword.arg is None: - raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey") + if ( + named_parameters is not None + or variadic_keyword is None + or not isinstance(keyword.value, ast.Name) + or keyword.value.id != variadic_keyword + ): + raise ValueError( + f"{fn.__name__} may unpack only its own **kwargs parameter " + "once in CompileKey" + ) + named_parameters = frozenset(candidate_names) + continue field_exprs.append((keyword.arg, keyword.value)) input_names.update( _collect_input_names(keyword.value, candidate_names, local_names) @@ -448,6 +516,7 @@ def _trace_compile_key_dispatch( tuple(field_exprs), globals_, frozenset(input_names), + named_parameters, defaults, ) @@ -482,10 +551,32 @@ class VllmJitKernel(Generic[CompileKeyT], ABC): CompileKey: type[CompileKeyT] def __init__(self) -> None: - self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + self._dispatch_trace = _trace_compile_key_dispatch(self.dispatch) + self._compiled_cache: dict[Any, Any] = {} def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT: - return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs) + return self._dispatch_trace.compile_key(self.CompileKey, kwargs) + + def _get_or_compile( + self, + compile_key: CompileKeyT, + *, + runtime_context: Mapping[str, Any] | None = None, + ) -> Any: + """Return a cached executor, compiling it on a monitored cache miss.""" + if compile_key not in self._compiled_cache: + self.compile(compile_key) + + try: + return self._compiled_cache[compile_key] + except KeyError as exc: + details = [f"compile_key={compile_key!r}"] + if runtime_context: + details.append(f"runtime_context={dict(runtime_context)!r}") + raise RuntimeError( + f"{type(self).__name__}.compile(...) did not cache its JIT " + f"executor ({', '.join(details)})" + ) from exc def _trace_dispatch( self, dispatch: CompileKeyDispatchFn[CompileKeyT] @@ -506,7 +597,11 @@ def traced( predicate_trace = ( _trace_warmup_predicate(_when) if _when is not None else None ) - input_names = compile_key_dispatch_trace.input_names + # Unmatched **kwargs fields also belong to the expansion space. + available_names = set(kwargs).union( + *(group.rows[0] for group in input_groups) + ) + input_names = compile_key_dispatch_trace.input_names_for(available_names) if predicate_trace is not None: input_names = input_names | predicate_trace.input_names expanded_input_groups = tuple( @@ -565,7 +660,105 @@ def compile(self, compile_key: CompileKeyT) -> None: """Compile one warmup key.""" raise NotImplementedError + def register_warmup(self, *args: Any, **kwargs: Any) -> None: + """Register this kernel with the active runner's warmup registry.""" + JitWarmupRegistry.register(self, *args, **kwargs) + def warmup(self, *args: Any, **kwargs: Any) -> None: """Compile this kernel's warmup keys.""" for compile_key in self.get_warmup_keys(*args, **kwargs): self.compile(compile_key) + + +class JitWarmupRegistry: + """Collect and compile JIT kernels selected during runner setup.""" + + _active: ContextVar[JitWarmupRegistry | None] = ContextVar( + "active_jit_warmup_registry", + default=None, + ) + + def __init__(self, vllm_config: Any) -> None: + self.vllm_config = vllm_config + self._registrations: dict[ + VllmJitKernel[Any], + list[tuple[tuple[Any, ...], dict[str, Any]]], + ] = {} + + @contextmanager + def activate(self) -> Iterator[None]: + """Collect registrations made in this context.""" + token = self._active.set(self) + try: + yield + finally: + self._active.reset(token) + + @classmethod + def register( + cls, + kernel: VllmJitKernel[Any], + *args: Any, + **kwargs: Any, + ) -> None: + """Register a kernel with the active registry, if one exists.""" + registry = cls._active.get() + if registry is not None: + registry._add(kernel, args, kwargs) + + def _add( + self, + kernel: VllmJitKernel[Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> None: + registrations = self._registrations.setdefault(kernel, []) + if ( + not args + and not kwargs + and any( + not registered_args and not registered_kwargs + for registered_args, registered_kwargs in registrations + ) + ): + return + registrations.append((args, kwargs)) + + def __len__(self) -> int: + return sum(len(registrations) for registrations in self._registrations.values()) + + def warmup(self) -> None: + """Expand registrations and compile each wrapper/key pair once.""" + from tqdm import tqdm + + from vllm.distributed import is_global_first_rank + + kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = [] + for kernel, registrations in self._registrations.items(): + compile_keys: dict[Any, None] = {} + for args, kwargs in registrations: + if not args and not kwargs: + args = (self.vllm_config,) + for compile_key in kernel.get_warmup_keys(*args, **kwargs): + compile_keys[compile_key] = None + if compile_keys: + kernel_items.append((kernel, compile_keys)) + + if not kernel_items: + return + + total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items) + with tqdm( + kernel_items, + desc=f"JIT kernel warmup ({total_keys} compile keys)", + disable=not is_global_first_rank(), + dynamic_ncols=True, + unit="kernel", + ) as progress: + for kernel, compile_keys in progress: + progress.set_postfix_str( + f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", + refresh=False, + ) + for compile_key in compile_keys: + kernel.compile(compile_key) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index b90975762175..29da528c93da 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -12,6 +12,42 @@ ) +def triton_scalar_specialization_rep(value: int) -> int: + """Return an integer with the same default Triton JIT specialization. + + For an ordinary integer argument, Triton's cache key contains its inferred + type (``i32``, ``i64``, or ``u64``) and one of three value classes: + + * ``1`` is specialized as the exact constant ``1``. + * Multiples of 16 receive a ``tt.divisibility = 16`` attribute. + * All other values have no value specialization. + + Warmup only needs one concrete value for each cache-key class. This helper + returns ``1`` for the exact-one class and otherwise returns a divisible or + generic representative while preserving the inferred integer type. + + This applies only to non-``constexpr`` integer arguments using Triton's + default specialization. Do not use it for arguments listed in + ``do_not_specialize`` or ``do_not_specialize_on_alignment``. + """ + if value == 1: + return 1 + + if -(1 << 31) <= value < (1 << 31): + divisible_rep = 16 + generic_rep = 2 + elif -(1 << 63) <= value < (1 << 63): + divisible_rep = 1 << 31 + generic_rep = (1 << 31) + 1 + elif 0 <= value < (1 << 64): + divisible_rep = 1 << 63 + generic_rep = (1 << 63) + 1 + else: + raise OverflowError(f"Integer {value} is outside Triton's scalar range") + + return divisible_rep if value % 16 == 0 else generic_rep + + @dataclass(frozen=True) class TritonWarmupTensor: # Compile-only tensor descriptor for Triton pointer specialization. @@ -167,6 +203,8 @@ def trace_triton_kernel_specialization_args( kernel: Callable[..., Any], ) -> tuple[str, ...]: function_def = get_function_source_node(kernel) + if not isinstance(function_def, ast.FunctionDef): + raise ValueError("Expected Triton kernel to be defined as a function") source_fn = getattr(kernel, "fn", kernel) arg_names = tuple(inspect.signature(source_fn).parameters) constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c35e1ac30c9a..e25a60c8445a 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -6,12 +6,14 @@ happen during model execution. """ +import time from typing import TYPE_CHECKING import torch import vllm.envs as envs from vllm.logger import init_logger +from vllm.model_executor.warmup.b12x_warmup import b12x_warmup from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( @@ -35,9 +37,6 @@ from vllm.model_executor.warmup.sparse_mla_triton_warmup import ( sparse_mla_triton_warmup, ) -from vllm.model_executor.warmup.v1_block_table_warmup import ( - warm_v1_block_table_kernels, -) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -102,26 +101,40 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): ) if not worker.use_v2_model_runner: - # Pooling models do not use the generation slot-mapping path. - if not worker.model_runner.is_pooling_model: - warm_v1_block_table_kernels(worker.model_runner) # The KV-block zeroing kernel is driven by the scheduler's # `new_block_ids_to_zero`, so no dummy run ever reaches it. zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None) if zeroer is not None: zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks) + if worker.vllm_config.kernel_config.enable_jit_warmup: + logger.info("JIT kernel warmup starting.") + jit_warmup_start = time.perf_counter() + try: + worker.model_runner.jit_warmup_registry.warmup() + except Exception: + logger.exception( + "JIT kernel warmup failed after %.2fs.", + time.perf_counter() - jit_warmup_start, + ) + raise + logger.info( + "JIT kernel warmup finished in %.2fs.", + time.perf_counter() - jit_warmup_start, + ) + qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config) + compilation_config = worker.vllm_config.compilation_config + cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or []) + # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder # layer per token; warm them across token sizes first so the first real # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside). deepseek_v4_mhc_warmup( worker.get_model(), max_tokens=worker.scheduler_config.max_num_batched_tokens, - cudagraph_capture_sizes=( - worker.vllm_config.compilation_config.cudagraph_capture_sizes or [] - ), + cudagraph_capture_sizes=cudagraph_capture_sizes, ) # Run next so input-prep kernels JIT against pristine runner state. @@ -156,6 +169,8 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): max_tokens = worker.scheduler_config.max_num_batched_tokens deep_gemm_warmup(model, max_tokens) + b12x_warmup(worker, cudagraph_capture_sizes) + minimax_m3_msa_warmup(worker) enable_flashinfer_autotune = ( diff --git a/vllm/model_executor/warmup/v1_block_table_warmup.py b/vllm/model_executor/warmup/v1_block_table_warmup.py deleted file mode 100644 index d49e1ba7cc89..000000000000 --- a/vllm/model_executor/warmup/v1_block_table_warmup.py +++ /dev/null @@ -1,29 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Warm up v1 block-table Triton kernels.""" - -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from vllm.v1.worker.gpu_model_runner import GPUModelRunner - -_SLOT_MAPPING_WARMUP_TOKENS = 8 - - -def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None: - """JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables.""" - - device = runner.device - block_table = runner.input_batch.block_table - num_tokens = min( - _SLOT_MAPPING_WARMUP_TOKENS, - runner.scheduler_config.max_num_batched_tokens, - ) - if num_tokens <= 0: - return - - query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) - positions = torch.arange(num_tokens, dtype=torch.int64, device=device) - block_table.compute_slot_mapping(1, query_start_loc, positions) diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 5ef6fbf03c7c..fe2085af1109 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -9,7 +9,6 @@ import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig -from vllm.distributed import tensor_model_parallel_all_reduce from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, @@ -37,6 +36,7 @@ get_pp_missing_layer_names, maybe_prefix, ) +from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm from vllm.models.common.ops.sequence_parallel import ( sp_all_gather, sp_padding_mask, @@ -119,9 +119,6 @@ def forward( hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None ) - if not is_sequence_parallel: - # Without sequence parallelism, the MoE output is left un-reduced. - hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Recycle the POST-final-norm hidden into the next draft step. The # residual-add is fused into the final RMSNorm so it is computed # exactly once, and the result is returned for both tuple positions: @@ -132,9 +129,15 @@ def forward( # is understood by both the V2 speculator (isinstance-tuple check) and # the legacy proposer (model_returns_tuple is True for the # DeepSeekMTPModel architecture). - hidden_states, _ = self.shared_head.norm(hidden_states, residual) if is_sequence_parallel: + hidden_states, _ = self.shared_head.norm(hidden_states, residual) hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]] + else: + # The MoE output is left un-reduced; fuse its all-reduce into the + # final norm, as the main model does at layer boundaries. + hidden_states, _ = fused_allreduce_rms_norm( + hidden_states, residual, self.shared_head.norm + ) return hidden_states, hidden_states @@ -212,6 +215,23 @@ def compute_logits( # second RMSNorm. return self.logits_processor(mtp_layer.shared_head.head, hidden_states) + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """Greedy draft token ids via per-rank argmax over the vocab shard. + + Saves the full-vocab all-gather ``compute_logits`` does; same tokens. + Name is fixed by the protocol the proposer probes for + (``use_local_argmax_reduction``). + """ + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + return self.logits_processor.get_top_tokens( + mtp_layer.shared_head.head, hidden_states + ) + class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -262,6 +282,14 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def get_top_tokens( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """See ``DeepseekV32MultiTokenPredictor.get_top_tokens``.""" + return self.model.get_top_tokens(hidden_states, spec_step_idx) + def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str: spec_layer_weight_names = [ "embed_tokens", diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index 23223acd2dba..5e7609bd1191 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -19,6 +19,7 @@ DeepseekV4SparseMLAMetadataBuilder, ) from vllm.platforms import current_platform +from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import ( CommonAttentionMetadata, @@ -36,6 +37,19 @@ from vllm.v1.worker.workspace import current_workspace_manager +def _trust_dsv4_extra_cache_nan_free( + kv_cache_dtype: str, + has_kv_transfer: bool, + has_extra_cache: bool, +) -> bool: + return ( + _ON_GFX950 + and kv_cache_dtype == "fp8_ds_mla" + and not has_kv_transfer + and has_extra_cache + ) + + def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor: lengths = lengths.to(dtype=torch.int32).contiguous() indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device) @@ -294,9 +308,13 @@ def _copy_ragged_to_graph_buffers( max_entries = max(num_rows * max_entries_per_row, 1) ragged_out = ragged_indices_buffer[:max_entries] - nnz = ragged_indices.numel() - if nnz > 0: - ragged_out[:nnz].copy_(ragged_indices, non_blocking=True) + source_entries = ragged_indices.numel() + if source_entries > 0: + ragged_out[:source_entries].copy_(ragged_indices, non_blocking=True) + if _ON_GFX950: + # Preserve the graph-stable base pointer while exposing source capacity + # to the sync-free split selector; indptr still carries the true NNZ. + ragged_out = ragged_out[: max(source_entries, 1)] return ragged_out, indptr_out @@ -306,6 +324,7 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata): c128a_decode_topk_ragged_indices: torch.Tensor | None = None c128a_decode_topk_ragged_indptr: torch.Tensor | None = None + for_cudagraph_capture: bool = False @dataclass @@ -370,6 +389,16 @@ def build( c128a_decode_topk_ragged_indptr=ragged_indptr, ) + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> DeepseekV4ROCMAiterMLASparseMetadata: + metadata = cast( + DeepseekV4ROCMAiterMLASparseMetadata, + super().build_for_cudagraph_capture(common_attn_metadata), + ) + metadata.for_cudagraph_capture = _ON_GFX950 + return metadata + class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): # Keep fused multi-step decode disabled until update_draft_decode_metadata() @@ -414,7 +443,9 @@ def build( and base.decode_swa_lens is not None ): ragged_indices, ragged_indptr = build_ragged_indices_from_dense( - base.decode_swa_indices.reshape(base.num_decode_tokens, -1), + base.decode_swa_indices.reshape( + base.num_decode_tokens, base.decode_swa_width + ), base.decode_swa_lens, ) ragged_indices, ragged_indptr = _copy_ragged_to_graph_buffers( @@ -423,9 +454,7 @@ def build( self.decode_swa_ragged_indices_buffer, self.decode_swa_ragged_indptr_buffer, base.num_decode_tokens, - # Actual dense width for this build: window_size (causal) or - # noncausal_index_width (DSpark non-causal draft). - base.decode_swa_indices.shape[-1], + base.decode_swa_width, ) return DeepseekV4ROCMAiterSparseSWAMetadata( @@ -451,7 +480,9 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention): backend_cls = DeepseekV4ROCMAiterMLASparseBackend def __init__(self, *args, **kwargs): + vllm_config = args[0] if args else kwargs["vllm_config"] super().__init__(*args, **kwargs) + self._has_kv_transfer = vllm_config.kv_transfer_config is not None # Block scale for the preshuffled weight; None = not preshuffled. self._wqa_wkv_scale: torch.Tensor | None = None self._wo_b_scale: torch.Tensor | None = None @@ -605,6 +636,13 @@ def forward_mqa( attn_metadata=rocm_metadata, swa_only=swa_only, output=output[:num_decode_tokens], + adaptive_splits=( + _ON_GFX950 + and not swa_only + and self.compress_ratio == 128 + and rocm_metadata is not None + and rocm_metadata.for_cudagraph_capture + ), ) def _forward_decode( @@ -615,6 +653,7 @@ def _forward_decode( attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None, swa_only: bool, output: torch.Tensor, + adaptive_splits: bool, ) -> None: num_decodes = swa_metadata.num_decodes num_decode_tokens = swa_metadata.num_decode_tokens @@ -666,6 +705,12 @@ def _forward_decode( nope_head_dim=self.nope_head_dim, rope_head_dim=self.rope_head_dim, output=output, + adaptive_splits=adaptive_splits, + extra_cache_nan_free=_trust_dsv4_extra_cache_nan_free( + self.kv_cache_dtype, + self._has_kv_transfer, + not swa_only and kv_cache is not None, + ), ) def _forward_prefill( diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 5b85daba52d2..128debf70cb1 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -56,6 +56,7 @@ from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.mla.indexer import ( DeepseekV4IndexerBackend, + dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache @@ -298,6 +299,13 @@ def __init__( eager_scratch_pool=eager_scratch_pool, ) + self._prepare_and_attn_fn = self._prepare_and_attn + if not vllm_config.use_v2_model_runner: + # MRV1's piecewise capture only tolerates the wide eager region: with + # the narrow one the attention input preparation stays in the captured + # graph and MRV1 produces garbage (#51430). + self._prepare_and_attn_fn = self._prepare_and_attn_eager + # Will be None on ROCm for now. self.aux_stream_list = aux_stream_list # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; @@ -379,6 +387,64 @@ def forward( self.eps, ) + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + @eager_break_during_capture + def _prepare_and_attn_eager( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly. + + The nested ``_sparse_indexer_and_attn`` break runs inline, since + ``add_eager`` clears ``_capturing`` before invoking this. + """ + self._prepare_and_attn( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """Attention input preparation followed by the sparse indexer and MLA. + + Only the latter runs in the eager break. + """ attn_metadata = get_forward_context().attn_metadata indexer = self.indexer compressor = self.compressor @@ -438,10 +504,6 @@ def project_query_and_cache_kv() -> torch.Tensor: positions, o_padded, ) - o = o_padded[:, : self.n_local_heads, :] - - # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). - return self._o_proj(o, positions) def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: # Override point: the ROCm layer preshuffles this weight in place, so @@ -670,6 +732,9 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: alignment=576 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, ) @@ -741,7 +806,7 @@ def __init__( self.q_lora_rank = q_lora_rank # 1536 self.compress_ratio = compress_ratio self.eager_scratch_pool = eager_scratch_pool - self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache + self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config) logger.info_once( "Using %s indexer cache for Lightning Indexer.", "MXFP4" if self.use_fp4_kv else "FP8", @@ -841,7 +906,10 @@ def forward( attn_metadata = get_forward_context().attn_metadata if isinstance(attn_metadata, dict): indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) - if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens: + if ( + indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens + and not torch.cuda.is_current_stream_capturing() + ): # candidates num smaller than topk, every candidate is selected # but we still need to build k cache compressor(compressed_kv_score, positions, rotary_emb) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index dc06dc91b22f..1f53c96e9367 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -477,15 +477,15 @@ def compute_global_topk_indices_and_lens( @triton.jit def _compute_global_topk_indices_and_lens_kernel( global_topk_indices_ptr, - global_topk_indices_stride, + global_topk_indices_stride: tl.constexpr, topk_lens_ptr, topk_indices_ptr, - topk_indices_stride, - topk, + topk_indices_stride: tl.constexpr, + topk: tl.constexpr, token_to_req_indices_ptr, block_table_ptr, - block_table_stride, - block_size, + block_table_stride: tl.constexpr, + block_size: tl.constexpr, is_valid_token_ptr, TRITON_BLOCK_SIZE: tl.constexpr, ): @@ -580,7 +580,7 @@ def combine_topk_swa_indices( return combined_indices, combined_lens -_COMBINE_TOPK_SWA_NUM_WORKERS = 128 +_COMBINE_TOPK_SWA_NUM_WORKERS = 256 # Representative pointer alignment variants for Triton pointer specialization. @@ -839,16 +839,17 @@ def build_flashinfer_mixed_sparse_indices( ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. - Produces ``sparse_indices`` of shape ``[num_tokens, window_size + - padded_topk]`` (the first ``window_size`` columns are SWA slot ids, the rest - are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length per - token). Decode tokens read precomputed SWA/compressed indices; prefill tokens - derive their SWA window from the position and translate local compressed - indices to global slots via the block tables. + Produces ``sparse_indices`` of shape ``[num_tokens, swa_index_width + + padded_topk]`` (the first ``swa_index_width`` columns are SWA slot ids, the + rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length + per token). Decode tokens read precomputed SWA/compressed indices; prefill + tokens derive their SWA window from the position and translate local + compressed indices to global slots via the block tables. """ assert decode_swa_indices.dtype == torch.int32 assert decode_swa_indices.dim() == 2 - assert decode_swa_indices.shape[-1] == window_size + swa_index_width = decode_swa_indices.shape[-1] + assert swa_index_width >= window_size if decode_compressed_topk_lens is not None: assert decode_compressed_topk_lens.dtype == torch.int32 assert prefill_topk_indices.dtype == torch.int32 @@ -897,7 +898,7 @@ def build_flashinfer_mixed_sparse_indices( padded_topk = max(topk, decode_compressed_topk) padded_topk = (padded_topk + 3) // 4 * 4 sparse_indices = torch.empty( - (num_tokens, window_size + padded_topk), + (num_tokens, swa_index_width + padded_topk), dtype=torch.int32, device=decode_swa_indices.device, ) @@ -907,7 +908,7 @@ def build_flashinfer_mixed_sparse_indices( if num_tokens == 0: return sparse_indices, sparse_topk_lens - window_block_size = triton.next_power_of_2(max(window_size, 1)) + window_block_size = triton.next_power_of_2(max(swa_index_width, 1)) topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 @@ -944,6 +945,7 @@ def build_flashinfer_mixed_sparse_indices( compressed_span, NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, + SWA_INDEX_WIDTH=swa_index_width, COMPRESS_RATIO=compress_ratio, TOP_K=topk, PADDED_TOP_K=padded_topk, @@ -1011,6 +1013,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( compressed_block_span, NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, + SWA_INDEX_WIDTH: tl.constexpr, COMPRESS_RATIO: tl.constexpr, TOP_K: tl.constexpr, PADDED_TOP_K: tl.constexpr, @@ -1024,9 +1027,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( token_idx = tl.program_id(0) if token_idx < NUM_DECODE_TOKENS: - for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < WINDOW_SIZE + mask = offset < SWA_INDEX_WIDTH values = tl.load( decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, mask=mask, @@ -1072,7 +1075,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + WINDOW_SIZE + + SWA_INDEX_WIDTH + offset, values, mask=mask, @@ -1086,7 +1089,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( else: compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + compressed_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len) return prefill_idx = token_idx - NUM_DECODE_TOKENS @@ -1102,9 +1105,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( swa_start_pos = pos - swa_len + 1 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < WINDOW_SIZE + mask = offset < SWA_INDEX_WIDTH pos_offset = swa_start_pos + offset block_indices = pos_offset // swa_block_size block_numbers = tl.load( @@ -1148,10 +1151,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + WINDOW_SIZE + + SWA_INDEX_WIDTH + offset, slot_ids, mask=mask, ) - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + topk_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len) diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index a2085cd220f1..4c9f464ef067 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -24,8 +24,14 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +if current_platform.is_rocm(): + from vllm.platforms.rocm import _ON_GFX950 +else: + _ON_GFX950 = False + from .fused_indexer_q import _fp32x2_to_fp4x2 @@ -61,12 +67,15 @@ def compress_norm_rope_store_triton( if head_dim == 512: kernel = _fused_kv_compress_norm_rope_insert_sparse_attn num_warps = 4 + kernel_kwargs = {"SANITIZE_CACHE_NANS": _ON_GFX950} elif use_fp4_cache: kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn num_warps = 1 + kernel_kwargs = {} else: kernel = _fused_kv_compress_norm_rope_insert_indexer_attn num_warps = 1 + kernel_kwargs = {} kernel[(num_actual,)]( # state cache @@ -103,6 +112,7 @@ def compress_norm_rope_store_triton( SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), num_warps=num_warps, + **kernel_kwargs, **pdl_kwargs, ) @@ -145,6 +155,7 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( TOKEN_STRIDE: tl.constexpr, # 576 for DeepseekV4 SCALE_DIM: tl.constexpr, # 8 for DeepseekV4 (7 real + 1 pad) KV_BLOCK_STRIDE: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, ): """Fused compress → RMSNorm → FP8 quant (nope) → RoPE → bf16 store (rope). @@ -261,7 +272,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( scale_idx = tl.arange(0, N_QUANT_BLOCKS) encoded = exponents + 127.0 - encoded = tl.maximum(tl.minimum(encoded, 255.0), 0.0) + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.maximum(tl.minimum(encoded, max_encoded), 0.0) tl.store( scale_ptr + scale_idx, encoded.to(tl.uint8), @@ -289,6 +301,8 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( new_even = even * cos_v - odd * sin_v new_odd = odd * cos_v + even * sin_v result = tl.interleave(new_even, new_odd) # [TRITON_BLOCK_SIZE] fp32 + if SANITIZE_CACHE_NANS: + result = tl.where(result == result, result, 0.0) # Store rotated rope portion as bf16 into the cache's bf16 area. bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) @@ -417,6 +431,7 @@ def _finalize_norm_rope_quant_store_sparse_attn( TOKEN_STRIDE: tl.constexpr, SCALE_DIM: tl.constexpr, KV_BLOCK_STRIDE: tl.constexpr, + SANITIZE_CACHE_NANS: tl.constexpr, ): """Stage 2: read compressed_kv[512] from scratch buffer, then RMSNorm + FP8 quant (nope) + RoPE + bf16 store @@ -474,7 +489,8 @@ def _finalize_norm_rope_quant_store_sparse_attn( tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) scale_idx = tl.arange(0, N_QUANT_BLOCKS) - encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0 + encoded = tl.maximum(tl.minimum(exponents + 127.0, max_encoded), 0.0) tl.store( scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS ) @@ -494,6 +510,8 @@ def _finalize_norm_rope_quant_store_sparse_attn( new_even = even * cos_v - odd * sin_v new_odd = odd * cos_v + even * sin_v result = tl.interleave(new_even, new_odd) + if SANITIZE_CACHE_NANS: + result = tl.where(result == result, result, 0.0) bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) rope_local = block - NOPE_HEAD_DIM is_rope = (block >= NOPE_HEAD_DIM) & mask @@ -564,6 +582,7 @@ def _launch_two_stage_sparse_attn_compressor( TOKEN_STRIDE=token_stride, SCALE_DIM=scale_dim, KV_BLOCK_STRIDE=kv_cache.stride(0), + SANITIZE_CACHE_NANS=_ON_GFX950, ) diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 27031e02c577..e73156c1c121 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -6,6 +6,7 @@ import torch +from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.attention import DeepseekV4Attention @@ -25,6 +26,9 @@ from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 from vllm.v1.attention.backend import MultipleOf +from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, +) if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -74,6 +78,19 @@ def _pad_to_supported_q_heads(num_heads: int) -> int: ) +def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int: + """Return the SM120 DSV4 SWA specialization needed by this model.""" + if not vllm_config.attention_config.use_non_causal: + return window_size + speculative_config = vllm_config.speculative_config + if speculative_config is None: + return window_size + return get_dspark_swa_index_width( + window_size, + speculative_config.num_speculative_tokens, + ) + + class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend): """FlashInfer backend using the DSv4 sparse metadata/cache layout. @@ -309,7 +326,7 @@ def _build_sparse_index_metadata( assert swa_metadata.block_table is not None decode_swa_indices = swa_metadata.decode_swa_indices.reshape( - num_decode_tokens, self.window_size + num_decode_tokens, swa_metadata.decode_swa_width ) decode_compressed_topk_lens = None decode_compressed_indices_are_local = False @@ -568,14 +585,18 @@ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: tma_aligned_scales=self._tma_aligned_scales, ) - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 + def __init__(self, vllm_config: VllmConfig, *args, **kwargs) -> None: + super().__init__(vllm_config, *args, **kwargs) + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config - if not has_flashinfer_sparse_mla_sm120(): + required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size) + if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk): raise RuntimeError( - "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires FlashInfer's " - "sparse MLA decode API." + "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer " + "DSV4 sparse MLA decode specialization for " + f"(num_q_heads={self.padded_heads}, top_k={required_topk}). " + "Install a FlashInfer build containing " + "flashinfer-ai/flashinfer#4380." ) self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe() # Per-tensor FP8 cache path scales. diff --git a/vllm/models/deepseek_v4/quant_config.py b/vllm/models/deepseek_v4/quant_config.py index 293d71f2f414..2ffedc97d315 100644 --- a/vllm/models/deepseek_v4/quant_config.py +++ b/vllm/models/deepseek_v4/quant_config.py @@ -192,8 +192,3 @@ def get_quant_method(self, layer, prefix): # expert_dtype == "fp8": fall through to Fp8Config which # returns Fp8MoEMethod with block-wise float32 scales. return super().get_quant_method(layer, prefix) - - def is_mxfp4_quant(self, prefix, layer): - if not isinstance(layer, RoutedExperts) or self.expert_dtype != "fp4": - return False - return self.moe_quant_algo != "NVFP4" diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index db8ab96e90ff..0475e4c96df7 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -257,13 +257,6 @@ def _build_c128a_metadata( assert cm.positions is not None, ( "positions is required for C128A metadata build" ) - active_topk_width = min( - max( - triton.next_power_of_2(max(cm.max_seq_len // self.compress_ratio, 1)), - _C128A_TOPK_ALIGNMENT, - ), - self.c128a_max_compressed, - ) block_size = self.kv_cache_spec.block_size // self.compress_ratio global_decode, decode_lens, prefill_local = build_c128a_topk_metadata( cm.positions[:num_total], @@ -276,7 +269,7 @@ def _build_c128a_metadata( self.c128a_global_decode_buffer, self.c128a_decode_lens_buffer, self.c128a_prefill_buffer, - max_compressed_tokens=active_topk_width, + max_compressed_tokens=self.c128a_max_compressed, ) result: dict[str, torch.Tensor | None] = {} @@ -322,30 +315,25 @@ def build_c128a_topk_metadata( Decode tokens: position → block_table lookup → global slot ids + topk_lens. Prefill tokens: position → local indices [0, ..., n-1, -1, ...]. - Writes into packed views of pre-allocated buffers for CUDA graph stability. + Writes into pre-allocated buffers for CUDA graph address stability. + Returns slices of the buffers. """ num_tokens = positions.shape[0] num_prefill_tokens = num_tokens - num_decode_tokens - # view(-1) as 1-d array and then expanded to - # [num_decode_tokens, max_compressed_tokens] - global_decode = global_decode_buffer.view(-1)[ - : num_decode_tokens * max_compressed_tokens - ].view(num_decode_tokens, max_compressed_tokens) + global_decode = global_decode_buffer[:num_decode_tokens] decode_lens = decode_lens_buffer[:num_decode_tokens] - prefill_local = prefill_buffer.view(-1)[ - : num_prefill_tokens * max_compressed_tokens - ].view(num_prefill_tokens, max_compressed_tokens) + prefill_local = prefill_buffer[:num_prefill_tokens] if num_tokens == 0: return global_decode, decode_lens, prefill_local _build_c128a_topk_metadata_kernel[(num_tokens,)]( global_decode_buffer, - max_compressed_tokens, + global_decode_buffer.stride(0), decode_lens_buffer, prefill_buffer, - max_compressed_tokens, + prefill_buffer.stride(0), positions, compress_ratio, max_compressed_tokens, diff --git a/vllm/models/dots3_note/nvidia/model.py b/vllm/models/dots3_note/nvidia/model.py index 83e675c4692c..e713911e619c 100644 --- a/vllm/models/dots3_note/nvidia/model.py +++ b/vllm/models/dots3_note/nvidia/model.py @@ -497,6 +497,7 @@ def __init__( layer_idx = int(prefix.split(sep=".")[-1]) self.layer_idx = layer_idx self.use_mha = False + self.use_sequence_parallel = False attention_cls = ( Dots3NoteSlidingAttention if config.layer_types[layer_idx] == "sliding_attention" @@ -536,11 +537,7 @@ def __init__( prefix=f"{prefix}.mlp", reduce_results=False, ) - self.use_sequence_parallel_moe = ( - parallel_config.use_sequence_parallel_moe - and parallel_config.pipeline_parallel_size == 1 - and isinstance(self.mlp, DeepseekV2MoE) - ) + self.use_sequence_parallel_moe = False self.tp_size = parallel_config.tensor_parallel_size self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( @@ -559,6 +556,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.config = config self.device = current_platform.device_type self.vocab_size = config.vocab_size + self.use_sequence_parallel = False self.is_v32 = True self._weight_block_size = getattr(quant_config, "weight_block_size", None) topk_indices_buffer = torch.empty( diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index b59293caebfa..e680180f26fa 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -282,31 +282,52 @@ def get_attn_backend(self) -> type[AttentionBackend]: def get_state_dtype( self, - ) -> tuple[torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, ...]: if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") - return MambaStateDtypeCalculator.kda_state_dtype( + base_dtypes = MambaStateDtypeCalculator.kda_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype ) + if self.cache_config.use_kda_recoverssm: + return MambaStateDtypeCalculator.append_kda_recoverssm_record( + base_dtypes, self.model_config.dtype + ) + return base_dtypes def get_state_shape( self, - ) -> tuple[tuple[int, ...], tuple[int, ...]]: - return MambaStateShapeCalculator.kda_state_shape( + ) -> tuple[tuple[int, ...], ...]: + base_shapes = MambaStateShapeCalculator.kda_state_shape( self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size, num_spec=self.num_spec, ) + if self.cache_config.use_kda_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + base_shapes, + self.num_heads, + self.head_dim, + tp_world_size=self.tp_size, + spec_query_len=1 + self.num_spec, + ) + return base_shapes def __init__( self, config: KimiLinearConfig, vllm_config: VllmConfig, prefix: str = "", + run_gemm_rs: bool = False, ) -> None: super().__init__(config, vllm_config, prefix) + self.use_recoverssm = self.cache_config.use_kda_recoverssm + if self.cache_config.use_replayssm and not self.use_recoverssm: + raise ValueError( + "Kimi-K3 supports --use-replayssm only with speculative decoding" + ) + self.spec_query_len = 1 + self.num_spec kda_config = config.linear_attn_config # type: ignore[attr-defined] assert kda_config is not None, "linear_attn_config must be set" @@ -370,7 +391,7 @@ def __init__( self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) # Keep a width-major copy for fused decode without changing the layout # consumed by the prefill and fallback decode kernels. - conv_state_dtype, _ = self.get_state_dtype() + conv_state_dtype = self.get_state_dtype()[0] decode_conv1d_weight = None if is_fused_kda_decode_supported( self.local_num_heads, @@ -470,7 +491,16 @@ def __init__( quant_config=self.quant_config, prefix=f"{prefix}.o_proj", ) - + self.run_gemm_rs = run_gemm_rs + if self.run_gemm_rs: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs + + self.run_gemm_rs = get_gemm_rs().can_run(self.o_proj) + if not self.run_gemm_rs: + logger.warning_once( + "GEMM-RS is disabled for %s due to an incompatible projection.", + prefix, + ) compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -511,6 +541,12 @@ def forward( core_attn_out=core_attn_out, ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") + if self.run_gemm_rs: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs + + gemm_rs = get_gemm_rs() + if gemm_rs.should_run(core_attn_out): + return gemm_rs(core_attn_out, self.o_proj.weight) return self.o_proj(core_attn_out)[0] @eager_break_during_capture @@ -551,7 +587,7 @@ def _forward( g1 = g1[:, :num_actual_tokens] beta = beta[:, :num_actual_tokens] - conv_state, recurrent_state = self.kv_cache + conv_state, recurrent_state, *recoverssm_records = self.kv_cache # The convolution kernels consume (..., dim, width - 1). if not is_conv_state_dim_first(): conv_state = conv_state.transpose(-1, -2) @@ -618,7 +654,11 @@ def _forward( assert spec_state_indices_tensor is not None assert spec_query_start_loc is not None spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes] - spec_max_query_len = spec_state_indices_tensor.size(-1) + spec_max_query_len = ( + self.spec_query_len + if self.use_recoverssm + else spec_state_indices_tensor.size(-1) + ) spec_conv_out = torch.empty_like(mixed_qkv_spec) mixed_qkv_spec = causal_conv1d_update( mixed_qkv_spec, @@ -643,21 +683,48 @@ def _forward( if m.num_prefills == 0 and m.num_decodes == 0 else None ) - core_attn_out_spec, _ = fused_recurrent_kda( - q=q_spec, - k=k_spec, - v=v_spec, - raw_g=g1_spec, - raw_beta=beta_spec, - A_log=self.A_log, - dt_bias=self.dt_bias, - lower_bound=self.gate_lower_bound, - initial_state=recurrent_state, - cu_seqlens=spec_cu_seqlens, - ssm_state_indices=spec_state_indices_tensor, - num_accepted_tokens=num_accepted_tokens, - out=spec_out, - ) + if self.use_recoverssm: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + kda_recoverssm_verify, + ) + + if len(recoverssm_records) != 2: + raise ValueError( + "KDA RecoverSSM requires correction and key/gate buffers" + ) + core_attn_out_spec = kda_recoverssm_verify( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + checkpoint_state=recurrent_state, + correction_cache=recoverssm_records[0], + kg_cache=recoverssm_records[1], + query_start_loc=spec_cu_seqlens, + state_indices=spec_state_indices_tensor[: m.num_spec_decodes, 0], + spec_query_len=self.spec_query_len, + out=spec_out, + ) + else: + core_attn_out_spec, _ = fused_recurrent_kda( + q=q_spec, + k=k_spec, + v=v_spec, + raw_g=g1_spec, + raw_beta=beta_spec, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=recurrent_state, + cu_seqlens=spec_cu_seqlens, + ssm_state_indices=spec_state_indices_tensor, + num_accepted_tokens=num_accepted_tokens, + out=spec_out, + ) # Prefill or plain-decode path. core_attn_out_non_spec = None diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 0bb4864d2918..2154eaa3c51f 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -6,13 +6,18 @@ ``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its prefill KDA kernel internally, so this builder omits the shared FLA chunk metadata construction. + +For Kimi-K3 speculative decoding, ``--use-replayssm`` selects the simplified +RecoverSSM path implemented here instead of the Mamba2 ReplaySSM kernel. """ -from dataclasses import dataclass +from dataclasses import dataclass, field from functools import cache +from typing import TYPE_CHECKING import torch +from vllm.config import VllmConfig from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import async_tensor_h2d @@ -22,6 +27,10 @@ GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.recoverssm_metadata import ( + RecoverSSMMetadata, + RecoverSSMPostprocessMetadata, +) from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, compute_causal_conv1d_metadata, @@ -29,6 +38,11 @@ ) from vllm.v1.kv_cache_interface import MambaSpec +if TYPE_CHECKING: + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + ) + @cache def _metadata_launch_pdl() -> bool: @@ -229,11 +243,103 @@ def stage_spec_decode_metadata( @dataclass -class KimiK3KDAMetadata(GDNAttentionMetadata): - pass +class KDARecoverSSMAlignMetadata: + block_table: torch.Tensor + num_computed_tokens: torch.Tensor + block_size: int + + +@dataclass +class KDARecoverSSMCommitMetadata: + state_indices: torch.Tensor + query_start_loc: torch.Tensor + request_indices: torch.Tensor | None + align: KDARecoverSSMAlignMetadata | None + + +@dataclass +class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + recoverssm_commit: KDARecoverSSMCommitMetadata | None = None + recoverssm_context: "KDARecoverSSMCommitContext | None" = field( + default=None, repr=False, compare=False + ) + + def commit_recoverssm_state( + self, num_accepted_tokens: torch.Tensor + ) -> RecoverSSMPostprocessMetadata | None: + commit = self.recoverssm_commit + if commit is None: + return None + context = self.recoverssm_context + assert context is not None + align = commit.align + context.commit( + num_accepted_tokens, + commit.state_indices[: self.num_spec_decodes, 0], + commit.query_start_loc[: self.num_spec_decodes + 1], + request_indices=commit.request_indices, + block_table=align.block_table if align is not None else None, + num_computed_tokens=( + align.num_computed_tokens if align is not None else None + ), + mamba_block_size=align.block_size if align is not None else None, + ) + if align is None: + return None + return RecoverSSMPostprocessMetadata( + num_spec_decodes=self.num_spec_decodes, + request_indices=commit.request_indices, + block_table=align.block_table, + num_computed_tokens=align.num_computed_tokens, + block_size=align.block_size, + ) class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): + def __init__( + self, + kv_cache_spec: MambaSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm + self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1 + self.recoverssm_num_accepted_tokens: torch.Tensor | None = None + self.recoverssm_context: KDARecoverSSMCommitContext | None = None + if self.use_recoverssm: + max_num_reqs = vllm_config.scheduler_config.max_num_seqs + self.spec_state_indices_tensor = torch.empty( + (max_num_reqs, 1), + dtype=torch.int32, + device=device, + ) + self.recoverssm_num_accepted_tokens = torch.ones( + max_num_reqs, + dtype=torch.int32, + device=device, + ) + + def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext": + context = self.recoverssm_context + if context is not None: + return context + + from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( + KDARecoverSSMCommitContext, + ) + + forward_context = self.vllm_config.compilation_config.static_forward_context + layers = [forward_context[layer_name] for layer_name in self.layer_names] + context = KDARecoverSSMCommitContext.create( + layers, + spec_query_len=1 + self.vllm_config.num_speculative_tokens, + max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, + ) + self.recoverssm_context = context + return context + def build( # type: ignore[override] self, common_prefix_len: int, @@ -263,14 +369,26 @@ def build( # type: ignore[override] num_spec_decodes = 0 else: spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 - # A nonnegative entry identifies a spec request. If no draft token - # was scheduled, process the whole batch as non-spec instead. - if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0: + if self.use_recoverssm: + assert m.is_prefilling is not None + assert m.is_prefilling.device.type == "cpu" + active_decode_mask_cpu = (~m.is_prefilling) & ( + query_start_loc_cpu.diff() > 0 + ) + spec_sequence_masks_cpu |= active_decode_mask_cpu + # Native KDA can use its regular decode path when no draft token + # was scheduled. RecoverSSM must preserve its extended conv window. + if ( + not self.use_recoverssm + and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() + == 0 + ): spec_sequence_masks_cpu = None num_spec_decodes = 0 else: num_spec_decodes = spec_sequence_masks_cpu.sum().item() + spec_request_indices = None if num_spec_decodes == 0: # The runner orders ordinary decodes before prefills. num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( @@ -289,6 +407,16 @@ def build( # type: ignore[override] assert spec_sequence_masks_cpu is not None assert num_accepted_tokens is not None query_lens_cpu = query_start_loc_cpu.diff() + if ( + self.use_recoverssm + and torch.any( + query_lens_cpu[spec_sequence_masks_cpu] > self.num_spec + 1 + ).item() + ): + raise ValueError( + "KDA RecoverSSM speculative decode query length exceeds " + f"its activation capacity ({self.num_spec + 1})" + ) num_query_tokens = query_start_loc_cpu[-1].item() # Exclude zero-length cudagraph padding from request-indexed @@ -325,7 +453,7 @@ def build( # type: ignore[override] non_spec_token_indx = None # Real requests precede trailing cudagraph padding. spec_state_indices_tensor = block_table_tensor[ - :num_spec_decodes, : self.num_spec + 1 + :num_spec_decodes, : self.spec_state_slots ] non_spec_state_indices_tensor = None # Padding trails real requests, so this prefix already contains @@ -339,6 +467,12 @@ def build( # type: ignore[override] spec_sequence_masks_gpu = async_tensor_h2d( spec_sequence_masks_cpu, device=query_start_loc.device ) + if self.use_recoverssm: + spec_request_indices = async_tensor_h2d( + spec_sequence_masks_cpu.nonzero(as_tuple=True)[0], + dtype=torch.int32, + device=query_start_loc.device, + ) spec_token_masks = torch.repeat_interleave( spec_sequence_masks_gpu, query_lens, @@ -351,10 +485,10 @@ def build( # type: ignore[override] non_spec_token_indx = index[:num_non_spec_tokens] spec_token_indx = index[num_non_spec_tokens:] - # Spec requests carry one state slot per speculative step; - # non-spec requests use only their current state slot. + # Native spec uses one state slot per step. RecoverSSM keeps + # only the current checkpoint slot. spec_state_indices_tensor = block_table_tensor[ - spec_sequence_masks_cpu, : self.num_spec + 1 + spec_sequence_masks_cpu, : self.spec_state_slots ] non_spec_state_indices_tensor = block_table_tensor[ active_non_spec_mask_cpu, 0 @@ -400,12 +534,18 @@ def build( # type: ignore[override] num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] + if self.use_recoverssm: + assert self.recoverssm_num_accepted_tokens is not None + num_accepted_tokens = self.recoverssm_num_accepted_tokens[ + :num_spec_decodes + ] + # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares # its own chunk indices. Only causal-convolution metadata is needed here. nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None if num_prefills > 0: has_initial_state = m.compute_num_computed_tokens() > 0 - if spec_sequence_masks_cpu is not None: + if num_spec_decodes > 0: has_initial_state = has_initial_state[active_non_spec_mask_cpu] assert non_spec_query_start_loc_cpu is not None nums_dict, batch_ptr, token_chunk_offset_ptr = ( @@ -426,7 +566,7 @@ def build( # type: ignore[override] and num_spec_decodes > 0 and num_prefills == 0 and num_decodes == 0 - and num_spec_decodes <= self.decode_cudagraph_max_bs + and batch_size <= self.spec_state_indices_tensor.shape[0] and num_spec_decode_tokens <= self.decode_cudagraph_max_bs ): # Equivalent PyTorch staging: @@ -463,6 +603,24 @@ def build( # type: ignore[override] :batch_size ] + recoverssm_commit = None + if self.use_recoverssm and num_spec_decodes > 0: + assert spec_state_indices_tensor is not None + assert spec_query_start_loc is not None + align = None + if self.kv_cache_spec.mamba_cache_mode == "align": + align = KDARecoverSSMAlignMetadata( + block_table=m.block_table_tensor, + num_computed_tokens=m.compute_num_computed_tokens(), + block_size=self.kv_cache_spec.block_size, + ) + recoverssm_commit = KDARecoverSSMCommitMetadata( + state_indices=spec_state_indices_tensor, + query_start_loc=spec_query_start_loc, + request_indices=spec_request_indices, + align=align, + ) + return KimiK3KDAMetadata( num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, @@ -480,6 +638,12 @@ def build( # type: ignore[override] spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, num_accepted_tokens=num_accepted_tokens, + recoverssm_commit=recoverssm_commit, + recoverssm_context=( + self._get_recoverssm_context() + if recoverssm_commit is not None + else None + ), nums_dict=nums_dict, batch_ptr=batch_ptr, token_chunk_offset_ptr=token_chunk_offset_ptr, diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 708104b0736a..2c67a8a4c4a6 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -134,6 +134,7 @@ def __init__( aux_stream: torch.cuda.Stream | None = None, use_rope: bool = False, non_causal_multi_token_decode: bool = False, + run_gemm_rs: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -270,6 +271,16 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.o_proj", ) + self.run_gemm_rs = run_gemm_rs + if self.run_gemm_rs: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs + + self.run_gemm_rs = get_gemm_rs().can_run(self.o_proj) + if not self.run_gemm_rs: + logger.warning_once( + "GEMM-RS is disabled for %s due to an incompatible projection.", + prefix, + ) # ---- Attention backend / impl / KV cache ---- self.quant_config = quant_config @@ -330,11 +341,6 @@ def __init__( "parallelism." ) self.dcp_world_size = parallel_config.decode_context_parallel_size - assert self.dcp_world_size <= 1 or self.rotary_emb is None, ( - "Kimi-K3 MultiHeadLatentAttention does not support RoPE with decode " - "context parallelism because gathered queries require gathered " - "positions." - ) self.dcp_manager: MLADCPManager | None = None if self.dcp_world_size > 1: query_dtype = ( @@ -389,6 +395,8 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: dtype=kv_cache_dtype, cache_dtype_str=self.kv_cache_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + # fp8_ds_mla: 656-byte custom layout; see flashmla_sparse.py. + state_content_bytes=656 if self.kv_cache_dtype == "fp8_ds_mla" else None, non_causal_multi_token_decode=self.non_causal_multi_token_decode, ) @@ -541,10 +549,13 @@ def forward( if gate is not None: attn_out = _gate_sigmoid_mul(attn_out, gate) - # ``o_proj`` (RowParallelLinear + out-of-place all-reduce) returns a - # fresh private tensor, so return it directly rather than copying into a - # caller buffer -- the previous ``output[:] = ...`` convention forced an - # extra [num_tokens, hidden] copy per layer. + if self.run_gemm_rs: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs + + gemm_rs = get_gemm_rs() + if gemm_rs.should_run(attn_out): + return gemm_rs(attn_out, self.o_proj.weight) + return self.o_proj(attn_out)[0] @eager_break_during_capture diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 59682fe32adf..f1166695fdeb 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -62,12 +62,14 @@ EagleModelMixin, HasInnerState, IsHybrid, + MambaStateShapes, MixtureOfExperts, SupportsEagle3, SupportsEncoderCudaGraph, SupportsMultiModal, SupportsPP, SupportsQuant, + SupportsReplaySSM, ) from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs from vllm.model_executor.models.kimi_k25_vit import ( @@ -133,6 +135,7 @@ def shard_sequence_parallel_mlp( hidden_size: int, intermediate_size: int, use_sequence_parallel: bool, + eligible: bool, ) -> bool: """Whether to TP-shard a sequence-parallel MLP instead of replicating it. @@ -140,7 +143,7 @@ def shard_sequence_parallel_mlp( the trade-off and :mod:`vllm.envs` for when it is worth enabling. """ enabled = envs.VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT - if not (use_sequence_parallel and enabled): + if not (use_sequence_parallel and eligible and enabled): return False tp_size = get_tensor_model_parallel_world_size() return ( @@ -148,6 +151,44 @@ def shard_sequence_parallel_mlp( ) +def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool: + if not envs.VLLM_KIMI_K3_GEMM_RS: + return False + + parallel_config = vllm_config.parallel_config + tp_size = parallel_config.tensor_parallel_size + if not use_sequence_parallel: + reason = "sequence parallelism is disabled" + elif parallel_config.use_ubatching: + reason = "ubatching is enabled" + elif vllm_config.model_config.dtype != torch.bfloat16: + reason = "the model dtype is not BF16" + elif not current_platform.is_cuda(): + reason = "the device is not CUDA" + elif not current_platform.is_device_capability_family(100): + reason = "the device is not SM100-family" + elif not 1 < tp_size <= 16: + reason = "TP size is not in the supported range 2-16" + elif 128 % tp_size != 0: + reason = "TP size does not divide 128" + else: + reason = None + + if reason is not None: + logger.warning_once("GEMM-RS was requested but is disabled because %s.", reason) + return False + + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import init_gemm_rs + + config = vllm_config.model_config.hf_text_config + init_gemm_rs( + max_M=vllm_config.scheduler_config.max_num_batched_tokens, + N=config.hidden_size, + ) + logger.info_once("GEMM-RS is enabled.") + return True + + class KimiMLP(nn.Module): """Dense / shared-expert MLP, optionally TP-sharded under sequence parallel. @@ -172,6 +213,8 @@ def __init__( quant_config: QuantizationConfig | None = None, reduce_results: bool = True, use_sequence_parallel: bool = False, + can_shard_sequence_parallel: bool = False, + run_gemm_rs: bool = False, prefix: str = "", activation_situ_beta: float | None = None, activation_situ_linear_beta: float | None = None, @@ -182,7 +225,9 @@ def __init__( hidden_size, intermediate_size, use_sequence_parallel, + can_shard_sequence_parallel, ) + self.run_gemm_rs = self.shard_sequence_parallel and run_gemm_rs replicate = use_sequence_parallel and not self.shard_sequence_parallel self.gate_up_proj = MergedColumnParallelLinear( @@ -204,6 +249,15 @@ def __init__( disable_tp=replicate, prefix=f"{prefix}.down_proj", ) + if self.run_gemm_rs: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs + + self.run_gemm_rs = get_gemm_rs().can_run(self.down_proj) + if not self.run_gemm_rs: + logger.warning_once( + "GEMM-RS is disabled for %s due to an incompatible projection.", + prefix, + ) if hidden_act == "silu": self.act_fn = SiluAndMul() elif hidden_act == "situ": @@ -226,6 +280,14 @@ def forward(self, x): x = sp_all_gather(x) gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) + + if self.run_gemm_rs: + from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs import get_gemm_rs + + gemm_rs = get_gemm_rs() + if gemm_rs.should_run(x): + return gemm_rs(x, self.down_proj.weight) + x, _ = self.down_proj(x) if self.shard_sequence_parallel: x = sp_reduce_scatter(x) @@ -426,8 +488,8 @@ def forward( symm_buffer, activation_clamp=activation_clamp, activation=self.activation, - activation_beta=self.activation_beta, - activation_linear_beta=self.activation_linear_beta, + situ_beta=self.activation_beta, + situ_linear_beta=self.activation_linear_beta, fast_math=fast_math, ) return y @@ -462,6 +524,7 @@ def __init__( prefix: str = "", layer_idx: int = 0, use_sequence_parallel: bool = False, + run_gemm_rs: bool = False, ): super().__init__() hidden_size = config.hidden_size @@ -546,6 +609,11 @@ def __init__( quant_config=quant_config, reduce_results=False, use_sequence_parallel=use_sequence_parallel, + # Only the MegaMoE path calls the shared experts directly; the + # FusedMoE path below hands them to the runner, which fuses + # their reduction and assumes the replicated layout. + can_shard_sequence_parallel=self.use_mega_moe, + run_gemm_rs=run_gemm_rs, prefix=f"{prefix}.shared_experts", activation_situ_beta=activation_situ_beta, activation_situ_linear_beta=activation_situ_linear_beta, @@ -767,6 +835,7 @@ def __init__( vllm_config: VllmConfig, prefix: str = "", aux_stream: torch.cuda.Stream | None = None, + run_gemm_rs: bool = False, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -801,6 +870,7 @@ def __init__( config, vllm_config, prefix=f"{prefix}.self_attn", + run_gemm_rs=run_gemm_rs, ) self._self_attn_writes_output = False else: @@ -837,6 +907,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", aux_stream=aux_stream, + run_gemm_rs=run_gemm_rs, ) self._self_attn_writes_output = False @@ -851,6 +922,7 @@ def __init__( prefix=f"{prefix}.block_sparse_moe", layer_idx=layer_idx, use_sequence_parallel=self.use_sequence_parallel, + run_gemm_rs=run_gemm_rs, ) self.mlp = self.block_sparse_moe else: @@ -861,6 +933,8 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.mlp", use_sequence_parallel=self.use_sequence_parallel, + can_shard_sequence_parallel=True, + run_gemm_rs=run_gemm_rs, activation_situ_beta=config.activation_situ_beta, activation_situ_linear_beta=config.activation_situ_linear_beta, ) @@ -991,15 +1065,18 @@ def forward( ) assert hidden_states is not None + M = None if self.use_sequence_parallel: hidden_states = sp_all_gather(hidden_states) # Remove SP padding before attention. hidden_states = hidden_states[: positions.shape[0]] + M = hidden_states.shape[0] # Attention. hidden_states = self._run_self_attn(positions, hidden_states) - if self.use_sequence_parallel: + # GEMM-RS returns the local sequence shard; standard O-proj preserves M. + if self.use_sequence_parallel and hidden_states.shape[0] == M: # Add SP padding if needed, and then perform reduce scatter. hidden_states = sp_reduce_scatter(hidden_states) @@ -1038,6 +1115,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vocab_size = config.vocab_size + # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP + # ranks to belong to one NVLink domain. + self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel) + if get_pp_group().is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, @@ -1058,6 +1139,7 @@ def get_layer(prefix: str): vllm_config, prefix, aux_stream=aux_stream, + run_gemm_rs=self.run_gemm_rs, ) self.start_layer, self.end_layer, self.layers = make_layers( @@ -1118,6 +1200,85 @@ def make_empty_intermediate_tensors( } ) + def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + super()._set_aux_hidden_state_layers(layers) + if self.use_attn_res: + # Emitted once, at configuration time. Which layers are tapped and + # which convention is in force are the two things you need to + # confirm from a running process, and neither is recoverable from + # the served output. + logger.info_once( + "Kimi-K3 aux hidden capture: layers=%s mode=%s " + "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)", + layers, + "attn_res_stream" if self._aux_attn_res_stream else "prefix_only", + int(self._aux_attn_res_stream), + ) + + @property + def _aux_attn_res_stream(self) -> bool: + return envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM + + def _capture_aux_hidden_stream( + self, + layer_idx: int, + prefix_sum: torch.Tensor, + pending_mlp_out: torch.Tensor | None, + block_residual: torch.Tensor, + ) -> torch.Tensor: + """Auxiliary feature tapped after ``layer_idx`` under AttnRes. + + The wire between layers only carries the current block's running prefix; + the committed blocks live in the bank. The value the next consumer + actually reads is the pre-norm AttnRes mixture over + ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were + trained against. ``attn_res`` with no delta, no block write and no + output norm computes exactly that and leaves both the prefix and the + bank untouched. + + Folding the pending MLP output into the prefix rather than passing it as + ``delta`` is deliberate: the kernel writes an applied delta back into + the prefix in place, which would double-add it into the live residual + stream. + """ + prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out + # `use_attn_res` is what constructs the norm and projection weights this + # reads; without it there is no mixture to compute and the attribute + # lookups below would raise. + if not (self._aux_attn_res_stream and self.use_attn_res): + return prefix + + if layer_idx + 1 < self.end_layer: + consumer = self.layers[layer_idx + 1] + score_norm = consumer.self_attention_res_norm + score_proj = consumer.self_attention_res_proj + num_blocks = consumer.prev_valid_blocks + elif get_pp_group().is_last_rank: + # Nothing downstream but the model's own output-side aggregation. + score_norm = self.output_attn_res_norm + score_proj = self.output_attn_res_proj + num_blocks = self.num_attn_res_blocks + else: + # Last layer of a non-final pipeline stage: the consumer lives on + # the next rank and the output-side aggregation only exists on the + # last one, so there is nothing here to mix against. Falling back + # to the running prefix keeps the tap defined rather than reaching + # for weights this rank does not construct. + return prefix + + return attn_res( + prefix, + None, + block_residual, + score_norm.weight, + score_proj.weight.squeeze(0), + None, + num_blocks=num_blocks, + block_write_idx=-1, + eps=score_norm.variance_epsilon, + output_norm_eps=0.0, + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -1185,7 +1346,10 @@ def forward( if (layer_idx + 1) in self.aux_hidden_state_layers: if self.use_attn_res: assert prefix_sum is not None - aux_hidden_state = prefix_sum + hidden_states + assert residual is not None + aux_hidden_state = self._capture_aux_hidden_stream( + layer_idx, prefix_sum, hidden_states, residual + ) else: assert residual is not None aux_hidden_state = hidden_states + residual @@ -1397,7 +1561,13 @@ def finalize_mega_moe_weights(self) -> None: class KimiLinearForCausalLM( - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid + nn.Module, + HasInnerState, + SupportsPP, + MixtureOfExperts, + IsHybrid, + SupportsEagle3, + SupportsReplaySSM, ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1451,15 +1621,20 @@ def forward( # type: ignore[override] def get_mamba_state_dtype_from_config( cls, vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.kda_state_dtype( + ) -> tuple[torch.dtype, ...]: + dtypes = MambaStateDtypeCalculator.kda_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype ) + if vllm_config.cache_config.use_kda_recoverssm: + dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record( + dtypes, vllm_config.model_config.dtype + ) + return dtypes @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: "VllmConfig" - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + ) -> MambaStateShapes: parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_config tp_size = parallel_config.tensor_parallel_size @@ -1468,13 +1643,22 @@ def get_mamba_state_shape_from_config( if vllm_config.speculative_config else 0 ) - return MambaStateShapeCalculator.kda_state_shape( + shapes = MambaStateShapeCalculator.kda_state_shape( tp_size, hf_config.linear_attn_config["num_heads"], hf_config.linear_attn_config["head_dim"], conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"], num_spec=num_spec, ) + if vllm_config.cache_config.use_kda_recoverssm: + return MambaStateShapeCalculator.append_kda_recoverssm_record( + shapes, + hf_config.linear_attn_config["num_heads"], + hf_config.linear_attn_config["head_dim"], + tp_world_size=tp_size, + spec_query_len=1 + num_spec, + ) + return shapes @classmethod def get_mamba_state_copy_func( @@ -1535,6 +1719,7 @@ class KimiK3ForConditionalGeneration( SupportsEagle3, HasInnerState, IsHybrid, + SupportsReplaySSM, ): """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text.""" diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs.py new file mode 100644 index 000000000000..a775cd7b1286 --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs.py @@ -0,0 +1,790 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""SM100 BF16 GEMM with a fused tensor-parallel reduce-scatter.""" + +# Based on CUTLASS's Blackwell distributed GEMM-RS example at dcf215a. +# See https://github.com/NVIDIA/cutlass/issues/3117 for memory semantics. + +from functools import cache + +import cutlass +import torch +import torch.distributed._symmetric_memory as symm_mem +from cuda.bindings.driver import CUstream +from cutlass import BFloat16, Int32, Int64, Uint16, cute, utils +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm, nvvm, vector +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import make_fake_stream, make_fake_tensor, make_ptr, nullptr +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils import get_smem_capacity_in_bytes + +from vllm.cute_utils import _tcgen05, mbarrier, simple_tma_copy, to_cta0_smem +from vllm.distributed import get_tp_group +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + + +@dsl_user_op +def nanosleep(ns: int, *, loc=None, ip=None) -> None: + nvvm.nanosleep(Int32(ns).ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + + +@dsl_user_op +def multimem_ld_reduce_16B(x: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + # NOTE: assume x is contiguous + assert x.element_type == BFloat16 + vec_type = ".v4.bf16x2" + + ptr = x.iterator.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip) + asm = ( + "multimem.ld_reduce.relaxed.gpu.global.add.acc::f32" + f"{vec_type} {{$0, $1, $2, $3}}, [$4];" + ) + struct = llvm.inline_asm( + llvm.StructType.get_literal([Int32.mlir_type] * 4), + [ptr], + asm, + "=r,=r,=r,=r,l", + has_side_effects=True, + loc=loc, + ip=ip, + ) + vec = vector.from_elements( + ir.VectorType.get([4], Int32.mlir_type, loc=loc), + [ + llvm.extractvalue(Int32.mlir_type, struct, [i], loc=loc, ip=ip) + for i in range(4) + ], + loc=loc, + ip=ip, + ) + ssa = cute.TensorSSA(vec, 4, Int32) + + y = cute.make_rmem_tensor(4, Int32) + y.store(ssa) + return cute.recast_tensor(y, x.element_type) + + +class Sm100GemmRsBF16: + def __init__( + self, rank: int, num_ranks: int, BN: int = 128, cta_group: int = 1 + ) -> None: + self.rank = rank + self.num_ranks = num_ranks + BM, BK = 128, 64 + self.cta_tile = (BM, BN, BK) + self.cta_group = cta_group + + smem_bytes = get_smem_capacity_in_bytes() + self.stage_size = (BM + (BN // cta_group)) * BK * 2 + self.num_stages = smem_bytes // self.stage_size + + @cute.jit + def prepare_tma( + self, tensor: cute.Tensor, BM: cutlass.Constexpr, BK: cutlass.Constexpr + ) -> cpasync.TmaInfo: + tma_group = ( + tcgen05.CtaGroup.TWO if self.cta_group == 2 else tcgen05.CtaGroup.ONE + ) + tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tma_group) + swizzle_128b = cute.make_swizzle(3, 4, 3) + layout = cute.make_layout( + (BM, BK, self.num_stages), + stride=(BK, 1, BM * BK), + ) + layout = cute.make_composed_layout(swizzle_128b, 0, layout) + return cpasync.make_tiled_tma_atom(tma_op, tensor, layout, (BM, BK)) + + @cute.jit + def __call__( + self, + A: cute.Tensor, + B: cute.Tensor, + partial_uc: cute.Tensor, + partial_mc_ptr: cute.Pointer, + output: cute.Tensor, + flags_uc: cute.Tensor, + flags_mc_ptr: cute.Pointer, + peer_flag_ptr: cute.Pointer, + grid_size: Int32, + stream: CUstream, + ) -> None: + N = B.shape[0] + BM, BN, BK = self.cta_tile + A_tma = self.prepare_tma(A, BM, BK) + B_tma = self.prepare_tma(B, BN // self.cta_group, BK) + padded_M = partial_uc.shape[0] + partial_mc = cute.make_tensor( + partial_mc_ptr, + cute.make_layout((padded_M, N), stride=(N, 1)), + ) + peer_flags = cute.make_tensor( + peer_flag_ptr, + cute.make_layout(self.num_ranks), + ) + + grid = (grid_size, 1, 1) + block = (10 * 32, 1, 1) + cluster = (self.cta_group, 1, 1) + self.kernel( + A_tma, + B_tma, + partial_uc, + partial_mc, + output, + flags_uc.iterator, + flags_mc_ptr, + peer_flags, + ).launch(grid=grid, block=block, cluster=cluster, stream=stream) + + @cute.kernel + def kernel( + self, + A_tma: cpasync.TmaInfo, + B_tma: cpasync.TmaInfo, + partial_uc: cute.Tensor, + partial_mc: cute.Tensor, + output: cute.Tensor, + flags_uc_ptr: cute.Pointer, + flags_mc_ptr: cute.Pointer, + peer_flags: cute.Tensor, + ) -> None: + tid, _, _ = cute.arch.thread_idx() + raw_bid, _, _ = cute.arch.block_idx() + num_bids, _, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(tid // 32) + + BM, BN, BK = self.cta_tile + cta_group = self.cta_group + num_stages = self.num_stages + num_ranks = self.num_ranks + + is_2cta = cta_group == 2 + cta_rank = raw_bid % self.cta_group + num_tmem_stages = 512 // BN + + smem = utils.SmemAllocator() + sA = smem.allocate_tensor( + BFloat16, + A_tma.smem_layout.outer, + byte_alignment=128, + swizzle=A_tma.smem_layout.inner, + ) + sB = smem.allocate_tensor( + BFloat16, + B_tma.smem_layout.outer, + byte_alignment=128, + swizzle=B_tma.smem_layout.inner, + ) + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + tmem_full_mbar = smem.allocate_array(Int64, num_tmem_stages) + tmem_empty_mbar = smem.allocate_array(Int64, num_tmem_stages) + taddr = smem.allocate(Int32, 4) + + # Named barriers + BAR_TMEM_ALLOC = 1 + BAR_EPILOGUE = 2 + BAR_COMM = 3 + + M, K = A_tma.tma_tensor.shape + N, _ = B_tma.tma_tensor.shape + grid_m = cute.ceil_div(M, BM) + # Keep 2-CTA clusters within a single N tile. + grid_m = cute.ceil_div(grid_m, cta_group) * cta_group + grid_n = cute.ceil_div(N, BN) + + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, cta_group) + cute.arch.mbarrier_init(tma_empty_mbar + i, 1) + for i in cutlass.range_constexpr(num_tmem_stages): + cute.arch.mbarrier_init(tmem_full_mbar + i, 1) + cute.arch.mbarrier_init(tmem_empty_mbar + i, 128 * cta_group) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(A_tma.atom) + cpasync.prefetch_descriptor(B_tma.atom) + + if cutlass.const_expr(is_2cta): + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + else: + cute.arch.sync_threads() + + total_tiles = grid_m * grid_n + + if warp_id == 9: + # TMA warp + tma_stage = 0 + parity = 1 + + if cutlass.const_expr(is_2cta): + tma_full_mbar_ = to_cta0_smem(tma_full_mbar) + else: + tma_full_mbar_ = tma_full_mbar + + # Select global-memory tiles. + # [(BM, BK), (M/BM, K/BK)] + gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK)) + gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cta_group, BK)) + + for bid in range(raw_bid, total_tiles, num_bids): + bid_m = bid % grid_m + bid_n = bid // grid_m + if cutlass.const_expr(cta_group == 2): + bid_n = bid_n * cta_group + cta_rank + + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1): + mbar = tma_full_mbar_ + tma_stage + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, parity) + + with cute.arch.elect_one(): + mbarrier.arrive_expect_tx(mbar, self.stage_size, "cluster") + simple_tma_copy( + A_tma.atom, + gA_tiles[None, (bid_m, iter_k)], + sA[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + B_tma.atom, + gB_tiles[None, (bid_n, iter_k)], + sB[None, None, tma_stage], + mbar, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + parity ^= 1 + + elif warp_id == 8: + # MMA warp + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=5 * 32) + + if cta_rank == 0: + tma_stage = 0 + tma_full_parity = 0 + tmem_stage = 0 + tmem_empty_parity = 1 + + MMA_M = BM * cta_group + MMA_N = BN + idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N) + sdesc = _tcgen05.make_sdesc_128B_swizzle(0) + multicast_mask = Uint16((1 << self.cta_group) - 1) + + for bid in range(raw_bid, total_tiles, num_bids): + cute.arch.mbarrier_wait( + tmem_empty_mbar + tmem_stage, tmem_empty_parity + ) + _tcgen05.fence_after_thread_sync() + + for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1): + d_tmem = BN * tmem_stage + a_addr = sA[None, None, tma_stage].iterator.toint() + b_addr = sB[None, None, tma_stage].iterator.toint() + a_desc = sdesc | (a_addr >> 4) + b_desc = sdesc | (b_addr >> 4) + + cute.arch.mbarrier_wait( + tma_full_mbar + tma_stage, tma_full_parity + ) + _tcgen05.fence_after_thread_sync() + + for mma_k in cutlass.range_constexpr(BK // 16): + enable_d = iter_k > 0 or mma_k > 0 + _tcgen05.mma_f16( + d_tmem, a_desc, b_desc, idesc, enable_d, cta_group + ) + a_desc += 32 >> 4 + b_desc += 32 >> 4 + _tcgen05.commit( + tma_empty_mbar + tma_stage, multicast_mask, cta_group + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_full_parity ^= 1 + + _tcgen05.commit( + tmem_full_mbar + tmem_stage, multicast_mask, cta_group + ) + + tmem_stage = (tmem_stage + 1) % num_tmem_stages + if tmem_stage == 0: + tmem_empty_parity ^= 1 + + elif warp_id >= 4: + # Communication warps + # Keep epilogue in warps 0-3 and communication in warps 4-7. + # Swapping the warpgroups can hang due to warp scheduling. + tid_ = tid % 128 + local_M = output.shape[0] + + # Offset in the [M, N] GEMM result. + rank_start = self.rank * local_M + rank_end = min(rank_start + local_M, M) + + st_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), BFloat16, num_bits_per_copy=128 + ) + + # Each thread issues one BF16x8 multimem reduction per vector. + vec_width = 8 + vec_cols = BN // vec_width + max_tile_rows = BM // num_ranks + vecs_per_tile = max_tile_rows * vec_cols + partial_vecs = cute.zipped_divide(partial_mc, (1, vec_width)) + output_vecs = cute.zipped_divide(output, (1, vec_width)) + + for tile_id in range(raw_bid, total_tiles, num_bids): + rs_bid_m = tile_id % grid_m + bid_n = tile_id // grid_m + + local_row_start = rs_bid_m * local_M // grid_m + local_row_end = (rs_bid_m + 1) * local_M // grid_m + global_row_start = rank_start + local_row_start + global_row_end = min(rank_start + local_row_end, M) + + if global_row_start < M: + # Map the current RS tile to the required GEMM tiles. + # Since an RS tile is smaller than a GEMM tile, each RS tile + # can overlap at most 2 GEMM tiles. + gemm_bid_m0 = global_row_start // BM + gemm_bid_m1 = max(global_row_end - 1, global_row_start) // BM + + # Poll local L2 with relaxed GPU-scope loads. The following + # multimem reduction reads every rank's L2 directly, so no + # acquire fence or L1 invalidation is needed. + if tid_ == 0: + # Poll the 1st GEMM tile. + flag_ptr = flags_uc_ptr + bid_n * grid_m + gemm_bid_m0 + arrivals = cute.arch.load( + flag_ptr, Int32, sem="relaxed", scope="gpu" + ) + while arrivals < num_ranks: + nanosleep(64) + arrivals = cute.arch.load( + flag_ptr, Int32, sem="relaxed", scope="gpu" + ) + elif tid_ == 32 and gemm_bid_m1 != gemm_bid_m0: + # Poll the 2nd GEMM tile with another warp if needed. + flag_ptr = flags_uc_ptr + bid_n * grid_m + gemm_bid_m1 + arrivals = cute.arch.load( + flag_ptr, Int32, sem="relaxed", scope="gpu" + ) + while arrivals < num_ranks: + nanosleep(64) + arrivals = cute.arch.load( + flag_ptr, Int32, sem="relaxed", scope="gpu" + ) + cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128) + + # Issue multimem.ld_reduce before storing any result. + reduced_vecs = [] + for vec_iter in cutlass.range_constexpr(vecs_per_tile // 128): + vec_idx = tid_ + vec_iter * 128 + + local_row = local_row_start + vec_idx // vec_cols + global_row = rank_start + local_row + col = bid_n * vec_cols + vec_idx % vec_cols + + reduced_vec = cute.make_rmem_tensor(vec_width, BFloat16) + if local_row < local_row_end and global_row < M: + tmp = multimem_ld_reduce_16B( + partial_vecs[None, (global_row, col)] + ) + reduced_vec.store(tmp.load()) + reduced_vecs.append(reduced_vec) + + # Store the result to local L2. + for vec_iter in cutlass.range_constexpr(vecs_per_tile // 128): + vec_idx = tid_ + vec_iter * 128 + + local_row = local_row_start + vec_idx // vec_cols + global_row = rank_start + local_row + col = bid_n * vec_cols + vec_idx % vec_cols + + if local_row < local_row_end and global_row < M: + cute.copy( + st_atom, + reduced_vecs[vec_iter], + output_vecs[None, (local_row, col)], + ) + cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128) + + # Release each GEMM tile consumed by this RS tile. The last + # consumer resets its producer flag for the next launch. + def release_gemm_tile( + gemm_bid_m, + tile_M, + logical_M, + rank_start, + rank_end, + local_M, + grid_m, + tile_N, + flags, + num_ranks, + ): + gemm_start = gemm_bid_m * tile_M + gemm_end = min(gemm_start + tile_M, logical_M) + local_start = max(gemm_start, rank_start) - rank_start + local_end = min(gemm_end, rank_end) - rank_start + + # Compute the number of consumers to identify the last + # arrival, which resets the producer flag. + first_rs_tile = ( + cute.ceil_div((local_start + 1) * grid_m, local_M) - 1 + ) + last_rs_tile = cute.ceil_div(local_end * grid_m, local_M) - 1 + num_consumers = last_rs_tile - first_rs_tile + 1 + + flag_ptr = flags + tile_N * grid_m + gemm_bid_m + old_count = cute.arch.atomic_add( + flag_ptr, Int32(1), sem="relaxed", scope="gpu" + ) + if old_count == num_ranks + num_consumers - 1: + cute.arch.store( + flag_ptr, Int32(0), sem="relaxed", scope="gpu" + ) + + if tid_ == 0: + # Arrive at the 1st GEMM tile. + release_gemm_tile( + gemm_bid_m0, + BM, + M, + rank_start, + rank_end, + local_M, + grid_m, + bid_n, + flags_uc_ptr, + num_ranks, + ) + elif tid_ == 32 and gemm_bid_m1 != gemm_bid_m0: + # Arrive at the 2nd GEMM tile if necessary. + release_gemm_tile( + gemm_bid_m1, + BM, + M, + rank_start, + rank_end, + local_M, + grid_m, + bid_n, + flags_uc_ptr, + num_ranks, + ) + + # Exit barrier. GPU scope is sufficient because the kernel writes + # local global memory and only needs to flush it to local L2. + cute.arch.barrier(barrier_id=BAR_COMM, number_of_threads=128) + if tid_ == 0: + exit_flag = total_tiles + raw_bid + utils.distributed.multimem_red_add1( + flags_mc_ptr + exit_flag, order="release", scope="gpu" + ) + utils.distributed.spin_lock_atom_cas_acquire_wait( + flags_uc_ptr + exit_flag, + expected_val=num_ranks, + reset_val=0, + scope="gpu", + ) + + else: + # Epilogue warps + warp_id_ = warp_id % 4 + tid_ = tid % 128 + + peer_flag_bases = cute.make_rmem_tensor(num_ranks, Int64) + if tid_ == 0: + for rank in cutlass.range_constexpr(num_ranks): + peer_flag_bases[rank] = cute.arch.load( + (peer_flags.iterator + rank).llvm_ptr, + Int64, + ) + + if warp_id_ == 0: + _tcgen05.alloc(taddr, cta_group) + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=5 * 32) + + WIDTH = cutlass.const_expr(16) + partial_vecs = cute.zipped_divide(partial_uc, (1, WIDTH)) + + bf16x16_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + BFloat16, + num_bits_per_copy=256, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + + tmem_stage = 0 + parity = 0 + + if cutlass.const_expr(is_2cta): + tmem_empty_mbar_ = to_cta0_smem(tmem_empty_mbar) + else: + tmem_empty_mbar_ = tmem_empty_mbar + + for bid in range(raw_bid, total_tiles, num_bids): + bid_m = bid % grid_m + bid_n = bid // grid_m + + if warp_id_ == 0: + cute.arch.mbarrier_wait(tmem_full_mbar + tmem_stage, parity) + cute.arch.barrier(barrier_id=BAR_EPILOGUE, number_of_threads=128) + _tcgen05.fence_after_thread_sync() + + for i in cutlass.range_constexpr(BN // WIDTH): + tcol = tmem_stage * BN + i * WIDTH + regs = _tcgen05.ld(warp_id_ * 32, tcol, "32x32b", WIDTH) + _tcgen05.wait_ld() + + if cutlass.const_expr(i == BN // WIDTH - 1): + _tcgen05.fence_before_thread_sync() + mbarrier.arrive(tmem_empty_mbar_ + tmem_stage, "cluster") + + tmp = cute.make_rmem_tensor(WIDTH, BFloat16) + tmp.store(regs.to(BFloat16)) + + global_row = bid_m * BM + tid_ + if global_row < M: + coord = (global_row, bid_n * (BN // WIDTH) + i) + cute.copy(bf16x16_atom, tmp, partial_vecs[None, coord]) + + cute.arch.barrier(barrier_id=BAR_EPILOGUE, number_of_threads=128) + + # Signal GEMM completion. GPU scope is sufficient because the + # partial output only needs to be flushed to local L2. + if tid_ == 0 and bid_m * BM < M: + gemm_start = bid_m * BM + gemm_end = min(gemm_start + BM, M) + first_owner = gemm_start // output.shape[0] + last_owner = (gemm_end - 1) // output.shape[0] + + # signal to all consuming ranks (can be more than 1) + for rank in cutlass.range_constexpr(num_ranks): + if first_owner <= rank and rank <= last_owner: + ptr = cute.make_ptr( + Int32, + peer_flag_bases[rank], + cute.AddressSpace.gmem, + assumed_align=16, + ) + ptr += bid_m + bid_n * grid_m + utils.distributed.red_add1( + ptr, order="release", scope="gpu" + ) + + tmem_stage = (tmem_stage + 1) % num_tmem_stages + if tmem_stage == 0: + parity ^= 1 + + if cutlass.const_expr(is_2cta): + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + else: + cute.arch.barrier(barrier_id=BAR_EPILOGUE, number_of_threads=128) + if warp_id_ == 0: + _tcgen05.dealloc(cta_group) + + @cache + @staticmethod + def compile(rank: int, num_ranks: int, BN: int, cta_group: int): + M = cute.sym_int() + padded_M = cute.sym_int() + N = cute.sym_int() + K = cute.sym_int() + local_M = cute.sym_int() + num_flags = cute.sym_int() + + A = make_fake_tensor( + BFloat16, (M, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16 + ) + B = make_fake_tensor( + BFloat16, (N, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16 + ) + partial = make_fake_tensor( + BFloat16, + (padded_M, N), + (cute.sym_int64(divisibility=16), 1), + assumed_align=32, + ) + partial_mc_ptr = nullptr(BFloat16, cute.AddressSpace.gmem, assumed_align=32) + output = make_fake_tensor( + BFloat16, + (local_M, N), + (cute.sym_int64(divisibility=16), 1), + assumed_align=32, + ) + flags = make_fake_tensor(Int32, (num_flags,), (1,), assumed_align=16) + flags_mc_ptr = nullptr(Int32, cute.AddressSpace.gmem, assumed_align=16) + peer_flag_ptr = nullptr(Int64, cute.AddressSpace.gmem, assumed_align=8) + + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = Sm100GemmRsBF16(rank, num_ranks, BN, cta_group) + return cute.compile( + kernel, + A, + B, + partial, + partial_mc_ptr, + output, + flags, + flags_mc_ptr, + peer_flag_ptr, + 128, + stream, + options="--enable-tvm-ffi", + ) + + +class GemmRS: + """Own the max-sized symmetric workspace for Kimi-K3 GEMM-RS launches. + + All TP ranks must belong to one NVLink domain for multimem instructions. + """ + + def __init__(self, *, max_M: int, N: int) -> None: + tp_group = get_tp_group() + group = tp_group.device_group + rank = tp_group.rank_in_group + world_size = tp_group.world_size + device = torch.device("cuda", torch.accelerator.current_device_index()) + + assert 1 < world_size <= 16 + assert 128 % world_size == 0 + assert max_M >= 128 and N % 128 == 0 + + max_M = (max_M + world_size - 1) // world_size * world_size + self.rank = rank + self.world_size = world_size + self.max_M = max_M + self.N = N + self.device = device + + self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device) + self.partial_handle = symm_mem.rendezvous(self.partial, group) + self.partial_mc_ptr = make_ptr( + BFloat16, + self.partial_handle.multicast_ptr, + cute.AddressSpace.gmem, + assumed_align=32, + ) + + grid_m = (max_M + 127) // 128 + cta_group = 2 if max_M >= 1024 or grid_m % 2 == 0 else 1 + grid_m = (grid_m + cta_group - 1) // cta_group * cta_group + self.num_sms = torch.cuda.get_device_properties(device).multi_processor_count + max_flags = grid_m * (N // 128) + self.num_sms + self.flags = symm_mem.empty(max_flags, dtype=torch.int32, device=device) + self.flags_handle = symm_mem.rendezvous(self.flags, group) + self.flags.zero_() + self.flags_mc_ptr = make_ptr( + Int32, + self.flags_handle.multicast_ptr, + cute.AddressSpace.gmem, + assumed_align=16, + ) + self.peer_flag_ptr = make_ptr( + Int64, + self.flags_handle.buffer_ptrs_dev, + cute.AddressSpace.gmem, + assumed_align=8, + ) + + assert self.partial_handle.multicast_ptr != 0 + assert self.flags_handle.multicast_ptr != 0 + torch.accelerator.synchronize(device) + tp_group.barrier() + + def can_run(self, linear: LinearBase) -> bool: + # Validate projection-invariant requirements once during model init. + # only supports BF16 for now + if not isinstance(linear.quant_method, UnquantizedLinearMethod): + return False + w = linear.weight + if w.ndim != 2: + return False + K = w.shape[1] + return ( + w.shape == (self.N, K) + and K % 64 == 0 + and w.dtype == torch.bfloat16 + and w.device == self.device + and w.is_contiguous() + ) + + def should_run(self, x: torch.Tensor) -> bool: + # Small-M shapes are supported but faster on the existing LL path. + return x.shape[0] >= 128 + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + assert x.ndim == 2 + M, K = x.shape + assert 0 < M <= self.max_M + assert w.shape == (self.N, K) and K % 64 == 0 + assert w.dtype == torch.bfloat16 + assert w.device == self.device + assert w.is_contiguous() + assert x.dtype == torch.bfloat16 + assert x.device == self.device + assert x.is_contiguous() + N = w.shape[0] + padded_M = (M + self.world_size - 1) // self.world_size + padded_M *= self.world_size + local_M = padded_M // self.world_size + + grid_m = (M + 127) // 128 + # Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024. + cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1 + grid_m = (grid_m + cta_group - 1) // cta_group * cta_group + BN = 256 if M * K >= 24 * 1024 * 1024 else 128 + assert N % BN == 0 + + num_tiles = grid_m * (N // BN) + num_ctas = min(num_tiles, self.num_sms) + num_ctas = num_ctas // cta_group * cta_group + assert self.flags.numel() >= num_tiles + num_ctas + + output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device) + compiled = Sm100GemmRsBF16.compile( + self.rank, + self.world_size, + BN, + cta_group, + ) + compiled( + x, + w, + self.partial[:padded_M], + self.partial_mc_ptr, + output, + self.flags, + self.flags_mc_ptr, + self.peer_flag_ptr, + num_ctas, + ) + return output + + +_gemm_rs: GemmRS | None = None + + +def init_gemm_rs(max_M: int, N: int) -> None: + """Collectively initialize the process-wide GEMM-RS state.""" + global _gemm_rs + if _gemm_rs is not None: + assert _gemm_rs.max_M >= max_M and _gemm_rs.N == N + return + _gemm_rs = GemmRS(max_M=max_M, N=N) + + +def get_gemm_rs() -> GemmRS: + assert _gemm_rs is not None, "GEMM-RS is not initialized" + return _gemm_rs diff --git a/vllm/models/kimi_k3/nvidia/ops/recoverssm.py b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py new file mode 100644 index 000000000000..c603ec7ed1fc --- /dev/null +++ b/vllm/models/kimi_k3/nvidia/ops/recoverssm.py @@ -0,0 +1,1067 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 RecoverSSM speculative verify and accepted-state recovery.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +@triton.jit +def _kda_gate( + raw_g, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND: tl.constexpr, +): + gate_input = raw_g + dt_bias + if USE_LOWER_BOUND: + return lower_bound * tl.sigmoid(A * gate_input) + softplus_gate = tl.where( + gate_input > 20.0, + gate_input, + tl.log(1.0 + tl.exp(gate_input)), + ) + return -A * softplus_gate + + +@triton.jit +def _kda_recurrent_step( + state, + k, + v, + raw_g, + raw_beta, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND: tl.constexpr, +): + normalized_k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) + gate = _kda_gate( + raw_g, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND, + ) + + state *= tl.exp(gate)[None, :] + correction = v - tl.sum(state * normalized_k[None, :], axis=1) + correction *= tl.sigmoid(raw_beta) + return state + correction[:, None] * normalized_k[None, :], correction + + +@triton.jit +def _kda_recoverssm_verify_kernel( + q_ptr, + k_ptr, + v_ptr, + raw_g_ptr, + raw_beta_ptr, + A_log_ptr, + dt_bias_ptr, + state_ptr, + correction_cache_ptr, + kg_cache_ptr, + out_ptr, + query_start_loc_ptr, + state_indices_ptr, + lower_bound, + null_block_id, + stride_q_token, + stride_k_token, + stride_v_token, + stride_g_token, + stride_beta_token, + stride_state_block, + stride_state_head, + stride_state_v, + stride_state_k, + stride_correction_block, + stride_correction_head, + stride_correction_pos, + stride_correction_dim, + stride_kg_block, + stride_kg_head, + stride_kg_pos, + stride_kg_dim, + stride_out_token, + stride_query_start_loc, + stride_state_indices, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + SPEC_QUERY_LEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + + bos = tl.load(query_start_loc_ptr + pid_b * stride_query_start_loc).to(tl.int64) + eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_query_start_loc).to( + tl.int64 + ) + query_len = eos - bos + state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to(tl.int64) + + offs_k = tl.arange(0, BK) + offs_v = pid_v * BV + tl.arange(0, BV) + mask_k = offs_k < K + mask_v = offs_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + + if state_idx <= null_block_id: + for token_offset in tl.static_range(SPEC_QUERY_LEN): + token_valid = token_offset < query_len + tl.store( + out_ptr + (bos + token_offset) * stride_out_token + pid_h * V + offs_v, + tl.zeros([BV], dtype=tl.float32), + mask=token_valid & mask_v, + ) + return + + state_ptrs = ( + state_ptr + + state_idx * stride_state_block + + pid_h * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + state = tl.load(state_ptrs, mask=mask_state, other=0.0).to(tl.float32) + A = tl.exp(tl.load(A_log_ptr + pid_h).to(tl.float32)) + + for token_offset in tl.static_range(SPEC_QUERY_LEN): + token_valid = token_offset < query_len + token = bos + token_offset + q = tl.load( + q_ptr + token * stride_q_token + pid_h * K + offs_k, + mask=token_valid & mask_k, + other=0.0, + ).to(tl.float32) + k = tl.load( + k_ptr + token * stride_k_token + pid_h * K + offs_k, + mask=token_valid & mask_k, + other=0.0, + ).to(tl.float32) + v = tl.load( + v_ptr + token * stride_v_token + pid_h * V + offs_v, + mask=token_valid & mask_v, + other=0.0, + ).to(tl.float32) + raw_g = tl.load( + raw_g_ptr + token * stride_g_token + pid_h * K + offs_k, + mask=token_valid & mask_k, + other=0.0, + ).to(tl.float32) + raw_beta = tl.load( + raw_beta_ptr + token * stride_beta_token + pid_h, + mask=token_valid, + other=0.0, + ).to(tl.float32) + + q *= tl.rsqrt(tl.sum(q * q) + 1e-6) * (K**-0.5) + dt_bias = tl.load(dt_bias_ptr + pid_h * K + offs_k, mask=mask_k, other=0.0).to( + tl.float32 + ) + updated_state, correction = _kda_recurrent_step( + state, + k, + v, + raw_g, + raw_beta, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND, + ) + state = tl.where(token_valid, updated_state, state) + + out = tl.sum(state * q[None, :], axis=1) + tl.store( + out_ptr + token * stride_out_token + pid_h * V + offs_v, + out, + mask=token_valid & mask_v, + ) + + correction_ptr = ( + correction_cache_ptr + + state_idx * stride_correction_block + + pid_h * stride_correction_head + + token_offset * stride_correction_pos + ) + tl.store( + correction_ptr + offs_v * stride_correction_dim, + correction, + mask=token_valid & mask_v, + ) + if pid_v == 0: + kg_ptr = ( + kg_cache_ptr + + state_idx * stride_kg_block + + pid_h * stride_kg_head + + token_offset * stride_kg_pos + ) + tl.store( + kg_ptr + offs_k * stride_kg_dim, + k, + mask=token_valid & mask_k, + ) + tl.store( + kg_ptr + (K + offs_k) * stride_kg_dim, + raw_g, + mask=token_valid & mask_k, + ) + + +@triton.heuristics( + { + "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, + "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, + } +) +@triton.jit +def _prepare_commit_plan_kernel( + num_accepted_ptr, + request_indices_ptr, + state_indices_ptr, + query_start_loc_ptr, + block_table_ptr, + num_computed_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + mamba_block_size, + block_table_width, + stride_num_accepted, + stride_request_indices, + stride_state_indices, + stride_query_start_loc, + stride_block_table_row, + stride_block_table_col, + stride_num_computed, + SPEC_QUERY_LEN: tl.constexpr, + HAS_REQUEST_INDICES: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + spec_idx = tl.program_id(0) + source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( + tl.int64 + ) + request_idx = spec_idx + if HAS_REQUEST_INDICES: + request_idx = tl.load( + request_indices_ptr + spec_idx * stride_request_indices + ).to(tl.int64) + num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( + tl.int32 + ) + bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) + eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( + tl.int64 + ) + query_len = (eos - bos).to(tl.int32) + commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) + commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) + + final_state_idx = source_state_idx + boundary_state_idx = null_block_id + boundary_recovery_len = 0 + if ALIGN_MODE: + num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( + tl.int32 + ) + final_num_computed = num_computed + commit_len + final_state_col = tl.minimum( + final_num_computed // mamba_block_size, block_table_width - 1 + ) + final_state_idx = tl.load( + block_table_ptr + + request_idx * stride_block_table_row + + final_state_col * stride_block_table_col + ).to(tl.int64) + next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size + crosses_boundary = final_num_computed >= next_boundary + boundary_recovery_len = next_boundary - num_computed + boundary_state_idx = tl.load( + block_table_ptr + + request_idx * stride_block_table_row + + (next_boundary // mamba_block_size - 1) * stride_block_table_col, + mask=crosses_boundary, + other=null_block_id, + ).to(tl.int64) + valid = (source_state_idx > null_block_id) & (commit_len > 0) + tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) + tl.store( + final_state_indices_ptr + spec_idx, + tl.where(valid, final_state_idx, null_block_id), + ) + tl.store( + boundary_state_indices_ptr + spec_idx, + tl.where(valid, boundary_state_idx, null_block_id), + ) + tl.store( + boundary_recovery_lens_ptr + spec_idx, + tl.where(valid, boundary_recovery_len, 0), + ) + + +@triton.jit +def _compact_conv_state_kernel( + conv_state_ref_ptr, + conv_state_base_addrs_ptr, + conv_state_block_strides_ptr, + conv_state_dim_strides_ptr, + conv_state_token_strides_ptr, + state_indices_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + null_block_id, + conv_dim, + conv_history_len, + stride_state_indices, + BLOCK_D: tl.constexpr, + BLOCK_HISTORY: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + pid_d = tl.program_id(0) + pid_b = tl.program_id(1) + pid_l = tl.program_id(2) + source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( + tl.int64 + ) + if source_state_idx <= null_block_id: + return + + commit_len = tl.load(commit_lens_ptr + pid_b) + if commit_len == 0: + return + final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) + boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) + + if final_state_idx <= null_block_id: + return + + base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) + block_stride = tl.load(conv_state_block_strides_ptr + pid_l) + dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) + token_stride = tl.load(conv_state_token_strides_ptr + pid_l) + conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) + source_ptr = conv_state_ptr + source_state_idx * block_stride + final_ptr = conv_state_ptr + final_state_idx * block_stride + + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + offs_h = tl.arange(0, BLOCK_HISTORY) + mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) + final_values = tl.load( + source_ptr + + offs_d[:, None] * dim_stride + + (commit_len - 1 + offs_h[None, :]) * token_stride, + mask=mask, + ) + if ALIGN_MODE: + boundary_values = tl.load( + source_ptr + + offs_d[:, None] * dim_stride + + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, + mask=mask & (boundary_state_idx > null_block_id), + ) + boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride + tl.store( + boundary_ptr + + offs_d[:, None] * dim_stride + + offs_h[None, :] * token_stride, + boundary_values, + mask=mask & (boundary_state_idx > null_block_id), + ) + tl.store( + final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, + final_values, + mask=mask, + ) + + +@triton.jit +def _commit_kda_state_kernel( + state_ref_ptr, + state_base_addrs_ptr, + state_block_strides_ptr, + correction_cache_ref_ptr, + correction_cache_base_addrs_ptr, + correction_cache_block_strides_ptr, + kg_cache_ref_ptr, + kg_cache_base_addrs_ptr, + kg_cache_block_strides_ptr, + A_log_ptr, + dt_bias_ptr, + state_indices_ptr, + commit_lens_ptr, + final_state_indices_ptr, + boundary_state_indices_ptr, + boundary_recovery_lens_ptr, + lower_bound, + null_block_id, + stride_state_head, + stride_state_v, + stride_state_k, + stride_correction_cache_head, + stride_correction_cache_pos, + stride_correction_cache_dim, + stride_kg_cache_head, + stride_kg_cache_pos, + stride_kg_cache_dim, + stride_A_layer, + stride_A_head, + stride_dt_bias_layer, + stride_dt_bias_head, + stride_dt_bias_dim, + stride_state_indices, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + NUM_HEADS: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, + ALIGN_MODE: tl.constexpr, +): + pid_v = tl.program_id(0) + pid_b = tl.program_id(1) + pid_lh = tl.program_id(2) + pid_l = pid_lh // NUM_HEADS + pid_h = pid_lh % NUM_HEADS + + source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( + tl.int64 + ) + if source_state_idx <= null_block_id: + return + commit_len = tl.load(commit_lens_ptr + pid_b) + if commit_len == 0: + return + final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) + boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) + + if final_state_idx <= null_block_id: + return + + state_base_addr = tl.load(state_base_addrs_ptr + pid_l) + state_block_stride = tl.load(state_block_strides_ptr + pid_l) + state_ptr = state_base_addr.to(tl.pointer_type(state_ref_ptr.dtype.element_ty)) + source_state_ptr = ( + state_ptr + source_state_idx * state_block_stride + pid_h * stride_state_head + ) + + correction_cache_base_addr = tl.load(correction_cache_base_addrs_ptr + pid_l) + correction_cache_block_stride = tl.load(correction_cache_block_strides_ptr + pid_l) + correction_cache_ptr = correction_cache_base_addr.to( + tl.pointer_type(correction_cache_ref_ptr.dtype.element_ty) + ) + correction_cache_ptr += ( + source_state_idx * correction_cache_block_stride + + pid_h * stride_correction_cache_head + ) + kg_cache_base_addr = tl.load(kg_cache_base_addrs_ptr + pid_l) + kg_cache_block_stride = tl.load(kg_cache_block_strides_ptr + pid_l) + kg_cache_ptr = kg_cache_base_addr.to( + tl.pointer_type(kg_cache_ref_ptr.dtype.element_ty) + ) + kg_cache_ptr += ( + source_state_idx * kg_cache_block_stride + pid_h * stride_kg_cache_head + ) + + offs_k = tl.arange(0, BK) + offs_v = pid_v * BV + tl.arange(0, BV) + mask_k = offs_k < K + mask_v = offs_v < V + mask_state = mask_v[:, None] & mask_k[None, :] + state_ptrs = ( + source_state_ptr + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + initial_state = tl.load(state_ptrs, mask=mask_state, other=0.0).to(tl.float32) + A = tl.exp( + tl.load(A_log_ptr + pid_l * stride_A_layer + pid_h * stride_A_head).to( + tl.float32 + ) + ) + + dt_bias = tl.load( + dt_bias_ptr + + pid_l * stride_dt_bias_layer + + pid_h * stride_dt_bias_head + + offs_k * stride_dt_bias_dim, + mask=mask_k, + other=0.0, + ).to(tl.float32) + final_decay = tl.full([BK], 1.0, tl.float32) + final_correction = tl.zeros([BV, BK], tl.float32) + boundary_decay = tl.full([BK], 1.0, tl.float32) + boundary_correction = tl.zeros([BV, BK], tl.float32) + + for reverse_offset in range(commit_len): + token_offset = commit_len - reverse_offset - 1 + correction_ptr = ( + correction_cache_ptr + token_offset * stride_correction_cache_pos + ) + kg_ptr = kg_cache_ptr + token_offset * stride_kg_cache_pos + k = tl.load( + kg_ptr + offs_k * stride_kg_cache_dim, + mask=mask_k, + other=0.0, + ).to(tl.float32) + correction = tl.load( + correction_ptr + offs_v * stride_correction_cache_dim, + mask=mask_v, + other=0.0, + ).to(tl.float32) + raw_g = tl.load( + kg_ptr + (K + offs_k) * stride_kg_cache_dim, + mask=mask_k, + other=0.0, + ).to(tl.float32) + normalized_k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) + gate = _kda_gate( + raw_g, + dt_bias, + A, + lower_bound, + USE_LOWER_BOUND, + ) + update = correction[:, None] * normalized_k[None, :] + decay = tl.exp(gate) + final_correction += update * final_decay[None, :] + final_decay *= decay + if ALIGN_MODE: + before_boundary = token_offset < boundary_recovery_len + boundary_correction += tl.where( + before_boundary, + update * boundary_decay[None, :], + 0.0, + ) + boundary_decay *= tl.where(before_boundary, decay, 1.0) + + state = initial_state * final_decay[None, :] + final_correction + if ALIGN_MODE: + boundary_ptrs = ( + state_ptr + + boundary_state_idx * state_block_stride + + pid_h * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + boundary_state = initial_state * boundary_decay[None, :] + boundary_correction + tl.store( + boundary_ptrs, + boundary_state, + mask=mask_state & (boundary_state_idx > null_block_id), + ) + + final_ptrs = ( + state_ptr + + final_state_idx * state_block_stride + + pid_h * stride_state_head + + offs_v[:, None] * stride_state_v + + offs_k[None, :] * stride_state_k + ) + tl.store(final_ptrs, state, mask=mask_state) + + +def kda_recoverssm_verify( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + raw_g: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float | None, + checkpoint_state: torch.Tensor, + correction_cache: torch.Tensor, + kg_cache: torch.Tensor, + query_start_loc: torch.Tensor, + state_indices: torch.Tensor, + spec_query_len: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Verify a KDA speculative window without modifying its checkpoint.""" + if q.ndim != 4 or q.shape[0] != 1: + raise ValueError("KDA RecoverSSM q must have shape [1, tokens, heads, dim]") + _, total_tokens, num_heads, key_dim = q.shape + value_dim = v.shape[-1] + if k.shape != q.shape or v.shape != (1, total_tokens, num_heads, value_dim): + raise ValueError("KDA RecoverSSM q, k, and v shapes are incompatible") + if raw_g.shape != q.shape or raw_beta.shape != (1, total_tokens, num_heads): + raise ValueError("KDA RecoverSSM gate or beta shape is incompatible") + if any(tensor.stride()[2:] != (key_dim, 1) for tensor in (q, k, raw_g)): + raise ValueError("KDA RecoverSSM q, k, and gate heads must be contiguous") + if v.stride()[2:] != (value_dim, 1) or raw_beta.stride(2) != 1: + raise ValueError("KDA RecoverSSM v and beta heads must be contiguous") + num_blocks = checkpoint_state.shape[0] + if checkpoint_state.shape[1:] != ( + num_heads, + value_dim, + key_dim, + ): + raise ValueError("KDA RecoverSSM checkpoint shape is incompatible") + expected_correction_shape = ( + num_blocks, + num_heads, + spec_query_len, + value_dim, + ) + if correction_cache.shape != expected_correction_shape: + raise ValueError( + f"KDA RecoverSSM correction buffer needs shape {expected_correction_shape}" + ) + expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) + if kg_cache.shape != expected_kg_shape: + raise ValueError( + f"KDA RecoverSSM key/gate buffer needs shape {expected_kg_shape}" + ) + if correction_cache.dtype != torch.float32: + raise ValueError("KDA RecoverSSM correction buffer must use float32") + if kg_cache.dtype != k.dtype: + raise ValueError("KDA RecoverSSM key/gate buffer must match activation dtype") + if A_log.shape != (num_heads,) or dt_bias.numel() != num_heads * key_dim: + raise ValueError("KDA RecoverSSM gate parameters are incompatible") + if not A_log.is_contiguous() or not dt_bias.is_contiguous(): + raise ValueError("KDA RecoverSSM gate parameters must be contiguous") + batch = state_indices.shape[0] + if query_start_loc.shape[0] != batch + 1: + raise ValueError("KDA RecoverSSM query metadata is incompatible") + if total_tokens > batch * spec_query_len: + raise ValueError( + "KDA RecoverSSM speculative decode input exceeds its activation capacity" + ) + if out is None: + out = torch.empty_like(v) + if out.shape != v.shape: + raise ValueError("KDA RecoverSSM output shape is incompatible") + if out.stride()[2:] != (value_dim, 1): + raise ValueError("KDA RecoverSSM output heads must be contiguous") + device = q.device + if any( + tensor.device != device + for tensor in ( + k, + v, + raw_g, + raw_beta, + A_log, + dt_bias, + checkpoint_state, + correction_cache, + kg_cache, + query_start_loc, + state_indices, + out, + ) + ): + raise ValueError("KDA RecoverSSM inputs must be on the same device") + if total_tokens == 0: + return out + + block_k = triton.next_power_of_2(key_dim) + block_v = min(triton.next_power_of_2(value_dim), 32) + grid = (triton.cdiv(value_dim, block_v), batch, num_heads) + _kda_recoverssm_verify_kernel[grid]( + q, + k, + v, + raw_g, + raw_beta, + A_log, + dt_bias, + checkpoint_state, + correction_cache, + kg_cache, + out, + query_start_loc, + state_indices, + lower_bound or 0.0, + NULL_BLOCK_ID, + q.stride(1), + k.stride(1), + v.stride(1), + raw_g.stride(1), + raw_beta.stride(1), + checkpoint_state.stride(0), + checkpoint_state.stride(1), + checkpoint_state.stride(2), + checkpoint_state.stride(3), + correction_cache.stride(0), + correction_cache.stride(1), + correction_cache.stride(2), + correction_cache.stride(3), + kg_cache.stride(0), + kg_cache.stride(1), + kg_cache.stride(2), + kg_cache.stride(3), + out.stride(1), + query_start_loc.stride(0), + state_indices.stride(0), + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + SPEC_QUERY_LEN=spec_query_len, + USE_LOWER_BOUND=lower_bound is not None, + num_warps=4, + num_stages=2, + ) + return out + + +@dataclass +class KDARecoverSSMCommitContext: + conv_states: tuple[torch.Tensor, ...] + conv_state_base_addrs: torch.Tensor + conv_state_block_strides: torch.Tensor + conv_state_dim_strides: torch.Tensor + conv_state_token_strides: torch.Tensor + conv_history_len: int + checkpoints: tuple[torch.Tensor, ...] + state_base_addrs: torch.Tensor + state_block_strides: torch.Tensor + correction_caches: tuple[torch.Tensor, ...] + correction_cache_base_addrs: torch.Tensor + correction_cache_block_strides: torch.Tensor + kg_caches: tuple[torch.Tensor, ...] + kg_cache_base_addrs: torch.Tensor + kg_cache_block_strides: torch.Tensor + commit_lens: torch.Tensor + final_state_indices: torch.Tensor + boundary_state_indices: torch.Tensor + boundary_recovery_lens: torch.Tensor + A_log: torch.Tensor + dt_bias: torch.Tensor + lower_bound: float | None + spec_query_len: int + + @classmethod + def create( + cls, + layers: Sequence[Any], + *, + spec_query_len: int, + max_num_reqs: int, + ) -> "KDARecoverSSMCommitContext": + if not layers: + raise ValueError("KDA RecoverSSM commit requires at least one layer") + if any(len(layer.kv_cache) != 4 for layer in layers): + raise ValueError( + "KDA RecoverSSM pages must contain conv, state, correction, " + "and key/gate" + ) + + conv_states = [layer.kv_cache[0] for layer in layers] + if not is_conv_state_dim_first(): + conv_states = [state.transpose(-1, -2) for state in conv_states] + checkpoints = [layer.kv_cache[1] for layer in layers] + correction_caches = [layer.kv_cache[2] for layer in layers] + kg_caches = [layer.kv_cache[3] for layer in layers] + A_log = [layer.A_log for layer in layers] + dt_bias = [ + layer.dt_bias.view(layer.local_num_heads, layer.head_dim) + for layer in layers + ] + lower_bounds = {layer.gate_lower_bound for layer in layers} + if len(lower_bounds) != 1: + raise ValueError("KDA RecoverSSM layers need matching gate bounds") + + state_ref = checkpoints[0] + if state_ref.ndim != 4: + raise ValueError("KDA RecoverSSM checkpoint must be four-dimensional") + num_blocks, num_heads, value_dim, key_dim = state_ref.shape + for state in checkpoints: + if ( + state.shape != state_ref.shape + or state.dtype != state_ref.dtype + or state.device != state_ref.device + or state.stride()[1:] != state_ref.stride()[1:] + ): + raise ValueError( + "KDA RecoverSSM layers need matching checkpoint layout" + ) + expected_correction_shape = ( + num_blocks, + num_heads, + spec_query_len, + value_dim, + ) + correction_ref = correction_caches[0] + for correction_cache in correction_caches: + if ( + correction_cache.shape != expected_correction_shape + or correction_cache.dtype != torch.float32 + or correction_cache.device != state_ref.device + or correction_cache.stride()[1:] != correction_ref.stride()[1:] + ): + raise ValueError( + "KDA RecoverSSM correction buffers need float32 shape " + f"{expected_correction_shape}" + ) + expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) + kg_ref = kg_caches[0] + for kg_cache in kg_caches: + if ( + kg_cache.shape != expected_kg_shape + or kg_cache.dtype != kg_ref.dtype + or kg_cache.device != state_ref.device + or kg_cache.stride()[1:] != kg_ref.stride()[1:] + ): + raise ValueError( + f"KDA RecoverSSM key/gate buffers need shape {expected_kg_shape}" + ) + if any(param.shape != (num_heads,) for param in A_log): + raise ValueError("KDA RecoverSSM A_log shape is incompatible") + if any(param.shape != (num_heads, key_dim) for param in dt_bias): + raise ValueError("KDA RecoverSSM dt_bias shape is incompatible") + + conv_ref = conv_states[0] + if conv_ref.ndim != 3: + raise ValueError("KDA RecoverSSM conv state must be three-dimensional") + conv_dim, conv_state_len = conv_ref.shape[1:] + conv_history_len = conv_state_len - spec_query_len + 1 + if conv_history_len <= 0: + raise ValueError("KDA RecoverSSM conv state is shorter than its window") + for conv_state in conv_states: + if ( + conv_state.shape != conv_ref.shape + or conv_state.dtype != conv_ref.dtype + or conv_state.device != state_ref.device + or conv_state.shape[0] != num_blocks + ): + raise ValueError("KDA RecoverSSM layers need matching conv state") + + device = state_ref.device + + def _base_addrs(tensors: Sequence[torch.Tensor]) -> torch.Tensor: + return torch.tensor( + [tensor.data_ptr() for tensor in tensors], + dtype=torch.int64, + device=device, + ) + + def _block_strides(tensors: Sequence[torch.Tensor]) -> torch.Tensor: + return torch.tensor( + [tensor.stride(0) for tensor in tensors], + dtype=torch.int64, + device=device, + ) + + return cls( + conv_states=tuple(conv_states), + conv_state_base_addrs=_base_addrs(conv_states), + conv_state_block_strides=_block_strides(conv_states), + conv_state_dim_strides=torch.tensor( + [state.stride(1) for state in conv_states], + dtype=torch.int64, + device=device, + ), + conv_state_token_strides=torch.tensor( + [state.stride(2) for state in conv_states], + dtype=torch.int64, + device=device, + ), + conv_history_len=conv_history_len, + checkpoints=tuple(checkpoints), + state_base_addrs=_base_addrs(checkpoints), + state_block_strides=_block_strides(checkpoints), + correction_caches=tuple(correction_caches), + correction_cache_base_addrs=_base_addrs(correction_caches), + correction_cache_block_strides=_block_strides(correction_caches), + kg_caches=tuple(kg_caches), + kg_cache_base_addrs=_base_addrs(kg_caches), + kg_cache_block_strides=_block_strides(kg_caches), + commit_lens=torch.empty(max_num_reqs, dtype=torch.int32, device=device), + final_state_indices=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + boundary_state_indices=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + boundary_recovery_lens=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + A_log=torch.stack(tuple(A_log)).contiguous(), + dt_bias=torch.stack(tuple(dt_bias)).contiguous(), + lower_bound=lower_bounds.pop(), + spec_query_len=spec_query_len, + ) + + def commit( + self, + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + query_start_loc: torch.Tensor, + request_indices: torch.Tensor | None = None, + block_table: torch.Tensor | None = None, + num_computed_tokens: torch.Tensor | None = None, + mamba_block_size: int | None = None, + ) -> None: + """Fold accepted KDA and convolution inputs into every layer.""" + batch = state_indices.shape[0] + if batch == 0: + return + if batch > self.commit_lens.shape[0]: + raise ValueError("KDA RecoverSSM commit batch exceeds its plan capacity") + if query_start_loc.shape[0] != batch + 1: + raise ValueError("KDA RecoverSSM commit metadata is incompatible") + if request_indices is not None and request_indices.shape[0] < batch: + raise ValueError("KDA RecoverSSM request mapping is too short") + align_args = (block_table, num_computed_tokens, mamba_block_size) + if any(arg is not None for arg in align_args) and any( + arg is None for arg in align_args + ): + raise ValueError("KDA RecoverSSM align metadata is incomplete") + if mamba_block_size is not None and mamba_block_size < self.spec_query_len: + raise ValueError( + "KDA RecoverSSM align block size must cover one speculative window" + ) + if block_table is not None and block_table.ndim != 2: + raise ValueError("KDA RecoverSSM block table must be two-dimensional") + device = self.checkpoints[0].device + if ( + any( + tensor.device != device + for tensor in ( + num_accepted_tokens, + state_indices, + query_start_loc, + ) + ) + or (request_indices is not None and request_indices.device != device) + or (block_table is not None and block_table.device != device) + or ( + num_computed_tokens is not None and num_computed_tokens.device != device + ) + ): + raise ValueError("KDA RecoverSSM commit inputs must be on the same device") + + block_table_stride = (0, 0) if block_table is None else block_table.stride() + num_computed_stride = ( + 0 if num_computed_tokens is None else num_computed_tokens.stride(0) + ) + + num_layers = len(self.checkpoints) + conv_ref = self.conv_states[0] + conv_dim = conv_ref.shape[1] + block_history = triton.next_power_of_2(self.conv_history_len) + _prepare_commit_plan_kernel[(batch,)]( + num_accepted_tokens, + request_indices, + state_indices, + query_start_loc, + block_table, + num_computed_tokens, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + mamba_block_size or 1, + block_table.shape[1] if block_table is not None else 1, + num_accepted_tokens.stride(0), + request_indices.stride(0) if request_indices is not None else 0, + state_indices.stride(0), + query_start_loc.stride(0), + block_table_stride[0], + block_table_stride[1], + num_computed_stride, + SPEC_QUERY_LEN=self.spec_query_len, + num_warps=1, + ) + _compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)]( + conv_ref, + self.conv_state_base_addrs, + self.conv_state_block_strides, + self.conv_state_dim_strides, + self.conv_state_token_strides, + state_indices, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + NULL_BLOCK_ID, + conv_dim, + self.conv_history_len, + state_indices.stride(0), + BLOCK_D=256, + BLOCK_HISTORY=block_history, + ALIGN_MODE=block_table is not None, + num_warps=4, + ) + + state_ref = self.checkpoints[0] + _, num_heads, value_dim, key_dim = state_ref.shape + block_k = triton.next_power_of_2(key_dim) + block_v = min(triton.next_power_of_2(value_dim), 32) + grid = ( + triton.cdiv(value_dim, block_v), + batch, + num_layers * num_heads, + ) + _commit_kda_state_kernel[grid]( + state_ref, + self.state_base_addrs, + self.state_block_strides, + self.correction_caches[0], + self.correction_cache_base_addrs, + self.correction_cache_block_strides, + self.kg_caches[0], + self.kg_cache_base_addrs, + self.kg_cache_block_strides, + self.A_log, + self.dt_bias, + state_indices, + self.commit_lens, + self.final_state_indices, + self.boundary_state_indices, + self.boundary_recovery_lens, + self.lower_bound or 0.0, + NULL_BLOCK_ID, + state_ref.stride(1), + state_ref.stride(2), + state_ref.stride(3), + self.correction_caches[0].stride(1), + self.correction_caches[0].stride(2), + self.correction_caches[0].stride(3), + self.kg_caches[0].stride(1), + self.kg_caches[0].stride(2), + self.kg_caches[0].stride(3), + self.A_log.stride(0), + self.A_log.stride(1), + self.dt_bias.stride(0), + self.dt_bias.stride(1), + self.dt_bias.stride(2), + state_indices.stride(0), + K=key_dim, + V=value_dim, + BK=block_k, + BV=block_v, + NUM_HEADS=num_heads, + USE_LOWER_BOUND=self.lower_bound is not None, + ALIGN_MODE=block_table is not None, + num_warps=4, + num_stages=2, + ) + + +__all__ = ["KDARecoverSSMCommitContext", "kda_recoverssm_verify"] diff --git a/vllm/models/minimax_m3/nvidia/model.py b/vllm/models/minimax_m3/nvidia/model.py index 68f9805b1fdc..a75eca2b630e 100644 --- a/vllm/models/minimax_m3/nvidia/model.py +++ b/vllm/models/minimax_m3/nvidia/model.py @@ -498,7 +498,9 @@ def __init__( set_default_quant_scales(self, register_buffer=True) # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main # cache (--attention-config '{"indexer_kv_dtype": ...}'). - self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype + self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype( + "bf16" + ) # Shared top-k buffer: the indexer writes the selected blocks into it and # the attend impl reads them back (so nothing crosses the eager break as a diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index 8d4d7090eded..c50ce7761438 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -671,13 +671,38 @@ def _pynvvc_frames_to_nhwc(frames: torch.Tensor) -> torch.Tensor: return frames.contiguous() +class _PyNvDecoderPool: + """Process-wide singleton managing PyNvVideoCodec decoder slot state. + + Prevents subclass counter shadowing (GHSA-j682-9xp5-rrf3) by storing + all mutable pool state in a single module-level instance rather than + in ClassVar attributes that get shadowed by Python's augmented + assignment semantics on subclasses. + """ + + def __init__(self) -> None: + self.slots: list[PyNvVideoCodecDecoderSlot] = [] + self.active: int = 0 + self.cond: threading.Condition = threading.Condition() + self.max_slots: int | None = None + + def configure(self, hw_decoders: int) -> None: + with self.cond: + if self.max_slots is None: + self.max_slots = hw_decoders + elif self.max_slots != hw_decoders: + raise RuntimeError( + "PyNvVideoCodec decoder count is already configured as " + f"{self.max_slots}, got {hw_decoders}" + ) + + +_pynv_decoder_pool = _PyNvDecoderPool() + + class PyNvVideoCodecVideoBackendMixin: """PyNvVideoCodec utilities for GPU-backed frame decode.""" - _decoder_slots: ClassVar[list[PyNvVideoCodecDecoderSlot]] = [] - _active_decoder_slots: ClassVar[int] = 0 - _decoder_slot_cond: ClassVar[threading.Condition] = threading.Condition() - _max_decoder_slots: ClassVar[int | None] = None _DEVICE_INDEX: ClassVar[int] = 0 @classmethod @@ -704,14 +729,7 @@ def _create_decoder_slot(cls) -> PyNvVideoCodecDecoderSlot: @classmethod def _configure_decoder_slots(cls, hw_decoders: object) -> None: hw_decoders = validate_pynvvideocodec_hw_decoders(hw_decoders) - with cls._decoder_slot_cond: - if cls._max_decoder_slots is None: - cls._max_decoder_slots = hw_decoders - elif cls._max_decoder_slots != hw_decoders: - raise RuntimeError( - "PyNvVideoCodec decoder count is already configured as " - f"{cls._max_decoder_slots}, got {hw_decoders}" - ) + _pynv_decoder_pool.configure(hw_decoders) @staticmethod @contextmanager @@ -729,28 +747,28 @@ def _torch_stream_context(stream): @classmethod @contextmanager def _borrow_decoder_slot(cls): + pool = _pynv_decoder_pool create_slot = False - with cls._decoder_slot_cond: - max_decoder_slots = cls._max_decoder_slots - if max_decoder_slots is None: + with pool.cond: + if pool.max_slots is None: raise RuntimeError("PyNvVideoCodec decoder slots are not configured") while True: - if cls._decoder_slots: - slot = cls._decoder_slots.pop() + if pool.slots: + slot = pool.slots.pop() break - if cls._active_decoder_slots < max_decoder_slots: - cls._active_decoder_slots += 1 + if pool.active < pool.max_slots: + pool.active += 1 create_slot = True break - cls._decoder_slot_cond.wait() + pool.cond.wait() if create_slot: try: slot = cls._create_decoder_slot() except Exception: - with cls._decoder_slot_cond: - cls._active_decoder_slots -= 1 - cls._decoder_slot_cond.notify() + with pool.cond: + pool.active -= 1 + pool.cond.notify() raise borrow_succeeded = False @@ -760,9 +778,9 @@ def _borrow_decoder_slot(cls): finally: if not borrow_succeeded: slot.invalidate() - with cls._decoder_slot_cond: - cls._decoder_slots.append(slot) - cls._decoder_slot_cond.notify() + with pool.cond: + pool.slots.append(slot) + pool.cond.notify() @staticmethod def _metadata_value(metadata, *names: str, default=None): diff --git a/vllm/outputs.py b/vllm/outputs.py index 986a315da527..29584e0e34cc 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -251,7 +251,7 @@ def __init__( self.finished = finished self.outputs = outputs - def __repr__(self): + def __repr__(self) -> str: return ( f"{type(self).__name__}(request_id={self.request_id!r}, " f"outputs={self.outputs!r}, " @@ -273,7 +273,7 @@ class EmbeddingOutput: embedding: list[float] @staticmethod - def from_base(pooling_output: PoolingOutput): + def from_base(pooling_output: PoolingOutput) -> "EmbeddingOutput": pooled_data = pooling_output.data if pooled_data.ndim != 1: raise ValueError("pooled_data should be a 1-D embedding vector") @@ -290,7 +290,9 @@ def __repr__(self) -> str: class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]): @staticmethod - def from_base(request_output: PoolingRequestOutput): + def from_base( + request_output: PoolingRequestOutput, + ) -> "EmbeddingRequestOutput": return EmbeddingRequestOutput( request_id=request_output.request_id, outputs=EmbeddingOutput.from_base(request_output.outputs), @@ -312,7 +314,7 @@ class ClassificationOutput: probs: list[float] @staticmethod - def from_base(pooling_output: PoolingOutput): + def from_base(pooling_output: PoolingOutput) -> "ClassificationOutput": # pooling_output shape: (num_classes) pooled_data = pooling_output.data if pooled_data.ndim != 1: @@ -330,7 +332,9 @@ def __repr__(self) -> str: class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]): @staticmethod - def from_base(request_output: PoolingRequestOutput): + def from_base( + request_output: PoolingRequestOutput, + ) -> "ClassificationRequestOutput": return ClassificationRequestOutput( request_id=request_output.request_id, outputs=ClassificationOutput.from_base(request_output.outputs), @@ -351,7 +355,7 @@ class ScoringOutput: score: float @staticmethod - def from_base(pooling_output: PoolingOutput): + def from_base(pooling_output: PoolingOutput) -> "ScoringOutput": # pooling_output shape: # classify task: (num_classes) num_classes == 1 # embed task: a scalar value @@ -367,7 +371,9 @@ def __repr__(self) -> str: class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]): @staticmethod - def from_base(request_output: PoolingRequestOutput): + def from_base( + request_output: PoolingRequestOutput, + ) -> "ScoringRequestOutput": return ScoringRequestOutput( request_id=request_output.request_id, outputs=ScoringOutput.from_base(request_output.outputs), diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index f7119d84ef8f..2eda37b432c3 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -367,6 +367,10 @@ def parse_delta( tool call extraction via internal stream state. """ + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + """Return the number of reasoning tokens in generated token IDs.""" + return 0 + class DelegatingParser(Parser): """ @@ -939,6 +943,12 @@ def parse_delta( return delta_message + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + """Count reasoning tokens through the configured reasoning parser.""" + if self._reasoning_parser is None: + return 0 + return self._reasoning_parser.count_reasoning_tokens(token_ids) + def _flush_engine_parsers( self, delta_message: DeltaMessage | None ) -> DeltaMessage | None: diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index b9c4e42b1704..946367e13709 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -47,6 +47,11 @@ class ParserEngineReasoningAdapter(ReasoningParser): def __init__(self, tokenizer: TokenizerLike, *args, **kwargs) -> None: super().__init__(tokenizer, *args, **kwargs) self._parser_engine = self._parser_engine_cls(tokenizer, **kwargs) # type: ignore[call-arg] + self._parser_engine_kwargs = kwargs + self._counting_parser_engine: ParserEngine | None = None + # TODO: Remove once Responses finalization reuses accumulated streaming + # parser results instead of reparsing the complete output. + self._streaming_count_valid = False @contextmanager def _skip_tool_parsing(self) -> Iterator[None]: @@ -71,6 +76,7 @@ def extract_reasoning( model_output: str, request: ChatCompletionRequest | ResponsesRequest, ) -> tuple[str | None, str | None]: + self._streaming_count_valid = False with self._skip_tool_parsing(): return self._parser_engine.extract_reasoning(model_output, request) @@ -83,6 +89,7 @@ def extract_reasoning_streaming( current_token_ids: Sequence[int], delta_token_ids: Sequence[int], ) -> DeltaMessage | None: + self._streaming_count_valid = True with self._skip_tool_parsing(): return self._parser_engine.extract_reasoning_streaming( previous_text, @@ -122,7 +129,18 @@ def get_streaming_fallback_content( return self._parser_engine.get_streaming_fallback_content(text, request) def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - return self._parser_engine.count_reasoning_tokens(token_ids) + if self._streaming_count_valid: + return self._parser_engine.count_reasoning_tokens(token_ids) + if not token_ids: + return 0 + if self._counting_parser_engine is None: + self._counting_parser_engine = self._parser_engine_cls( + self.model_tokenizer, **self._parser_engine_kwargs + ) # type: ignore[call-arg] + self._counting_parser_engine._single_pass_parse( + self.model_tokenizer.decode(token_ids), token_ids + ) + return self._counting_parser_engine.count_reasoning_tokens(token_ids) class ParserEngineToolAdapter(ToolParser): diff --git a/vllm/parser/engine/events.py b/vllm/parser/engine/events.py index f138fb248f45..08eaf2d01fc7 100644 --- a/vllm/parser/engine/events.py +++ b/vllm/parser/engine/events.py @@ -24,3 +24,4 @@ class SemanticEvent: type: EventType value: str = "" tool_index: int = -1 + token_count: int = 0 diff --git a/vllm/parser/engine/incremental_lexer.py b/vllm/parser/engine/incremental_lexer.py index 31e9bd4a3b2c..34345fe1d8f4 100644 --- a/vllm/parser/engine/incremental_lexer.py +++ b/vllm/parser/engine/incremental_lexer.py @@ -24,6 +24,7 @@ class TerminalDef: class LexToken: terminal: str value: str + token_count: int = 0 class LexerShape: @@ -104,6 +105,7 @@ def __init__( self.terminals = shape.terminals self.content_terminal = content_terminal self.buffer = "" + self._token_counts: list[int] = [] self._literal_strings = shape.literal_strings self._max_literal_len = shape.max_literal_len @@ -114,15 +116,23 @@ def __init__( def reset(self) -> None: self.buffer = "" + self._token_counts.clear() - def feed(self, text: str) -> list[LexToken]: + def feed( + self, + text: str, + token_texts: tuple[str, ...] = (), + token_count: int = 0, + ) -> list[LexToken]: + char_token_counts = self._char_token_counts(text, token_texts, token_count) if not self.buffer and self._has_only_literals and self._literal_first_chars: for ch in text: if ch in self._literal_first_chars: break else: - return [LexToken(self.content_terminal, text)] + return [LexToken(self.content_terminal, text, sum(char_token_counts))] self.buffer += text + self._token_counts.extend(char_token_counts) return self._drain() def flush(self) -> list[LexToken]: @@ -130,10 +140,46 @@ def flush(self) -> list[LexToken]: if self.buffer: tokens.extend(self._drain(final=True)) if self.buffer: - tokens.append(LexToken(self.content_terminal, self.buffer)) + tokens.append( + LexToken(self.content_terminal, self.buffer, sum(self._token_counts)) + ) self.buffer = "" + self._token_counts.clear() return tokens + @staticmethod + def _char_token_counts( + text: str, + token_texts: tuple[str, ...], + token_count: int, + ) -> list[int]: + counts = [0] * len(text) + if not text: + return counts + if token_texts: + pos = 0 + assigned = 0 + for token_text in token_texts: + if not token_text: + continue + found = text.find(token_text, pos) + if found < 0: + continue + counts[found] += 1 + assigned += 1 + pos = found + len(token_text) + missing = token_count - assigned + if missing > 0: + counts[0] += missing + elif token_count: + counts[0] = token_count + return counts + + def _pop_token_count(self, length: int) -> int: + token_count = sum(self._token_counts[:length]) + del self._token_counts[:length] + return token_count + def _drain(self, *, final: bool = False) -> list[LexToken]: tokens: list[LexToken] = [] first_chars = self._literal_first_chars @@ -150,8 +196,11 @@ def _drain(self, *, final: bool = False) -> list[LexToken]: has_potential = True break if not has_potential: - tokens.append(LexToken(content_terminal, self.buffer)) + tokens.append( + LexToken(content_terminal, self.buffer, sum(self._token_counts)) + ) self.buffer = "" + self._token_counts.clear() break best_match: tuple[str, str, int] | None = None @@ -175,7 +224,13 @@ def _drain(self, *, final: bool = False) -> list[LexToken]: longer_match = True break if not longer_match: - tokens.append(LexToken(best_match[0], best_match[1])) + tokens.append( + LexToken( + best_match[0], + best_match[1], + self._pop_token_count(best_match[2]), + ) + ) self.buffer = self.buffer[best_match[2] :] continue break @@ -183,15 +238,33 @@ def _drain(self, *, final: bool = False) -> list[LexToken]: break if best_match is not None: - tokens.append(LexToken(best_match[0], best_match[1])) + tokens.append( + LexToken( + best_match[0], + best_match[1], + self._pop_token_count(best_match[2]), + ) + ) self.buffer = self.buffer[best_match[2] :] else: content_end = self._find_content_boundary() if content_end > 0: - tokens.append(LexToken(content_terminal, self.buffer[:content_end])) + tokens.append( + LexToken( + content_terminal, + self.buffer[:content_end], + self._pop_token_count(content_end), + ) + ) self.buffer = self.buffer[content_end:] else: - tokens.append(LexToken(content_terminal, self.buffer[0])) + tokens.append( + LexToken( + content_terminal, + self.buffer[0], + self._pop_token_count(1), + ) + ) self.buffer = self.buffer[1:] return tokens diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 048e714cb4ae..5ed837d22df1 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -625,23 +625,8 @@ def get_streaming_fallback_content( return None def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: - start_id = self._reasoning_start_token_id - end_id = self._reasoning_end_token_id - if start_id is None or end_id is None: - return 0 - count = 0 - depth = 0 - for token_id in token_ids: - if token_id == start_id: - depth += 1 - continue - if token_id == end_id: - if depth > 0: - depth -= 1 - continue - if depth > 0: - count += 1 - return count + """Return reasoning tokens observed by the parser engine so far.""" + return self._engine.reasoning_token_count # ── Single-pass parse helper ──────────────────────────────────────── diff --git a/vllm/parser/engine/streaming_parser_engine.py b/vllm/parser/engine/streaming_parser_engine.py index 157a1796e854..ec21b05f4d3d 100644 --- a/vllm/parser/engine/streaming_parser_engine.py +++ b/vllm/parser/engine/streaming_parser_engine.py @@ -167,6 +167,17 @@ def __init__( self.skip_tool_parsing = False self.reset(initial_state=initial_state) + @property + def reasoning_token_count(self) -> int: + return self._reasoning_token_count + + def _record_reasoning_tokens(self, events: Sequence[SemanticEvent]) -> None: + self._reasoning_token_count += sum( + event.token_count + for event in events + if event.type == EventType.REASONING_CHUNK + ) + def _reset_args_state(self) -> None: self._args_buffer: str = "" self._args_safe_end: int = 0 @@ -186,6 +197,7 @@ def reset(self, initial_state: ParserState | None = None) -> None: ) self.tool_index = -1 self._ever_had_token_ids = False + self._reasoning_token_count = 0 # DO NOT reset skip_tool_parsing here — callers set it before # calling methods that trigger reset() (e.g. extract_reasoning), # and clearing it silently breaks non-streaming tool-call-as- @@ -193,6 +205,7 @@ def reset(self, initial_state: ParserState | None = None) -> None: self._scanner.reset() self._lexer.reset() self._message_header_buffer = "" + self._message_header_token_count = 0 self._in_skipped_tool_span = False self._reset_args_state() @@ -218,18 +231,30 @@ def feed( has_special = True break if not has_special: - return self._emit_for_state(delta_text) + events = self._emit_for_state( + delta_text, token_count=len(delta_token_ids) + ) + self._record_reasoning_tokens(events) + return events scanner_items = self._scanner.scan(delta_text, delta_token_ids) if len(scanner_items) == 1 and isinstance(scanner_items[0], TextChunk): - lex_tokens = self._lexer.feed(scanner_items[0].text) + item = scanner_items[0] + lex_tokens = self._lexer.feed(item.text, item.token_texts, item.token_count) if len(lex_tokens) == 1 and lex_tokens[0].terminal == CONTENT_TERMINAL: - text = lex_tokens[0].value - return self._emit_for_state(text) - return self._process_lex_tokens(lex_tokens) + events = self._emit_for_state( + lex_tokens[0].value, + token_count=lex_tokens[0].token_count, + ) + else: + events = self._process_lex_tokens(lex_tokens) + self._record_reasoning_tokens(events) + return events - return self._process_scanner_items(scanner_items) + events = self._process_scanner_items(scanner_items) + self._record_reasoning_tokens(events) + return events def _process_scanner_items( self, items: Sequence[LexerInput] @@ -240,7 +265,18 @@ def _process_scanner_items( events.extend(self._process_lex_tokens(self._lexer.flush())) events.extend(self._on_terminal(item.terminal, item.text)) elif isinstance(item, TextChunk): - events.extend(self._process_lex_tokens(self._lexer.feed(item.text))) + if not item.text and item.token_count: + events.extend( + self._emit_for_state("", token_count=item.token_count) + ) + else: + events.extend( + self._process_lex_tokens( + self._lexer.feed( + item.text, item.token_texts, item.token_count + ) + ) + ) return events def finish(self) -> list[SemanticEvent]: @@ -285,11 +321,14 @@ def finish(self) -> list[SemanticEvent]: EventType.TEXT_CHUNK, value=self._message_header_buffer, tool_index=self.tool_index, + token_count=self._message_header_token_count, ) ) self._message_header_buffer = "" + self._message_header_token_count = 0 self.state = ParserState.CONTENT + self._record_reasoning_tokens(events) return events def parse_complete(self, text: str) -> list[SemanticEvent]: @@ -303,9 +342,11 @@ def _process_lex_tokens(self, tokens: list[LexToken]) -> list[SemanticEvent]: strict = self._token_id_terminal_names if self._ever_had_token_ids else None for tok in tokens: if tok.terminal == CONTENT_TERMINAL or (strict and tok.terminal in strict): - events.extend(self._on_content(tok.value)) + events.extend(self._on_content(tok.value, tok.token_count)) else: - events.extend(self._on_terminal(tok.terminal, tok.value)) + events.extend( + self._on_terminal(tok.terminal, tok.value, tok.token_count) + ) return events _TOOL_STATES = frozenset( @@ -317,7 +358,9 @@ def _process_lex_tokens(self, tokens: list[LexToken]) -> list[SemanticEvent]: } ) - def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: + def _on_terminal( + self, terminal: str, value: str, token_count: int = 0 + ) -> list[SemanticEvent]: key = (self.state, terminal) transition = self.config.transitions.get(key) @@ -327,7 +370,7 @@ def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: # The projected skip state may not define the wrapper closer. if self.skip_tool_parsing and terminal in self._tool_exit_terminals: self._in_skipped_tool_span = False - return self._emit_for_state(value) + return self._emit_for_state(value, token_count) if self.skip_tool_parsing and terminal in self._tool_terminals: # Inkling reuses one terminal for tool, text, and reasoning exits. @@ -345,6 +388,7 @@ def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: leaving_message_header = self.state == ParserState.MESSAGE_HEADER if leaving_message_header: self._message_header_buffer = "" + self._message_header_token_count = 0 # A tool terminal that implicitly ends reasoning must report # that even from the header state, or the reasoning pass never # hands the block to the tool pass. @@ -381,13 +425,14 @@ def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: return [] if transition.skip_in_token_id_mode and self._ever_had_token_ids: - return self._emit_for_state(value) + return self._emit_for_state(value, token_count) - return self._apply_transition(transition, value) + return self._apply_transition(transition, value, token_count) - def _emit_for_state(self, text: str) -> list[SemanticEvent]: + def _emit_for_state(self, text: str, token_count: int = 0) -> list[SemanticEvent]: if self.state == ParserState.MESSAGE_HEADER: self._message_header_buffer += text + self._message_header_token_count += token_count return [] if self.state == ParserState.TOOL_ARGS: if self.config.tool_args_json: @@ -397,26 +442,36 @@ def _emit_for_state(self, text: str) -> list[SemanticEvent]: EventType.ARG_VALUE_CHUNK, value=text, tool_index=self.tool_index, + token_count=token_count, ) ] content_type = self.config.content_events.get(self.state) if content_type is not None: - return [SemanticEvent(content_type, value=text, tool_index=self.tool_index)] + return [ + SemanticEvent( + content_type, + value=text, + tool_index=self.tool_index, + token_count=token_count, + ) + ] return [] - def _on_content(self, text: str) -> list[SemanticEvent]: + def _on_content(self, text: str, token_count: int = 0) -> list[SemanticEvent]: if not text: return [] - return self._emit_for_state(text) + return self._emit_for_state(text, token_count) def _apply_transition( self, transition: Transition, value: str, + token_count: int = 0, ) -> list[SemanticEvent]: events: list[SemanticEvent] = [] previous_state = self.state message_header = "" + message_header_token_count = 0 if ( self.state == ParserState.TOOL_ARGS @@ -434,7 +489,9 @@ def _apply_transition( if previous_state == ParserState.MESSAGE_HEADER: message_header = self._message_header_buffer + message_header_token_count = self._message_header_token_count self._message_header_buffer = "" + self._message_header_token_count = 0 self.state = transition.next_state @@ -454,6 +511,12 @@ def _apply_transition( event_type, value=event_value, tool_index=self.tool_index, + token_count=( + message_header_token_count + if previous_state == ParserState.MESSAGE_HEADER + and event_type == EventType.TEXT_CHUNK + else token_count + ), ) ) diff --git a/vllm/parser/engine/token_id_scanner.py b/vllm/parser/engine/token_id_scanner.py index abcc2e2baec6..60a64dcae281 100644 --- a/vllm/parser/engine/token_id_scanner.py +++ b/vllm/parser/engine/token_id_scanner.py @@ -14,6 +14,8 @@ @dataclass(slots=True) class TextChunk: text: str + token_texts: tuple[str, ...] = () + token_count: int = 0 @dataclass(slots=True) @@ -49,18 +51,30 @@ def __init__( self.tokenizer = tokenizer self._token_text_cache: dict[int, str] = {} self._deferred_terminals: list[PreLexedTerminal] = [] + self._deferred_prefix_token_counts: list[int] = [] + self._deferred_trailing_token_count = 0 self._deferred_post_text: str = "" def reset(self) -> None: """Clear mutable state for reuse. Preserves the token text cache.""" self._deferred_terminals.clear() + self._deferred_prefix_token_counts.clear() + self._deferred_trailing_token_count = 0 self._deferred_post_text = "" def _decode_token(self, token_id: int) -> str: if token_id not in self._token_text_cache: - self._token_text_cache[token_id] = self.tokenizer.decode([token_id]) + token_text = self.tokenizer.decode([token_id]) + self._token_text_cache[token_id] = ( + token_text if isinstance(token_text, str) else "" + ) return self._token_text_cache[token_id] + def _decode_tokens(self, token_ids: Sequence[int]) -> tuple[str, ...]: + if self.tokenizer is None: + return () + return tuple(self._decode_token(tid) for tid in token_ids) + _EMPTY: tuple[LexerInput, ...] = () def scan( @@ -70,13 +84,23 @@ def scan( ) -> Sequence[LexerInput]: prefix_items: list[LexerInput] = [] effective_text = delta_text + deferred_trailing_count = 0 if self._deferred_terminals: - prefix_items, effective_text = self._resolve_deferred(delta_text) + prefix_items, effective_text, deferred_trailing_count = ( + self._resolve_deferred(delta_text, len(delta_token_ids)) + ) if not self.token_id_to_terminal: if effective_text: - prefix_items.append(TextChunk(effective_text)) + token_texts = self._decode_tokens(delta_token_ids) + prefix_items.append( + TextChunk( + effective_text, + token_texts, + deferred_trailing_count + len(token_texts), + ) + ) return prefix_items has_special = False @@ -89,14 +113,23 @@ def scan( if not has_special: if effective_text: if not prefix_items: - return [TextChunk(effective_text)] - prefix_items.append(TextChunk(effective_text)) + token_texts = self._decode_tokens(delta_token_ids) + return [TextChunk(effective_text, token_texts, len(token_texts))] + token_texts = self._decode_tokens(delta_token_ids) + prefix_items.append( + TextChunk( + effective_text, + token_texts, + deferred_trailing_count + len(token_texts), + ) + ) return prefix_items or self._EMPTY - token_texts = [self._decode_token(tid) for tid in delta_token_ids] + decoded_token_texts = [self._decode_token(tid) for tid in delta_token_ids] results: list[LexerInput] = [] text_accum: list[str] = [] + token_text_accum: list[str] = [] for idx, tid in enumerate(delta_token_ids): terminal = self.token_id_to_terminal.get(tid) @@ -104,16 +137,31 @@ def scan( if text_accum: joined = "".join(text_accum) if joined: - results.append(TextChunk(joined)) + results.append( + TextChunk( + joined, + tuple(token_text_accum), + len(token_text_accum), + ) + ) text_accum.clear() - results.append(PreLexedTerminal(terminal, tid, token_texts[idx])) + token_text_accum.clear() + results.append( + PreLexedTerminal(terminal, tid, decoded_token_texts[idx]) + ) else: - text_accum.append(token_texts[idx]) + text_accum.append(decoded_token_texts[idx]) + token_text_accum.append(decoded_token_texts[idx]) if text_accum: joined = "".join(text_accum) if joined: - results.append(TextChunk(joined)) + results.append( + TextChunk(joined, tuple(token_text_accum), len(token_text_accum)) + ) + + if deferred_trailing_count: + results.insert(0, TextChunk("", token_count=deferred_trailing_count)) if effective_text: results = self._recover_holdback_text(effective_text, results) @@ -124,9 +172,15 @@ def scan( # transition before the preceding text has arrived. The # deferred terminals will be resolved against the actual # delta_text in a subsequent scan() or flushed by finish(). + prefix_token_count = 0 for r in results: + if isinstance(r, TextChunk): + prefix_token_count += r.token_count if isinstance(r, PreLexedTerminal): self._deferred_terminals.append(r) + self._deferred_prefix_token_counts.append(prefix_token_count) + prefix_token_count = 0 + self._deferred_trailing_token_count += prefix_token_count results = [] return prefix_items + results @@ -136,16 +190,30 @@ def flush_pending(self) -> list[LexerInput]: return [] results: list[LexerInput] = [] if self._deferred_post_text: - results.append(TextChunk(self._deferred_post_text)) + prefix_count = ( + self._deferred_prefix_token_counts[0] + if self._deferred_prefix_token_counts + else 0 + ) + results.append( + TextChunk(self._deferred_post_text, token_count=prefix_count) + ) self._deferred_post_text = "" results.extend(self._deferred_terminals) + if self._deferred_trailing_token_count: + results.append( + TextChunk("", token_count=self._deferred_trailing_token_count) + ) self._deferred_terminals.clear() + self._deferred_prefix_token_counts.clear() + self._deferred_trailing_token_count = 0 return results def _resolve_deferred( self, delta_text: str, - ) -> tuple[list[LexerInput], str]: + current_token_count: int = 0, + ) -> tuple[list[LexerInput], str, int]: """Resolve deferred terminals against new delta_text. When a previous ``scan()`` deferred a terminal (its text hadn't @@ -160,7 +228,11 @@ def _resolve_deferred( should be scanned with the current delta's token IDs. """ deferred = self._deferred_terminals + prefix_token_counts = self._deferred_prefix_token_counts + trailing_token_count = self._deferred_trailing_token_count self._deferred_terminals = [] + self._deferred_prefix_token_counts = [] + self._deferred_trailing_token_count = 0 results: list[LexerInput] = [] remaining = delta_text @@ -171,10 +243,13 @@ def _resolve_deferred( # Duplicate-text deferred terminals resolve left-to-right via # find(); correct when each terminal text appears once in sequence. - for terminal in deferred: + for idx, terminal in enumerate(deferred): + prefix_token_count = prefix_token_counts[idx] pos = remaining.find(terminal.text) if pos > 0: - results.append(TextChunk(remaining[:pos])) + results.append( + TextChunk(remaining[:pos], token_count=prefix_token_count) + ) results.append(terminal) remaining = remaining[pos + len(terminal.text) :] elif pos == 0: @@ -185,10 +260,16 @@ def _resolve_deferred( # only the terminal provides a reliable split point. if remaining: self._deferred_post_text += remaining + prefix_token_count += current_token_count remaining = "" self._deferred_terminals.append(terminal) + self._deferred_prefix_token_counts.append(prefix_token_count) + + if self._deferred_terminals: + self._deferred_trailing_token_count = trailing_token_count + trailing_token_count = 0 - return results, remaining + return results, remaining, trailing_token_count def _recover_holdback_text( self, @@ -257,6 +338,16 @@ def _rebuild_from_anchors( if not anchors: return [TextChunk(delta_text)] + token_groups: list[list[str]] = [[] for _ in range(len(anchors) + 1)] + count_groups = [0] * (len(anchors) + 1) + group_idx = 0 + for item in results: + if isinstance(item, PreLexedTerminal): + group_idx += 1 + else: + token_groups[group_idx].extend(item.token_texts) + count_groups[group_idx] += item.token_count + # Resolve positions right-to-left: each anchor gets the # rightmost occurrence that is still before the next anchor. positions: list[int] = [-1] * len(anchors) @@ -274,7 +365,13 @@ def _rebuild_from_anchors( pos = positions[i] if pos >= consumed: if pos > consumed: - new_results.append(TextChunk(delta_text[consumed:pos])) + new_results.append( + TextChunk( + delta_text[consumed:pos], + tuple(token_groups[i]), + count_groups[i], + ) + ) new_results.append(anchor) consumed = pos + len(anchor.text) else: @@ -291,6 +388,13 @@ def _rebuild_from_anchors( self._deferred_post_text += delta_text[consumed:] consumed = len(delta_text) self._deferred_terminals.append(anchor) + self._deferred_prefix_token_counts.append(count_groups[i]) if consumed < len(delta_text): - new_results.append(TextChunk(delta_text[consumed:])) + new_results.append( + TextChunk( + delta_text[consumed:], + tuple(token_groups[-1]), + count_groups[-1], + ) + ) return new_results diff --git a/vllm/parser/harmony.py b/vllm/parser/harmony.py index 3cc493b65923..c0be3ea0da9b 100644 --- a/vllm/parser/harmony.py +++ b/vllm/parser/harmony.py @@ -388,11 +388,9 @@ def _normalize_recipient(recipient: str | None) -> str | None: # Harmomy's stop tokens are <|return|>, <|call|>, <|endoftext|> -# <|return|> is represented as "" since it's the default stop token, which xgrammar -# disallows under constraints, leading to bad or infinite generation. -# StreamableParser doesn't consider <|endoftext|> as a message end, so it's excluded -# TODO: Remove <|call|> once #50595 lands. -_END_TAG = ["<|end|>", "<|call|>", ""] +# They are represented as "" since xgrammar disallows stop tokens while +# under constraints, leading to bad or infinite generations. +_END_TAG = ["<|end|>", ""] _FINAL_BEGIN = "<|channel|>final{constrain}<|message|>" _TOOL_CALL_CHANNELS = [ "<|channel|>commentary", diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index 865f749a79c9..40af5b801f3e 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -9,6 +9,7 @@ if TYPE_CHECKING: from vllm.parser.abstract_parser import Parser + from vllm.parser.engine.parser_engine import ParserEngine from vllm.reasoning import ReasoningParser from vllm.tool_parsers import ToolParser @@ -17,10 +18,30 @@ class ParserManager: """ - Provides a unified Parser by composing individual reasoning and tool - parsers from their respective registries. + Provides a unified Parser from the reasoning and tool parser registries. + + Parser engine adapters backed by the same engine are collapsed back into + that engine. Other parser pairs are composed through ``DelegatingParser``. """ + @staticmethod + def _get_parser_engine_cls( + parser_cls: type[object] | None, + ) -> type[ParserEngine] | None: + if parser_cls is None: + return None + parser_engine_cls = getattr(parser_cls, "_parser_engine_cls", None) + if parser_engine_cls is None: + return None + + from vllm.parser.engine.parser_engine import ParserEngine + + if not isinstance(parser_engine_cls, type) or not issubclass( + parser_engine_cls, ParserEngine + ): + return None + return parser_engine_cls + @classmethod def get_tool_parser( cls, @@ -84,8 +105,8 @@ def get_parser( """ Get a Parser that handles both reasoning and tool parsing. - Composes individual reasoning and tool parsers into a single - DelegatingParser subclass. + Reuses a shared parser engine when possible, otherwise composes the + individual parsers into a ``DelegatingParser`` subclass. Args: tool_parser_name: The name of the tool parser. @@ -116,6 +137,11 @@ def get_parser( HarmonyParser.tool_parser_cls = tool_parser_cls return HarmonyParser + reasoning_engine_cls = cls._get_parser_engine_cls(reasoning_parser_cls) + tool_engine_cls = cls._get_parser_engine_cls(tool_parser_cls) + if reasoning_engine_cls is not None and reasoning_engine_cls is tool_engine_cls: + return reasoning_engine_cls + if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3": from vllm.parser.kimi_k3 import KimiK3Parser diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py index ac536aff00c5..718e2520c4f2 100644 --- a/vllm/platforms/__init__.py +++ b/vllm/platforms/__init__.py @@ -160,24 +160,32 @@ def _is_amd_zen_cpu() -> bool: def cpu_platform_plugin() -> str | None: - is_cpu = False logger.debug("Checking if CPU platform is available.") - try: - is_cpu = vllm_version_matches_substr("cpu") - if is_cpu: - logger.debug( - "Confirmed CPU platform is available because vLLM is built with CPU." - ) - if not is_cpu: - import sys - - is_cpu = sys.platform.startswith("darwin") + is_cpu = envs.VLLM_TARGET_DEVICE == "cpu" + if is_cpu: + logger.debug( + "Confirmed CPU platform is available because " + "VLLM_TARGET_DEVICE is set to CPU." + ) + else: + try: + is_cpu = vllm_version_matches_substr("cpu") if is_cpu: logger.debug( - "Confirmed CPU platform is available because the machine is MacOS." + "Confirmed CPU platform is available because vLLM is built " + "with CPU." ) - except Exception as e: - logger.debug("CPU platform is not available because: %s", str(e)) + if not is_cpu: + import sys + + is_cpu = sys.platform.startswith("darwin") + if is_cpu: + logger.debug( + "Confirmed CPU platform is available because the machine " + "is MacOS." + ) + except Exception as e: + logger.debug("CPU platform is not available because: %s", str(e)) if not is_cpu: return None @@ -209,6 +217,15 @@ def cpu_platform_plugin() -> str | None: def resolve_current_platform_cls_qualname() -> str: + # An explicit CPU target is authoritative. Native CPU-only CI jobs reuse + # an accelerator wheel and can run on accelerator hosts, so probing every + # plugin would otherwise activate both CPU and the host accelerator. + if envs.VLLM_TARGET_DEVICE == "cpu": + cpu_platform_cls_qualname = cpu_platform_plugin() + assert cpu_platform_cls_qualname is not None + logger.debug("Explicitly selected CPU platform.") + return cpu_platform_cls_qualname + platform_plugins = load_plugins_by_group(PLATFORM_PLUGINS_GROUP) activated_plugins = [] diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 4eeb98a79c79..5633e160cc2c 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -70,6 +70,10 @@ def supported_dtypes(self) -> list[torch.dtype]: # x86/aarch64 CPU has supported both bf16 and fp16 natively. return [torch.bfloat16, torch.float16, torch.float32] + @classmethod + def check_runner_kv_caches_multi_layer(cls) -> None: + pass + @classmethod def get_device_name(cls, device_id: int = 0) -> str: return "cpu" diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index aaaccfae752f..bfe08a577213 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -229,6 +229,10 @@ def import_kernels(cls) -> None: with contextlib.suppress(ImportError): import vllm._qutlass_C # noqa: F401 + @classmethod + def check_runner_kv_caches_multi_layer(cls) -> None: + pass + @property def supported_dtypes(self) -> list[torch.dtype]: if self.has_device_capability(80): diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index f0b93da2922c..c3ade53ca139 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -612,7 +612,6 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: For hybrid models, also aligns block_size with mamba page sizes. """ from vllm.config.cache import CacheConfig - from vllm.config.vllm import set_current_vllm_config cache_config = vllm_config.cache_config model_config = vllm_config.model_config @@ -627,10 +626,9 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: # Phase 1: Pick block size from backend (skip if user set --block-size) if not cache_config.user_specified_block_size: - with set_current_vllm_config(vllm_config): - preferred = backend_cls.get_preferred_block_size( - CacheConfig.DEFAULT_BLOCK_SIZE - ) + preferred = backend_cls.get_preferred_block_size_for_config( + CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config + ) if preferred != CacheConfig.DEFAULT_BLOCK_SIZE: logger.info( "Setting kv cache block size to %d for %s backend.", @@ -689,13 +687,15 @@ def _align_heterogeneous_kv_block_size( def per_token_page_bytes(dtype: "torch.dtype", cache_dtype: str) -> int: """Bytes one token occupies in one layer, for the given dtype.""" - return FullAttentionSpec( + spec = FullAttentionSpec( block_size=1, num_kv_heads=model_config.get_num_kv_heads(parallel_config), head_size=model_config.get_head_size(), dtype=dtype, kv_quant_mode=get_kv_quant_mode(cache_dtype), - ).page_size_bytes + ) + # The backend owns its packing + return backend_cls.customize_spec(spec).page_size_bytes primary_dtype = ( STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype] @@ -811,23 +811,18 @@ def _align_hybrid_block_size( # when all attention layers are TQ. With mixed skip+TQ the skip # layers still use the standard layout — take max so mamba # padding covers the largest actual page. - from vllm.model_executor.layers.quantization.turboquant.config import ( - TurboQuantConfig, + from vllm.v1.attention.backends.turboquant_attn import ( + TurboQuantAttentionBackend, ) - from vllm.v1.kv_cache_interface import TQFullAttentionSpec - tq_cfg = TurboQuantConfig.from_cache_dtype( - cache_config.cache_dtype, model_config.get_head_size() - ) - tq_page = TQFullAttentionSpec( + tq_spec = FullAttentionSpec( block_size=1, num_kv_heads=model_config.get_num_kv_heads(parallel_config), head_size=model_config.get_head_size(), - head_size_v=model_config.get_head_size(), dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, - tq_slot_size=tq_cfg.slot_size_aligned, - ).page_size_bytes + ) + tq_page = TurboQuantAttentionBackend.customize_spec(tq_spec).page_size_bytes if cache_config.kv_cache_dtype_skip_layers: skip_page = FullAttentionSpec( block_size=1, @@ -842,12 +837,15 @@ def _align_hybrid_block_size( else: attn_page_size_1_token = tq_page else: - attn_page_size_1_token = FullAttentionSpec( + attn_spec = FullAttentionSpec( block_size=1, num_kv_heads=model_config.get_num_kv_heads(parallel_config), head_size=model_config.get_head_size(), dtype=kv_cache_dtype, kv_quant_mode=kv_quant_mode, + ) + attn_page_size_1_token = backend_cls.customize_spec( + attn_spec ).page_size_bytes # Compute mamba page size @@ -1200,6 +1198,23 @@ def support_static_graph_mode(cls) -> bool: """ return False + @classmethod + def check_runner_kv_caches_multi_layer(cls) -> None: + """ + Check whether the platform's ModelRunner can handle multiple attention + layers that share the same layer index (e.g. cross attention and self + attention in the same decoder block of an encoder-decoder model such as + BART). + + Platforms that have verified that their ``runner_kv_caches`` is not + impacted by this case should override this to a no-op. Otherwise the + default implementation raises ``NotImplementedError``. + """ + raise NotImplementedError( + "Multiple attention layers with the same layer index are not " + "supported on the current platform." + ) + @classmethod def support_deep_gemm(cls) -> bool: """ diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index d7045a7431fd..9072c7142699 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -547,6 +547,10 @@ def import_kernels(cls) -> None: with contextlib.suppress(ImportError): import vllm._rocm_C # noqa: F401 + @classmethod + def check_runner_kv_caches_multi_layer(cls) -> None: + pass + @classmethod def is_pin_memory_available(cls) -> bool: if in_wsl(): diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index dde7a223f7c1..ddfde225225b 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -134,6 +134,10 @@ def import_kernels(cls) -> None: with contextlib.suppress(ImportError): import vllm._moe_C # noqa: F401 + @classmethod + def check_runner_kv_caches_multi_layer(cls) -> None: + pass + @classmethod def get_attn_backend_cls( cls, diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 6c761147c902..8b52c4fca8ae 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -120,6 +120,10 @@ "olmo3_reasoning_parser", "Olmo3ReasoningParser", ), + "muse_glimmer": ( + "muse_glimmer_reasoning_parser", + "MuseGlimmerReasoningParser", + ), "qwen3": ( "qwen3_engine_reasoning_parser", "Qwen3ParserReasoningAdapter", diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py index 340f7e022c3c..dffea49731c2 100644 --- a/vllm/reasoning/cohere_command_reasoning_parser.py +++ b/vllm/reasoning/cohere_command_reasoning_parser.py @@ -16,7 +16,7 @@ raise ImportError( "The Cohere reasoning parser requires the `cohere_melody` " "package, which is not installed. Install it with:\n" - " pip install cohere_melody" + " pip install 'cohere-melody>=0.11.1'" ) from e @@ -132,7 +132,7 @@ def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: tool_dictionary[tool_name] = f"{tool_name} ::= {tool_name}root\n{tool_grammar}" # Emitted grammar shape: # root ::= tools - # tools ::= ws "[" ws tool ws ("," ws tool)* ws "]" ws + # tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws # ws ::= (" " | "\t" | "\n")* # tool ::= | | ... (one alternative per input) # ::= root (per-tool xgrammar rules) @@ -140,7 +140,7 @@ def collect_tool_schema(tool_schema: list[CohereNormalizedTool]) -> str: tool_alternatives = "tool ::= " + " | ".join(tool_dictionary.keys()) tool_rules = "\n ".join(tool_dictionary.values()) grammar = f"""root ::= tools - tools ::= ws "[" ws tool ws ("," ws tool)* ws "]" ws + tools ::= ws "[" ws tool (ws "," ws tool)* ws "]" ws ws ::= (" " | "\\t" | "\\n")* {tool_alternatives} {tool_rules} @@ -699,7 +699,9 @@ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd3().stream_non_grounded_answer(), + streaming_opts=( + PyFilterOptions().cmd3().no_tools().stream_non_grounded_answer() + ), unary_opts=PyFilterOptions().cmd3().no_tools(), **kwargs, ) @@ -710,7 +712,9 @@ def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): super().__init__( tokenizer, *args, - streaming_opts=PyFilterOptions().cmd4().stream_non_grounded_answer(), + streaming_opts=( + PyFilterOptions().cmd4().no_tools().stream_non_grounded_answer() + ), unary_opts=PyFilterOptions().cmd4().no_tools(), **kwargs, ) diff --git a/vllm/reasoning/muse_glimmer_reasoning_parser.py b/vllm/reasoning/muse_glimmer_reasoning_parser.py new file mode 100644 index 000000000000..d3abedca4b86 --- /dev/null +++ b/vllm/reasoning/muse_glimmer_reasoning_parser.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reasoning-content parser for MuseGlimmer. +Port of the ``reasoning_content`` rule from the HuggingFace MuseGlimmer +``MUSE_GLIMMER_RESPONSE_SCHEMA`` (synced with internal master). MuseGlimmer emits +chain-of-thought in ``to=self`` channels delimited by ``<|message|>`` ... ``<|eom|>``: + to=self<|message|>...reasoning...<|eom|> +A turn may contain several ``to=self`` blocks interleaved with tool calls, and a +tool call or final answer follows in its own channel. +Because MuseGlimmer's framing markers (``<|message|>``, ``<|eom|>``) are not guaranteed +to be single vocab tokens across every checkpoint's tokenizer, this parser works +on the decoded text with regexes rather than the single start/end-token base class. +Usage: ``--reasoning-parser muse_glimmer`` +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence + +import regex as re + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.reasoning.abs_reasoning_parsers import ReasoningParser + +_EOM = "<|eom|>" +_EOT = "<|eot|>" +_FUNCTION_CALLS_OPEN = "" +_REASONING_OPEN = "to=self<|message|>" +_ASSISTANT_TURN_OPEN = "<|start|>assistant" +# A channel header: ``to=<|message|>`` where recipient is ``self`` +# (reasoning), ``user`` (final answer) or ``[.]`` (tool call). +_CHANNEL_HEADER_RE = re.compile(r"to=(?P[^\s<]+)<\|message\|>") +_HEADER_PAT = r"to=[^\s<]+<\|message\|>" +# Collapse the gap between reasoning blocks so multiple to=self spans join. +_COLLAPSE_RE = re.compile( + r"<\|eom\|>(?:(?!to=self<\|message\|>).)*?to=self<\|message\|>", re.DOTALL +) +_REASONING_RE = re.compile(r"to=self<\|message\|>(.*?)<\|eom\|>", re.DOTALL) +_CONTENT_RE = re.compile( + r"to=user<\|message\|>(.*?)(?=<\|eot\|>|<\|eom\|>|$)", re.DOTALL +) +# Strip a CLOSED reasoning span (header .. <|eom|>). +_STRIP_REASONING_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>.*?<\|eom\|>", re.DOTALL +) +# An UNTERMINATED trailing reasoning span. The model sometimes leaves the +# analysis channel WITHOUT emitting <|eom|>, writing a bare +# ``to=<|message|>`` header instead (observed deterministically for a call +# with EMPTY arguments on a tool that has optional parameters; reproduced on +# other engines too, so it is a model-side defect, not engine-specific). +# +# These two patterns MUST therefore stop at the next channel header rather than +# running to end-of-text. An unbounded ``...$`` version consumes the real tool +# call along with the reasoning: `is_reasoning_end` then never fires, the parser +# never leaves the reasoning phase, the tool parser is never invoked, and the +# entire generation is dropped (empty reasoning, empty content, no tool call). +_STRIP_OPEN_REASONING_RE = re.compile( + r"(?:<\|start\|>assistant\s*)?to=self<\|message\|>" + r"(?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*" + r"(?=" + _HEADER_PAT + r"|$)", + re.DOTALL, +) +_OPEN_REASONING_RE = re.compile( + r"to=self<\|message\|>((?:(?!<\|eom\|>)(?!" + _HEADER_PAT + r").)*)" + r"(?=" + _HEADER_PAT + r"|$)", + re.DOTALL, +) +# Markers whose PREFIX could appear at the tail of an OPEN (still-streaming) body. +_HOLDBACK_MARKERS = (_EOM, _EOT, "<|start|>", "<|message|>") +# A trailing fragment that could still grow into a channel header (" t", " to", +# " to=", " to=skill"). Without this the recipient name leaks into reasoning and +# then has to be un-emitted once ``<|message|>`` arrives. +_OPEN_TAIL_HEADER_RE = re.compile(r"[\s](?:t|to|to=[^\s<]*)$") + + +def _current_assistant_turn(text: str) -> str: + """Return only the text generated in the current assistant turn. + ``is_reasoning_end`` is evaluated on the PROMPT token-ids at stream start, + and an MuseGlimmer prompt legitimately contains ATEM markers (``render_tool_defs`` + writes a literal ```` example into the system message, + and prior assistant turns may carry real tool calls). Anchoring on the last + channel-open keeps prompt text from deciding the phase. + """ + idx = text.rfind(_ASSISTANT_TURN_OPEN) + return text[idx + len(_ASSISTANT_TURN_OPEN) :] if idx != -1 else text + + +def _trim_open_body(body: str) -> str: + """Hold back any tail of a still-growing body that could still be framing. + Iterated to a fixpoint because the two cases compose: `" to=skill<"` needs + the partial-marker trim (``<``) before the partial-header trim can see + `" to=skill"`. Trimming only once leaks the recipient name as reasoning. + """ + while True: + trimmed = body + for marker in _HOLDBACK_MARKERS: + for k in range(min(len(marker) - 1, len(trimmed)), 0, -1): + if trimmed.endswith(marker[:k]): + trimmed = trimmed[:-k] + break + else: + continue + break + header_tail = _OPEN_TAIL_HEADER_RE.search(trimmed) + if header_tail is not None: + trimmed = trimmed[: header_tail.start()] + if trimmed == body: + return body + body = trimmed + + +class MuseGlimmerReasoningParser(ReasoningParser): + def __init__(self, tokenizer, *args, **kwargs) -> None: + super().__init__(tokenizer, *args, **kwargs) + # Cursors over what was ACTUALLY emitted. Diffing a freshly reclassified + # `previous_text` is unsafe: a classified body legitimately shrinks when + # a partial header becomes recognisable, and diffing against the shrunken + # value re-emits text that already went out. + self._emitted_reasoning: str = "" + self._emitted_content: str = "" + self._tool_handoff_done: bool = False + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Preserve MuseGlimmer's ATEM framing tokens in the decoded output. + vLLM's serving default is ``skip_special_tokens=True``, which strips + ``<|start|>`` / ``<|message|>`` / ``<|eom|>`` / ``<|eot|>`` before the + parsers run, collapsing reasoning into content and breaking channel + scoping. Unlike the base tool-parser hook we do NOT touch + ``structured_outputs`` -- MuseGlimmer emits native ATEM, not JSON. + """ + request.skip_special_tokens = False + return request + + def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: + """Whether the model has left reasoning and opened a TOOL channel. + A ``to=user`` answer is NOT a reason to leave the reasoning phase -- this + parser surfaces that content itself. Only a real tool channel switches + the ``DelegatingParser`` phase machine over to the tool parser. + Both closed and unterminated reasoning spans are stripped before the + check, so an ```` the model merely echoes inside its CoT + never flips the phase. + """ + try: + text = self.model_tokenizer.decode(input_ids) + except Exception: + return False + remainder = self._tool_channel_remainder(text) + return _FUNCTION_CALLS_OPEN in remainder or " bool: + return self.is_reasoning_end(list(input_ids)) + + def extract_content_ids(self, input_ids: list[int]) -> list[int]: + # Content-id slicing is unreliable for multi-token markers; the serving + # path uses extract_reasoning() for the final split. + return [] + + @classmethod + def _scoped_turn(cls, text: str) -> str: + """Current assistant turn with reasoning spans removed.""" + scoped = _current_assistant_turn(text) + scoped = _STRIP_REASONING_RE.sub("", scoped) + return _STRIP_OPEN_REASONING_RE.sub("", scoped) + + @classmethod + def _tool_channel_remainder(cls, text: str) -> str: + """Text from the first tool-channel header onward, framing INCLUDED. + ``DelegatingParser.parse_delta`` rebuilds ``current_text`` from whatever + this parser returns as ``.content`` on the transition delta and commits + it; anything not returned is destroyed. It must start AT the + ``to=<|message|>`` header -- handing over the text after the header + loses the recipient, and the tool parser then sees a bare ``<|message|>``, + classifies it as the content channel, and leaks the ATEM markup. + """ + scoped = cls._scoped_turn(text) + for match in _CHANNEL_HEADER_RE.finditer(scoped): + if match.group("recipient") not in ("self", "user"): + return scoped[match.start() :] + return "" + + @staticmethod + def _classify_bodies(text: str) -> tuple[str, str]: + """Split ``text`` into (reasoning_body, content_body), channel-aware. + Framing markers and tool channels contribute nothing -- the tool parser + owns those. A body ends at ``<|eom|>`` / ``<|eot|>``, at the next channel + header, or at end-of-text (an OPEN body, which is held back). + """ + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + pos = 0 + n = len(text) + while pos < n: + match = _CHANNEL_HEADER_RE.search(text, pos) + if not match: + break + recipient = match.group("recipient") + body_start = match.end() + eom = text.find(_EOM, body_start) + eot = text.find(_EOT, body_start) + terminators = [p for p in (eom, eot) if p != -1] + next_header = _CHANNEL_HEADER_RE.search(text, body_start) + if next_header is not None: + terminators.append(next_header.start()) + body_end = min(terminators) if terminators else n + body = text[body_start:body_end] + if not terminators: + body = _trim_open_body(body) + if recipient == "self": + reasoning_parts.append(body) + elif ( + # Never surface tool XML echoed into a user channel. + recipient == "user" + and _FUNCTION_CALLS_OPEN not in body + and " str | None: + """Promote un-surfaced content when the stream ends mid-reasoning. + ``DelegatingParser.finalize_generation`` calls this when + ``reasoning_ended`` is still False. Returns only the channel-classified + ``to=user`` body, and only the portion not already streamed. + """ + _, content_body = self._classify_bodies(previous_text) + remainder = content_body[len(self._emitted_content) :] + return remainder or None + + def extract_reasoning( + self, + model_output: str, + request: ChatCompletionRequest | ResponsesRequest, + ) -> tuple[str | None, str | None]: + collapsed = _COLLAPSE_RE.sub("\n", model_output) + matches = _REASONING_RE.findall(collapsed) + reasoning = "\n".join(matches) if matches else None + # Truncation fallback: generation stopped inside a to=self block, so + # there is no closing <|eom|>. Bounded at the next channel header so a + # real tool call that follows a header-less channel switch is not + # absorbed into the reasoning field. + open_match = _OPEN_REASONING_RE.search(model_output) + if open_match and open_match.group(1): + partial = open_match.group(1) + reasoning = f"{reasoning}\n{partial}" if reasoning else partial + # Content is everything that is not a reasoning block. In a + # reasoning+tool-call turn there is no to=user answer, but the tool + # channels MUST be forwarded -- the unified parser runs the tool parser + # on this returned `content`, not on the original model_output. + remainder = _STRIP_REASONING_RE.sub("", model_output) + remainder = _STRIP_OPEN_REASONING_RE.sub("", remainder) + if " DeltaMessage | None: + """Channel-aware streaming split of reasoning vs content. + Classifies the full ``current_text`` and emits only what has not been + emitted yet, so no framing token is ever surfaced and a delta straddling + a channel boundary only contributes the portion inside a real body. + """ + curr_reason, curr_content = self._classify_bodies(current_text) + reasoning_delta = "" + if curr_reason.startswith(self._emitted_reasoning) and len(curr_reason) > len( + self._emitted_reasoning + ): + reasoning_delta = curr_reason[len(self._emitted_reasoning) :] + self._emitted_reasoning = curr_reason + content_delta = "" + if curr_content.startswith(self._emitted_content) and len(curr_content) > len( + self._emitted_content + ): + content_delta = curr_content[len(self._emitted_content) :] + self._emitted_content = curr_content + # Hand the tool channel to the tool parser exactly once, starting at its + # header. parse_delta discards anything not returned here. + # + # This MUST fire on the same delta where is_reasoning_end() flips, i.e. + # only once the tool channel actually contains ATEM. Emitting it earlier + # -- when only the bare `to=<|message|>` header has arrived -- keeps + # the parser in the reasoning phase, so parse_delta never replaces this + # DeltaMessage with the tool parser's and the header is delivered to the + # client as visible content. + handoff = "" + if not self._tool_handoff_done: + remainder = self._tool_channel_remainder(current_text) + if _FUNCTION_CALLS_OPEN in remainder or "=0.11.1'` or build from " "https://github.com/cohere-ai/melody." ) from e diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index 4d3b9911355a..712869c2bfdf 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -26,7 +26,7 @@ render_for_completion, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest -from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve import create_error_response from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.inputs import ( EngineInput, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 41cfabbecf46..2f0af9e1b43b 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -406,6 +406,7 @@ def from_optional( skip_clone: bool = False, repetition_detection: RepetitionDetectionParams | None = None, logprob_token_ids: list[int] | None = None, + routed_experts_prompt_start: int = 0, ) -> "SamplingParams": if logit_bias is not None: # Fast path uses a dict comprehension; on failure we iterate once @@ -468,6 +469,7 @@ def from_optional( extra_args=extra_args, skip_clone=skip_clone, repetition_detection=repetition_detection, + routed_experts_prompt_start=routed_experts_prompt_start, ) def __post_init__(self) -> None: @@ -1027,6 +1029,15 @@ def _validate_structured_outputs( "structured_outputs.json_object must be True if set; omit " "structured_outputs to disable structured outputs" ) + # Reject a regex containing a NUL byte early, in every backend mode. A + # NUL is never meaningful in a regex pattern and is not handled by the + # regex-to-grammar conversion. Checked here, before backend selection, + # so it is a clean 400 rather than a silent fallback in the default + # "auto" mode. + if self.structured_outputs.regex and "\x00" in self.structured_outputs.regex: + raise VLLMValidationError( + "structured_outputs.regex must not contain a NUL character ('\\x00')" + ) from vllm.v1.structured_output.backend_guidance import ( has_guidance_unsupported_json_features, @@ -1081,7 +1092,7 @@ def _validate_structured_outputs( try: validate_xgrammar_grammar(self) self.structured_outputs._backend = "xgrammar" - except ValueError: + except VLLMValidationError: # The request either failed validation # or includes some jsonschema feature(s) that # are not supported in xgrammar. @@ -1092,7 +1103,12 @@ def _validate_structured_outputs( so_params = self.structured_outputs if not skip_guidance and so_params.json: if isinstance(so_params.json, str): - schema = json_mod.loads(so_params.json) + try: + schema = json_mod.loads(so_params.json) + except json_mod.JSONDecodeError as e: + raise VLLMValidationError( + "Invalid JSON grammar specification." + ) from e else: schema = so_params.json skip_guidance = has_guidance_unsupported_json_features(schema) diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index c8004cb57173..e1d0d965f056 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -278,9 +278,13 @@ def fused_recurrent_gated_delta_rule_packed_decode_kernel( BV: tl.constexpr, SOFTPLUS_THRESHOLD: tl.constexpr, USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + SPLIT_BATCH_HEAD_GRID: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hv = i_nh // HV, i_nh % HV + if SPLIT_BATCH_HEAD_GRID: + i_v, i_hv, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2) + else: + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) o_k = tl.arange(0, BK) @@ -446,7 +450,9 @@ def fused_recurrent_gated_delta_rule_packed_decode( stride_indices_seq = ssm_state_indices.stride(0) NV = triton.cdiv(V, BV) - grid = (NV, B * HV) + # CUDA limits grid Y/Z dimensions to 65535. + split_batch_head_grid = B * HV > 65535 + grid = (NV, HV, B) if split_batch_head_grid else (NV, B * HV) fused_recurrent_gated_delta_rule_packed_decode_kernel[grid]( mixed_qkv=mixed_qkv, a=a, @@ -472,6 +478,7 @@ def fused_recurrent_gated_delta_rule_packed_decode( BV=BV, SOFTPLUS_THRESHOLD=20.0, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + SPLIT_BATCH_HEAD_GRID=split_batch_head_grid, num_warps=num_warps, num_stages=num_stages, ) diff --git a/vllm/third_party/flash_linear_attention/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py index 810d32c18b85..cc6f54d8af0b 100644 --- a/vllm/third_party/flash_linear_attention/ops/index.py +++ b/vllm/third_party/flash_linear_attention/ops/index.py @@ -9,6 +9,7 @@ # ruff: noqa: E501 import torch +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.triton_utils import triton from .utils import tensor_cache @@ -21,17 +22,18 @@ def prepare_lens(cu_seqlens: torch.Tensor) -> torch.Tensor: @tensor_cache def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor: - indices = torch.cat( - [ - torch.arange(n) - for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() - ] + # This will be fixed by https://github.com/vllm-project/vllm/pull/51540. + with gpu_sync_allowed(): + 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 torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) @tensor_cache def prepare_chunk_offsets(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor: return torch.cat( - [cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)] + [cu_seqlens.new_zeros(1), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)] ).cumsum(-1) diff --git a/vllm/tilelang_utils/__init__.py b/vllm/tilelang_utils/__init__.py new file mode 100644 index 000000000000..803409b1ec65 --- /dev/null +++ b/vllm/tilelang_utils/__init__.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import functools +from collections.abc import Callable +from functools import cache +from typing import TYPE_CHECKING, Any + +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_tilelang + +if TYPE_CHECKING or current_platform.is_cuda(): + if not has_tilelang(): + raise ImportError( + "tilelang is required for mhc but is not installed. Install it with " + "`pip install tilelang`." + ) + import tilelang + import tilelang.language as T +else: + tilelang = None # type: ignore[assignment] + T = None # type: ignore[assignment] + + +def _ensure_tilelang_imported() -> None: + """Bind the `tilelang` and `T` module globals, importing them if needed. + + On ROCm, this runs on the first kernel call instead of at import time. + + Raises: + ImportError: If TileLang is not installed. + """ + global T, tilelang + + if tilelang is not None: + return + if not has_tilelang(): + raise ImportError( + "tilelang is required for mhc but is not installed. Install it with " + "`pip install tilelang`." + ) + import tilelang as tilelang_module + import tilelang.language as tilelang_language + + tilelang = tilelang_module + T = tilelang_language + + +@cache +def _get_pass_configs() -> dict[Any, Any]: + _ensure_tilelang_imported() + pass_configs: dict[Any, Any] = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + } + if current_platform.is_cuda(): + pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10 + return pass_configs + + +def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]: + """Apply `tilelang.jit`, deferring until first call on ROCm. + + ROCm defers JIT decoration so importing the caller's module does not + require TileLang immediately. CUDA keeps the eager decoration behavior. + + The kernel body parsed by TileLang references `T` as an unqualified + global, so on the deferred ROCm path this rebinds `T`/`tilelang` in the + decorated function's own module globals once they become available. + """ + if not current_platform.is_rocm(): + _ensure_tilelang_imported() + return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function) + + compiled_kernel: Callable[..., Any] | None = None + + @functools.wraps(kernel_function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal compiled_kernel + if compiled_kernel is None: + _ensure_tilelang_imported() + kernel_function.__globals__["tilelang"] = tilelang + kernel_function.__globals__["T"] = T + compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())( + kernel_function + ) + return compiled_kernel(*args, **kwargs) + + return wrapper diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 535e8e27816f..9c218a5eebdc 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -154,6 +154,10 @@ "olmo3_tool_parser", "Olmo3PythonicToolParser", ), + "muse_glimmer": ( + "muse_glimmer_tool_parser", + "MuseGlimmerToolParser", + ), "openai": ( "gptoss_tool_parser", "GptOssToolParser", diff --git a/vllm/tool_parsers/cohere_command_tool_parser.py b/vllm/tool_parsers/cohere_command_tool_parser.py index 6ce753b993c5..558d71ccc5b7 100644 --- a/vllm/tool_parsers/cohere_command_tool_parser.py +++ b/vllm/tool_parsers/cohere_command_tool_parser.py @@ -9,7 +9,7 @@ raise ImportError( "The Cohere tool parser requires the `cohere_melody` " "package, which is not installed. Install it with:\n" - " pip install cohere_melody" + " pip install 'cohere-melody>=0.11.1'" ) from e from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -130,8 +130,8 @@ def __init__( ): super().__init__( tokenizer, - streaming_opts=PyFilterOptions().cmd3(), - unary_opts=PyFilterOptions().cmd3(), + streaming_opts=PyFilterOptions().cmd3().start_in_answer(), + unary_opts=PyFilterOptions().cmd3().start_in_answer(), ) @@ -143,6 +143,6 @@ def __init__( ): super().__init__( tokenizer, - streaming_opts=PyFilterOptions().cmd4(), - unary_opts=PyFilterOptions().cmd4(), + streaming_opts=PyFilterOptions().cmd4().start_in_answer(), + unary_opts=PyFilterOptions().cmd4().start_in_answer(), ) diff --git a/vllm/tool_parsers/muse_glimmer_tool_parser.py b/vllm/tool_parsers/muse_glimmer_tool_parser.py new file mode 100644 index 000000000000..3a3842b8a5dd --- /dev/null +++ b/vllm/tool_parsers/muse_glimmer_tool_parser.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ATEM tool-call parser for MuseGlimmer. + +Faithful port of the MuseGlimmer ``response_schema`` tool-call contract from the +HuggingFace MuseGlimmer export (``convert_muse_glimmer_weights_to_hf.py``: +``MUSE_GLIMMER_RESPONSE_SCHEMA``). + +MuseGlimmer emits tool calls in an XML-ish ATEM format inside channel-scoped messages: + + <|start|>assistant to=self<|message|>...reasoning...<|eom|> + <|start|>assistant to=.<|message|> + + + value + + <|eom|> # non-final call + <|start|>assistant to=user<|message|>...final answer...<|eot|> + +Channel scoping is essential: an ```` echoed inside a ``to=self`` +reasoning block or a ``to=user`` final answer must NOT be parsed as a real tool +call. + +Rather than *subtracting* reasoning/answer spans with regex substitutions (the +approach the HF ``response_schema`` uses, which is safe only on a complete, +well-formed turn), this parser *segments* the output into messages and then +*selects* the tool-channel bodies. On a complete turn the two are equivalent; +on a truncated or damaged turn, subtraction can delete a valid tool call +(an unterminated ``to=self`` block makes the non-greedy strip run to the next +``<|eom|>``, which belongs to the tool-call message) whereas selection cannot. + +Usage: ``--enable-auto-tool-choice --tool-call-parser muse_glimmer`` +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Sequence + +import regex as re +from openai.types.responses import ToolChoiceFunction +from transformers import PreTrainedTokenizerBase + +from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + ExtractedToolCallInformation, + FunctionCall, + ToolCall, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.logger import init_logger +from vllm.tool_parsers.abstract_tool_parser import ( + Tool, + ToolParser, +) + +logger = init_logger(__name__) + +# --- Message framing ------------------------------------------------------- +# An assistant message header. All three parts are optional except the +# <|message|> terminator: +# "<|start|>assistant to=get_weather<|message|>" -- after an <|eom|> boundary +# " to=self<|message|>" -- first message of a turn +# (the prompt already ended +# with "<|start|>assistant") +# "<|message|>" -- bare recipient (public CoT +# / untagged content) +_MSG_HEADER_RE = re.compile( + r"(?:<\|start\|>\s*assistant)?[^\S\n]*(?:to=(?P[A-Za-z0-9_.\-]+))?<\|message\|>" +) +_MSG_END_RE = re.compile(r"<\|eom\|>|<\|eot\|>") + +# Recipients whose bodies are NOT tool calls. +_REASONING_RECIPIENT = "self" +_USER_RECIPIENT = "user" + +# Structural markers that must never reach the client. A streamed body is held +# back by up to len(marker)-1 characters so a marker split across two chunks is +# not emitted as content. +_STRUCTURAL_MARKERS = ("<|eom|>", "<|eot|>", "<|start|>", "<|message|>") +_MAX_MARKER_LEN = max(len(m) for m in _STRUCTURAL_MARKERS) +# A trailing " to=NAME" that could still grow into a bare message header (the +# first message of a turn has no <|start|> prefix -- the prompt ends with +# "<|start|>assistant", so the model's first emitted text is " to=self<|message|>"). +_OPEN_TAIL_TO_RE = re.compile(r"[^\S\n]+to=[A-Za-z0-9_.\-]*$") + +# --- Tool-call extraction (unchanged from MUSE_GLIMMER_RESPONSE_SCHEMA) ------------- +_INVOKE_RE = re.compile(r"()", re.DOTALL) +_NAME_RE = re.compile(r']*?\bname="([^"]+)"') +_PARAM_RE = re.compile( + r']*?\bname="(?P[^"]+)"[^>]*?>(?P.*?)', + re.DOTALL, +) +_FUNCTION_CALLS_OPEN = "" + + +def _decode_value(raw: str): + """JSON-decode a parameter value when possible, else keep the raw string. + + Mirrors the schema's ``x-parser: json`` with ``allow_non_json: True``. + """ + try: + return json.loads(raw) + except (json.JSONDecodeError, ValueError): + return raw + + +def _iter_messages(text: str) -> Iterator[tuple[str | None, str, bool]]: + """Segment *text* into assistant messages. + + Yields ``(recipient, body, closed)`` per message, where ``recipient`` is + ``None`` for a bare ``<|message|>`` header and ``closed`` is False for a + message that has not (yet) seen ``<|eom|>`` / ``<|eot|>``. + + A message is also terminated by the start of the NEXT header. Without that, + a reasoning block whose ``<|eom|>`` is missing (truncation, or a chunk + dropped at the reasoning -> tool transition) would absorb the tool-call + message that follows it and the call would be lost -- the same defect the + subtractive regexes have. + """ + pos = 0 + while pos < len(text): + header = _MSG_HEADER_RE.search(text, pos) + if header is None: + return + body_start = header.end() + end = _MSG_END_RE.search(text, body_start) + nxt = _MSG_HEADER_RE.search(text, body_start) + body_end = end.start() if end is not None else len(text) + closed = end is not None + if nxt is not None and nxt.start() < body_end: + body_end = nxt.start() + closed = False + next_pos = nxt.start() + else: + next_pos = end.end() if end is not None else len(text) + body = text[body_start:body_end] + # A body can never legitimately contain <|start|>. Seeing one means the + # next header is only partially generated (its <|message|> has not + # arrived), so the regex above could not recognise it yet. Cut there, + # otherwise the streamed body would grow to include the next header and + # then shrink back once it completes. + start_tok = body.find("<|start|>") + if start_tok != -1: + body = body[:start_tok] + closed = False + yield header.group("rcpt"), body, closed + pos = next_pos + + +def _trailing_partial_marker_len(text: str) -> int: + """Length of the longest suffix of *text* that prefixes a structural marker.""" + max_overlap = min(len(text), _MAX_MARKER_LEN - 1) + for overlap in range(max_overlap, 0, -1): + suffix = text[-overlap:] + if any(marker.startswith(suffix) for marker in _STRUCTURAL_MARKERS): + return overlap + return 0 + + +def _safe_open_body(body: str) -> str: + """Trim the tail of a still-growing body to what is safe to emit now. + + Holds back anything that could still turn out to be structural, so the + emitted prefix only ever grows. Chunks under speculative decoding are large + enough that markers routinely straddle them. + """ + tail_to = _OPEN_TAIL_TO_RE.search(body) + if tail_to is not None: + return body[: tail_to.start()] + partial = _trailing_partial_marker_len(body) + return body[: len(body) - partial] if partial else body + + +class MuseGlimmerToolParser(ToolParser): + # MuseGlimmer emits ATEM markup around tool-call arguments. The generic + # named/required tool_choice path in vllm/parser/abstract_parser.py assigns + # the raw model text straight to FunctionCall.arguments, which leaks that + # framing and yields invalid JSON; the "required" branch then silently + # swallows the ValidationError and returns tool_calls=null. Opting out + # routes named/required through extract_tool_calls / + # extract_tool_calls_streaming -- the same path "auto" already uses. + supports_required_and_named = False + + def __init__( + self, + tokenizer: PreTrainedTokenizerBase, + tools: list[Tool] | None = None, + ) -> None: + super().__init__(tokenizer, tools) + # Streaming cursors. vLLM constructs one ToolParser per request, so + # instance state is per-stream. + self._streamed_content_len: int = 0 + self._streamed_reasoning_len: int = 0 + self._emitted_tool_calls: int = 0 + + def adjust_request( + self, request: ChatCompletionRequest | ResponsesRequest + ) -> ChatCompletionRequest | ResponsesRequest: + """Force special tokens through, and keep JSON guided decoding off. + + ``skip_special_tokens`` defaults to True on both ChatCompletionRequest + and ResponsesRequest. Every rule in this parser keys off ``<|message|>`` + / ``<|eom|>`` / ``<|eot|>`` / ``<|start|>``, so with the default the + channel framing is stripped before we see it: message segmentation + finds nothing, reasoning-channel invokes are indistinguishable from real + ones, and the raw ATEM markup falls through to the client as content. + Set it unconditionally and FIRST -- not only for tools requests, and not + relying on the reasoning parser to have set it. + + For required/named ``tool_choice`` the base hook installs a JSON schema + constraint (ToolParser.adjust_request -> get_json_schema_from_tools). + MuseGlimmer emits ATEM XML, so under that constraint it writes JSON *inside* + the tool channel -- ``[{"name": ...`` -- with no + ```` for this parser to find. Skip the base hook so those + choices decode natively, the same as "auto". + """ + request.skip_special_tokens = False + + tool_choice = getattr(request, "tool_choice", None) + if request.tools and ( + tool_choice == "required" + or isinstance( + tool_choice, (ChatCompletionNamedToolChoiceParam, ToolChoiceFunction) + ) + ): + return request + return super().adjust_request(request) + + # ---------------- channel selection ---------------- + + @classmethod + def _tool_channel_text(cls, text: str) -> str: + """Concatenate the bodies of messages addressed to a tool. + + Falls back to the whole text when no message header is present at all -- + that means the framing never reached us (``skip_special_tokens`` was on, + or the chunk carrying the header was dropped upstream), and scanning + everything is strictly better than returning nothing. + """ + bodies = [ + body + for rcpt, body, _closed in _iter_messages(text) + if rcpt is not None + and rcpt != _REASONING_RECIPIENT + and rcpt != _USER_RECIPIENT + ] + if bodies: + return "\n".join(bodies) + if _MSG_HEADER_RE.search(text) is None and ( + _FUNCTION_CALLS_OPEN in text or " tuple[str, str, bool, bool]: + """Return ``(content, reasoning, content_open, reasoning_open)``. + + The ``*_open`` flags say whether that channel's LAST message is still + being generated; only then must the caller hold back a partial + structural marker. Tracking them per channel matters: a closed + reasoning block whose text happens to end in ``<`` would otherwise stay + permanently truncated while a later content message is open. + """ + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + content_open = False + reasoning_open = False + for rcpt, body, closed in _iter_messages(text): + if rcpt == _REASONING_RECIPIENT: + reasoning_parts.append(body) + reasoning_open = not closed + elif rcpt is None or rcpt == _USER_RECIPIENT: + content_parts.append(body) + content_open = not closed + return ( + "".join(content_parts), + "".join(reasoning_parts), + content_open, + reasoning_open, + ) + + # ---------------- tool name binding ---------------- + + @staticmethod + def _registered_names(request: ChatCompletionRequest | None) -> set[str]: + """Names of the tools the client registered on this request.""" + names: set[str] = set() + tools = getattr(request, "tools", None) if request is not None else None + for t in tools or []: + fn = getattr(t, "function", None) or t + name = getattr(fn, "name", None) + if name is None and isinstance(fn, dict): + name = fn.get("name") + if name: + names.add(name) + return names + + @staticmethod + def _normalize_name(emitted: str, registered: set[str]) -> str: + """Map an emitted ATEM invoke name back to a registered tool name. + + When a client registers a BARE name (e.g. ``get_weather``) the shipped + chat template renders the valid recipient as ``"get_weather.*"``, and + the model duly emits ``get_weather.get_weather``. Collapsing that + doubled form is safe: head and tail are identical and the collapsed + name is registered. + + Anything else is passed through unchanged. Matching on the trailing + segment alone is NOT safe -- an emitted ``weather.get`` against a + registered ``{calendar.get}`` has a unique leaf match and would silently + dispatch the wrong tool. + """ + if not registered or emitted in registered: + return emitted + head, sep, tail = emitted.partition(".") + if sep and head == tail and head in registered: + return head + logger.warning( + "MuseGlimmer: emitted tool name %r does not match any registered tool; " + "passing through unchanged.", + emitted, + ) + return emitted + + @classmethod + def _parse_tool_calls( + cls, text: str, registered: set[str] | None = None + ) -> list[ToolCall]: + registered = registered or set() + scoped = cls._tool_channel_text(text) + tool_calls: list[ToolCall] = [] + for invoke in _INVOKE_RE.findall(scoped): + name_m = _NAME_RE.search(invoke) + if not name_m: + continue + name = cls._normalize_name(name_m.group(1), registered) + args: dict = {} + for pm in _PARAM_RE.finditer(invoke): + args[pm.group("key")] = _decode_value(pm.group("value")) + tool_calls.append( + ToolCall( + function=FunctionCall( + name=name, + arguments=json.dumps(args, ensure_ascii=False), + ) + ) + ) + return tool_calls + + @classmethod + def _extract_content(cls, text: str) -> str | None: + """Return the user-facing body, or the raw text when unframed.""" + content, reasoning, _c_open, _r_open = cls._visible_channels(text) + if content: + return content + # No framing at all -> the whole thing is plain content. + if not reasoning and _MSG_HEADER_RE.search(text) is None: + return text or None + return None + + # ---------------- non-streaming ---------------- + + def extract_tool_calls( + self, model_output: str, request: ChatCompletionRequest + ) -> ExtractedToolCallInformation: + if ( + _FUNCTION_CALLS_OPEN not in model_output + and "... parsed -- typically a truncated + # call (finish_reason='length'/abort). Log it: silently + # returning "no tool call" here is indistinguishable from the + # model choosing not to call one, which makes this failure mode + # invisible in production. + logger.warning( + "MuseGlimmer: tool channel opened but no complete " + "parsed (truncated tool call?); returning content only." + ) + return ExtractedToolCallInformation( + tools_called=False, + tool_calls=[], + content=self._extract_content(model_output), + ) + return ExtractedToolCallInformation( + tools_called=True, + tool_calls=tool_calls, + content=self._extract_content(model_output), + ) + except Exception: + logger.exception("Error extracting MuseGlimmer ATEM tool calls.") + return ExtractedToolCallInformation( + tools_called=False, tool_calls=[], content=model_output + ) + + # ---------------- streaming ---------------- + + def extract_tool_calls_streaming( + self, + previous_text: str, + current_text: str, + delta_text: str, + previous_token_ids: Sequence[int], + current_token_ids: Sequence[int], + delta_token_ids: Sequence[int], + request: ChatCompletionRequest, + ) -> DeltaMessage | None: + """Incremental ATEM streaming for tool calls AND content. + + This parser owns every delta once reasoning has ended: in + vllm/parser/abstract_parser.py::parse_delta the "pass through as + content" fallback is guarded by ``not self._in_tool_call_phase(state)``, + and ``_in_tool_call_phase`` is simply ``tool_parser is not None and + state.reasoning_ended``. So with a tool parser loaded that fallback is + dead code, and returning None here DISCARDS the delta. Anything we do + not emit -- including the ``to=user`` final answer -- never reaches the + client. Hence content is emitted here, not left to the reasoning parser. + + Tool calls are surfaced only when an ```` block becomes + complete: the XML is opaque until closed and MuseGlimmer parameters are not + incremental JSON, so there is nothing meaningful to stream before then. + """ + if not previous_text: + # First delta of the tool phase (parse_delta resets previous_text + # to "" when it hands the stream over). Reset the cursors. + self._streamed_content_len = 0 + self._streamed_reasoning_len = 0 + self._emitted_tool_calls = 0 + + try: + registered = self._registered_names(request) + calls = self._parse_tool_calls(current_text, registered) + content, reasoning, content_open, reasoning_open = self._visible_channels( + current_text + ) + + # Trim the tail of a channel that is still growing, so the emitted + # prefix never shrinks between deltas. + if content_open: + content = _safe_open_body(content) + if reasoning_open: + reasoning = _safe_open_body(reasoning) + + content_delta = content[self._streamed_content_len :] + reasoning_delta = reasoning[self._streamed_reasoning_len :] + + tool_deltas: list[DeltaToolCall] = [] + for i in range(self._emitted_tool_calls, len(calls)): + fn = calls[i].function + tool_deltas.append( + DeltaToolCall( + index=i, + type="function", + id=make_tool_call_id(), + function=DeltaFunctionCall( + name=fn.name, + arguments=fn.arguments, + ).model_dump(exclude_none=True), + ) + ) + + if not content_delta and not reasoning_delta and not tool_deltas: + return None + + self._streamed_content_len = len(content) + self._streamed_reasoning_len = len(reasoning) + self._emitted_tool_calls = len(calls) + + message = DeltaMessage() + if content_delta: + message.content = content_delta + if reasoning_delta: + message.reasoning = reasoning_delta + if tool_deltas: + message.tool_calls = tool_deltas + return message + except Exception: + logger.exception("Error extracting MuseGlimmer ATEM streaming tool calls.") + return None diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 8bbba8df24a6..87271c258825 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -102,6 +102,10 @@ def __getitem__(self, key): kimi_linear="KimiLinearConfig", kimi_vl="KimiVLConfig", kimi_k25="KimiK25Config", + muse_glimmer="MuseGlimmerConfig", + muse_glimmer_text="MuseGlimmerTextConfig", + muse_glimmer_vision="MuseGlimmerVisionConfig", + muse_glimmer_assistant="MuseGlimmerAssistantConfig", kimi_k3="KimiK3Config", RefinedWeb="RWConfig", # For tiiuae/falcon-40b(-instruct) RefinedWebModel="RWConfig", # For tiiuae/falcon-7b(-instruct) diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index a6d477831dcb..c38cc4b4b8cc 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -74,6 +74,10 @@ "KimiLinearConfig": "vllm.transformers_utils.configs.kimi_linear", "KimiVLConfig": "vllm.transformers_utils.configs.kimi_vl", "KimiK25Config": "vllm.transformers_utils.configs.kimi_k25", + "MuseGlimmerConfig": "vllm.transformers_utils.configs.muse_glimmer", + "MuseGlimmerTextConfig": "vllm.transformers_utils.configs.muse_glimmer", + "MuseGlimmerVisionConfig": "vllm.transformers_utils.configs.muse_glimmer", + "MuseGlimmerAssistantConfig": "vllm.transformers_utils.configs.muse_glimmer", "KimiK3Config": "vllm.transformers_utils.configs.kimi_k3", "KimiK3VisionConfig": "vllm.transformers_utils.configs.kimi_k3", "NemotronConfig": "vllm.transformers_utils.configs.nemotron", @@ -162,6 +166,10 @@ "KimiLinearConfig", "KimiVLConfig", "KimiK25Config", + "MuseGlimmerConfig", + "MuseGlimmerTextConfig", + "MuseGlimmerVisionConfig", + "MuseGlimmerAssistantConfig", "KimiK3Config", "KimiK3VisionConfig", "NemotronConfig", diff --git a/vllm/transformers_utils/configs/muse_glimmer.py b/vllm/transformers_utils/configs/muse_glimmer.py new file mode 100644 index 000000000000..9022c84cf0cc --- /dev/null +++ b/vllm/transformers_utils/configs/muse_glimmer.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MuseGlimmer model configuration for vLLM. + +Native vLLM copy of the MuseGlimmer HuggingFace configs +(``configuration_muse_glimmer.py``). MuseGlimmer's +``model_type`` is not yet registered in released transformers, so vLLM ships +this config so partners can serve MuseGlimmer checkpoints *without* trust_remote_code. + +The field set and defaults are kept byte-for-byte in sync with the HF reference +so a checkpoint's ``config.json`` deserializes identically here. Both the text +and vision configs are consumed by the native multimodal serving path. +""" + +from __future__ import annotations + +from transformers import Qwen3Config +from transformers.configuration_utils import PretrainedConfig + + +def _default_no_rope_layers(num_hidden_layers: int) -> list[int]: + # iRoPE mask: NoPE every 4 layers, counted backward from the last layer. + stride = 4 + return [ + 0 if (num_hidden_layers - 1 - i) % stride == 0 else 1 + for i in range(num_hidden_layers) + ] + + +class MuseGlimmerTextConfig(PretrainedConfig): + model_type = "muse_glimmer_text" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 202_048, + hidden_size: int = 6656, + intermediate_size: int = 19968, + num_hidden_layers: int = 52, + num_attention_heads: int = 32, + num_key_value_heads: int = 2, + head_dim: int = 128, + hidden_activation: str = "silu", + max_position_embeddings: int = 16_384, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-5, + use_cache: bool = True, + pad_token_id: int | None = None, + eos_token_id: int | list[int] | None = 200_001, + bos_token_id: int | None = 200_000, + tie_word_embeddings: bool = False, + rope_parameters: dict | None = None, + rope_theta: float | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + query_pre_attn_scalar: int = 256, + sliding_window: int | None = 2048, + layer_types: list[str] | None = None, + final_logit_softcapping: float | None = 20.0, + attn_logit_softcapping: float | None = None, + use_bidirectional_attention: bool | None = None, + # MuseGlimmer-specific + qk_scale_factor: float = 43.7840518911, + use_qk_norm: bool = True, + use_attn_output_gate: bool = True, + output_multiplier: float = 0.19611613513818404, + normalize_tok_embeddings: bool = True, + post_norm_eps: float = 1e-8, + no_rope_layers: list[int] | None = None, + **kwargs, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_activation = hidden_activation + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.query_pre_attn_scalar = query_pre_attn_scalar + self.sliding_window = sliding_window + self.final_logit_softcapping = final_logit_softcapping + self.attn_logit_softcapping = attn_logit_softcapping + self.use_bidirectional_attention = use_bidirectional_attention + + # RoPE: accept either an explicit rope_parameters dict (HF 5.x) or a + # bare rope_theta; normalize to rope_parameters for vLLM's get_rope. + if rope_parameters is None: + theta = rope_theta if rope_theta is not None else 500_000.0 + rope_parameters = {"rope_type": "default", "rope_theta": theta} + self.rope_parameters = rope_parameters + # vLLM reads rope_theta off the config in some codepaths. + self.rope_theta = rope_parameters.get("rope_theta", 500_000.0) + + # MuseGlimmer-specific fields + self.qk_scale_factor = qk_scale_factor + self.use_qk_norm = use_qk_norm + self.use_attn_output_gate = use_attn_output_gate + self.output_multiplier = output_multiplier + self.normalize_tok_embeddings = normalize_tok_embeddings + self.post_norm_eps = post_norm_eps + + self.no_rope_layers = ( + no_rope_layers + if no_rope_layers is not None + else _default_no_rope_layers(num_hidden_layers) + ) + if layer_types is None: + layer_types = [ + "full_attention" if self.no_rope_layers[i] == 0 else "sliding_attention" + for i in range(num_hidden_layers) + ] + self.layer_types = layer_types + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class MuseGlimmerVisionConfig(PretrainedConfig): + model_type = "muse_glimmer_vision" + + def __init__( + self, + patch_size: int = 14, + pos_emb_height: int = 32, + pos_emb_width: int = 32, + num_attention_heads: int = 16, + num_hidden_layers: int = 50, + hidden_size: int = 1536, + intermediate_size: int = 8960, + hidden_act: str = "gelu", + merge_kernel_size: int = 2, + rope_parameters: dict | None = None, + max_position_embeddings: int = 32 * 32, + output_dim: int = 6144, + patch_temporal: int = 2, + adapter_dim: int = 4096, + layer_norm_eps: float = 1e-5, + layer_types: list[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.patch_size = patch_size + self.pos_emb_height = pos_emb_height + self.pos_emb_width = pos_emb_width + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.merge_kernel_size = merge_kernel_size + self.rope_parameters = rope_parameters + self.max_position_embeddings = max_position_embeddings + self.output_dim = output_dim + self.patch_temporal = patch_temporal + self.adapter_dim = adapter_dim + self.layer_norm_eps = layer_norm_eps + if layer_types is None: + stride = 4 + layer_types = [ + "full_attention" + if (i + 1) % stride == 0 or i == num_hidden_layers - 1 + else "sliding_attention" + for i in range(num_hidden_layers) + ] + self.layer_types = layer_types + + +class MuseGlimmerConfig(PretrainedConfig): + model_type = "muse_glimmer" + sub_configs = { + "text_config": MuseGlimmerTextConfig, + "vision_config": MuseGlimmerVisionConfig, + } + + # --- Flat (legacy-converter) -> canonical normalization ------------------- + # MuseGlimmer checkpoints exist in two config layouts in the wild: + # + # * CANONICAL (current HF converter, transformers 5.15): nested + # ``text_config`` / ``vision_config`` sub-dicts with canonical field + # names (``hidden_activation``, ``final_logit_softcapping``, + # ``rope_parameters``, ``vision_config.hidden_size`` ...). + # + # * FLAT (older converter, e.g. Ruan's ``rl_v1/hf``, transformers 5.9): + # every field is a top-level key, with different names + # (``hidden_act``, ``output_soft_cap_temp``, ``rope_theta``, + # ``vision_latent_dim`` ...) and NO ``text_config`` nesting. + # + # Without normalization a flat config silently deserializes to an + # ALL-DEFAULT text config (every checkpoint value ignored) — a dangerous + # correctness bug: a checkpoint whose arch differs from the defaults would + # load into a wrong-shaped model with no error. So when we detect a flat + # config we hoist the flat fields into ``text_config`` / ``vision_config`` + # and rename them to canonical names. + + # flat text field name -> canonical MuseGlimmerTextConfig field name + _FLAT_TEXT_RENAMES = { + "hidden_act": "hidden_activation", + "output_soft_cap_temp": "final_logit_softcapping", + } + # canonical MuseGlimmerTextConfig constructor params that may appear flat + _FLAT_TEXT_KEYS = frozenset( + { + "vocab_size", + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_activation", + "hidden_act", + "max_position_embeddings", + "initializer_range", + "rms_norm_eps", + "use_cache", + "tie_word_embeddings", + "rope_parameters", + "rope_theta", + "attention_bias", + "attention_dropout", + "query_pre_attn_scalar", + "sliding_window", + "layer_types", + "final_logit_softcapping", + "output_soft_cap_temp", + "attn_logit_softcapping", + "use_bidirectional_attention", + "qk_scale_factor", + "use_qk_norm", + "use_attn_output_gate", + "output_multiplier", + "normalize_tok_embeddings", + "post_norm_eps", + "no_rope_layers", + } + ) + # flat vision field name -> canonical MuseGlimmerVisionConfig field name + _FLAT_VISION_RENAMES = { + "vision_latent_dim": "hidden_size", + "vision_heads": "num_attention_heads", + "vision_layers": "num_hidden_layers", + "vision_output_dim": "output_dim", + "vision_patch_size": "patch_size", + "vision_patch_temporal": "patch_temporal", + "vision_adapter_dim": "adapter_dim", + "vision_pos_emb_grid_h": "pos_emb_height", + "vision_pos_emb_grid_w": "pos_emb_width", + "vision_downsample_factor": "merge_kernel_size", + } + + @classmethod + def _looks_flat(cls, kwargs: dict) -> bool: + # Flat if there is no explicit text_config but there ARE text-level + # fields at the top level (e.g. hidden_size / num_hidden_layers). + if kwargs.get("text_config") is not None: + return False + return any(k in kwargs for k in ("hidden_size", "num_hidden_layers")) + + @classmethod + def _normalize_flat(cls, kwargs: dict) -> dict: + kwargs = dict(kwargs) + text: dict = {} + for key in list(kwargs.keys()): + if key in cls._FLAT_TEXT_KEYS: + canon = cls._FLAT_TEXT_RENAMES.get(key, key) + text.setdefault(canon, kwargs.pop(key)) + + vision_mlp_ratio = kwargs.pop("vision_mlp_ratio", None) + sparse_factor = kwargs.pop("vision_sparse_attention_factor", None) + vision: dict = {} + for key in list(kwargs.keys()): + if key in cls._FLAT_VISION_RENAMES: + vision.setdefault(cls._FLAT_VISION_RENAMES[key], kwargs.pop(key)) + + if vision_mlp_ratio is not None: + hidden_size = vision.get("hidden_size", 1536) + vision["intermediate_size"] = int(vision_mlp_ratio * hidden_size) + if sparse_factor is not None: + sparse_factor = int(sparse_factor) + if sparse_factor <= 0: + raise ValueError("vision_sparse_attention_factor must be positive") + num_layers = vision.get("num_hidden_layers", 50) + vision["layer_types"] = [ + "full_attention" + if (layer_idx + 1) % sparse_factor == 0 or layer_idx == num_layers - 1 + else "sliding_attention" + for layer_idx in range(num_layers) + ] + + if text: + kwargs["text_config"] = text + if vision: + kwargs["vision_config"] = vision + return kwargs + + def __init__( + self, + text_config: dict | MuseGlimmerTextConfig | None = None, + vision_config: dict | MuseGlimmerVisionConfig | None = None, + image_token_id: int = 200092, + video_token_id: int = 200091, + **kwargs, + ) -> None: + # Detect + fold a flat config into nested text/vision before building + # the sub-configs. (image/video token ids are handled below; a flat + # ``patch_token_id`` alias is mapped too.) + if text_config is None and self._looks_flat(kwargs): + kwargs = self._normalize_flat(kwargs) + text_config = kwargs.pop("text_config", None) + vision_config = kwargs.pop("vision_config", vision_config) + # flat image-token alias + if "patch_token_id" in kwargs and "image_token_id" not in kwargs: + image_token_id = kwargs.pop("patch_token_id") + + if text_config is None: + self.text_config = MuseGlimmerTextConfig() + elif isinstance(text_config, dict): + self.text_config = MuseGlimmerTextConfig(**text_config) + else: + self.text_config = text_config + + if vision_config is None: + self.vision_config = MuseGlimmerVisionConfig() + elif isinstance(vision_config, dict): + self.vision_config = MuseGlimmerVisionConfig(**vision_config) + else: + self.vision_config = vision_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.patch_token_id = image_token_id + super().__init__(**kwargs) + + +class MuseGlimmerAssistantConfig(Qwen3Config): + """Config for the Muse Glimmer DFlash draft head. + + The head is Qwen3-shaped and runs on vLLM's generic ``qwen3_dflash`` + implementation, so this derives from ``Qwen3Config``. It cannot BE + ``Qwen3Config``, because two of the checkpoint's values do not survive that + class: + + * ``sliding_window`` is gated behind ``use_sliding_window``, which defaults + to False -- ``Qwen3Config(sliding_window=2048).sliding_window`` is None. + The checkpoint declares ``sliding_window: 2048`` and five + ``sliding_attention`` layers, so the window silently disappears and the + DFlash path then raises "sliding attention requires a window size". + * ``vocab_size`` is absent from the checkpoint, so Qwen3's default of + 151936 applies instead of Muse Glimmer's 202048. That one is *silent*: it + builds an all-zero ``draft_id_to_target_id`` remap, and also puts + pad/bos/eos/mask token ids out of range. + + Both defaults are set here so an unmodified checkpoint loads correctly. + """ + + model_type = "muse_glimmer_assistant" + + def __init__( + self, + vocab_size: int = 202048, + use_sliding_window: bool = True, + dflash_config: dict | None = None, + **kwargs, + ) -> None: + super().__init__( + vocab_size=vocab_size, use_sliding_window=use_sliding_window, **kwargs + ) + # muse_glimmer_assistant uses non-causal attention. + self.dflash_config = {"causal": False} | (dflash_config or {}) + + +__all__ = [ + "MuseGlimmerTextConfig", + "MuseGlimmerVisionConfig", + "MuseGlimmerConfig", + "MuseGlimmerAssistantConfig", +] diff --git a/vllm/transformers_utils/processors/cosmos3_edge.py b/vllm/transformers_utils/processors/cosmos3_edge.py index 2644c6b17e59..95ca2cccd531 100644 --- a/vllm/transformers_utils/processors/cosmos3_edge.py +++ b/vllm/transformers_utils/processors/cosmos3_edge.py @@ -8,12 +8,16 @@ from torchvision.transforms import InterpolationMode from transformers import AutoTokenizer from transformers.feature_extraction_utils import BatchFeature -from transformers.image_utils import ChannelDimension, PILImageResampling, SizeDict +from transformers.image_utils import ( + ChannelDimension, + PILImageResampling, + SizeDict, + get_image_size, +) from transformers.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor from transformers.models.qwen3_vl.video_processing_qwen3_vl import ( Qwen3VLVideoProcessor, Qwen3VLVideoProcessorInitKwargs, - get_image_size, smart_resize, ) from transformers.models.siglip2.image_processing_siglip2 import ( @@ -22,6 +26,7 @@ ) from transformers.processing_utils import Unpack from transformers.utils import TensorType +from transformers.video_processing_utils import BaseVideoProcessor from transformers.video_utils import group_videos_by_shape, reorder_videos @@ -225,8 +230,11 @@ def _preprocess( max_pixels=size.longest_edge, ) stacked_videos = stacked_videos.view(B * T, C, H, W) - stacked_videos = self.resize( - stacked_videos, + # The target size is already computed, so bypass Qwen3-VL's + # dynamic resize, which expects an unflattened video tensor. + stacked_videos = BaseVideoProcessor.resize( + self, + image=stacked_videos, size=SizeDict(height=resized_height, width=resized_width), resample=resample, ) diff --git a/vllm/transformers_utils/processors/muse_glimmer.py b/vllm/transformers_utils/processors/muse_glimmer.py new file mode 100644 index 000000000000..c1a5326ea1f1 --- /dev/null +++ b/vllm/transformers_utils/processors/muse_glimmer.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HuggingFace processor for the MuseGlimmer multimodal model (text + image + video). + +This packages MuseGlimmer's vision preprocessing and token-span expansion the standard +HF way, so you can do: + + processor = AutoProcessor.from_pretrained(hf_dir, trust_remote_code=True) + messages = [{"role": "user", "content": [ + {"type": "image"}, + {"type": "text", "text": "Describe this image."}, + ]}] + text = processor.apply_chat_template(messages, add_generation_prompt=True) + inputs = processor(text=text, images=[pil_image], return_tensors="pt") + out = model.generate(**inputs, max_new_tokens=256) + +It reproduces the exact token layout the model expects: + + image -> <|image_start|> + <|patch|> * N + <|image_end|> + video -> <|vid_start|> + ( "Time: X.Xs" + <|video|> * P [+ <|vid_frame_separator|>] )* + + <|vid_end|> + +where N = floor(H/(patch*ds)) * floor(W/(patch*ds)) (<= 4096) for images and +P = the same grid count per temporal frame-group (<= 144) for video, with the +target (H, W) chosen by MuseGlimmerImageProcessor.compute_image_size / +MuseGlimmerVideoProcessor.compute_video_frame_size (the same grid logic as +modeling_muse_glimmer.MuseGlimmerVisionEncoder). + +The chat template emits one sentinel per media item -- ``<|image|>`` for images, +``<|video|>`` for videos -- which ``MuseGlimmerProcessor.__call__`` expands into the +spans above using each media item's computed token count. Vision features are +returned in ``pixel_values`` (a list of per-image / per-frame-group tensors) in +the order the sentinels appear, matching how ``MuseGlimmerModel`` merges them at the +``<|patch|>`` / ``<|video|>`` positions. +""" + +from __future__ import annotations + +import itertools +import math +from pathlib import Path + +import torch +from PIL import Image +from torchvision import transforms as T +from transformers.feature_extraction_utils import BatchFeature +from transformers.image_processing_utils import BaseImageProcessor +from transformers.processing_utils import ProcessorMixin +from transformers.video_processing_utils import BaseVideoProcessor + +# Single-token sentinels emitted by the chat template, one per media item. +IMAGE_SENTINEL = "<|image|>" # id 200090; absent from the image/video spans +VIDEO_SENTINEL = "<|video|>" # id 200091; also the per-frame token (see __call__) + +# Jinja chat template (multimodal superset of convert_muse_glimmer_to_hf.CHAT_TEMPLATE). +# Message content may be a plain string OR a list of parts +# ({"type": "text"|"image"|"video", ...}); image/video parts render as a single +# sentinel that __call__ expands. +MUSE_GLIMMER_MM_CHAT_TEMPLATE = ( + "{{- bos_token -}}" + "{%- macro render_parts(content) -%}" + "{%- if content is string -%}{{- content -}}" + "{%- else -%}" + "{%- for part in content -%}" + "{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}" + "{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}" + "{%- elif part['type'] == 'text' -%}{{- part['text'] -}}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- endif -%}" + "{%- endmacro -%}" + # At inference (add_generation_prompt), inject a default system message when + # the caller supplied none. Full-transcript rendering + # (add_generation_prompt=False) is left unchanged. + "{%- set ns = namespace(has_system=false) -%}" + "{%- for m in messages -%}" + "{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt and not ns.has_system -%}" + "{{- '<|start|>system<|message|>You are a helpful assistant.<|eot|>' -}}" + "{%- endif -%}" + "{%- for message in messages -%}" + "{%- set role = message['role'] -%}" + "{%- if role == 'assistant' -%}" + "{%- set recipient = message.get('recipient') -%}" + "{%- set end_turn = message.get('end_turn') -%}" + "{%- if end_turn is none -%}" + "{%- set end_turn = not (recipient and recipient != 'user') -%}" + "{%- endif -%}" + "{{- '<|start|>assistant' -}}" + "{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}" + "{{- '<|message|>' -}}{{- render_parts(message['content']) -}}" + "{{- ('<|eot|>' if end_turn else '<|eom|>') -}}" + "{%- elif role == 'tool' -%}" + "{%- set name = message.get('name', '') -%}" + # Tool content is emitted as-is (string body or interleaved image/text parts). + # Any wrapper is baked into the SFT data content; the + # tokenizer does not add it, so the template must not either. + "{{- '<|start|>tool ' + name + '<|message|>' -}}" + "{{- render_parts(message['content']) -}}" + "{{- '<|eot|>' -}}" + "{%- else -%}" + "{%- set header = role -%}" + "{%- if message.get('name') -%}" + "{%- set header = role + ' ' + message['name'] -%}" + "{%- endif -%}" + "{{- '<|start|>' + header + '<|message|>' -}}" + "{{- render_parts(message['content']) -}}" + "{{- '<|eot|>' -}}" + "{%- endif -%}" + "{%- endfor -%}" + "{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}" +) + + +def _grid_size( + img_w: int, img_h: int, patch_hw: int, max_tokens: int +) -> tuple[int, int, int]: + """Pick the integer (H, W) grid closest to the aspect ratio under the token cap. + + Replicates MuseGlimmerVisionEncoder._compute_grid_size + (modeling_muse_glimmer.py) so the processor needs no torch model import. + Returns (target_h, target_w, n_tokens). + """ + i_nph = img_h / patch_hw + i_npw = img_w / patch_hw + ratio = i_npw / i_nph if i_nph > 0 else 1.0 + if i_nph * i_npw > max_tokens: + i_nph = (max_tokens / ratio) ** 0.5 + i_npw = i_nph * ratio + candidates = list( + set( + itertools.product( + [math.floor(i_nph), math.ceil(i_nph)], + [math.floor(i_npw), math.ceil(i_npw)], + ) + ) + ) + candidates = [ + (nph, npw) + for nph, npw in candidates + if nph >= 1 and npw >= 1 and nph * npw <= max_tokens + ] + if not candidates: + candidates = [(max(1, round(i_nph)), max(1, round(i_npw)))] + nph, npw = min(candidates, key=lambda c: abs(c[0] / c[1] - img_h / img_w)) + return nph * patch_hw, npw * patch_hw, nph * npw + + +class MuseGlimmerImageProcessor(BaseImageProcessor): + """Resize + normalize MuseGlimmer images and compute patch-token counts. + + Variable-resolution: each image is resized to the grid that best matches its + aspect ratio under the per-image token cap, then normalized with mean/std 0.5. + Returns per-image tensors (not stacked) because MuseGlimmer consumes a list of + variable-size images. Video frames are handled by MuseGlimmerVideoProcessor. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + patch_size: int = 14, + downsample_factor: int = 2, + max_image_tokens: int = 4096, + image_mean: float = 0.5, + image_std: float = 0.5, + **kwargs, + ): + super().__init__(**kwargs) + self.patch_size = patch_size + self.downsample_factor = downsample_factor + self.max_image_tokens = max_image_tokens + self.image_mean = image_mean + self.image_std = image_std + + def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor: + # Functional normalize -- no stored Normalize object, so the processor + # stays JSON-serializable for to_dict() / save_pretrained. + return T.functional.normalize( + T.functional.to_tensor(image), + [self.image_mean] * 3, + [self.image_std] * 3, + ) + + # -- size computation (mirrors modeling_muse_glimmer.MuseGlimmerVisionEncoder) -- + def compute_image_size(self, img_w: int, img_h: int) -> tuple[int, int, int]: + ph = self.patch_size * self.downsample_factor + return _grid_size(img_w, img_h, ph, self.max_image_tokens) + + # -- preprocessing --------------------------------------------------------- + def preprocess_image(self, image: Image.Image) -> tuple[torch.Tensor, int]: + """Return (pixel tensor [3, H, W], n_patch_tokens) for one image.""" + image = image.convert("RGB") + target_h, target_w, n_tokens = self.compute_image_size( + image.width, image.height + ) + image = image.resize((target_w, target_h), Image.LANCZOS) + return self._to_norm_tensor(image), n_tokens + + +# torchcodec is the training-faithful video decoder; other decoders (torchvision / +# PyAV) seek + color-convert differently and diverge from training. Optional at +# import time so text/image paths load without it; decode_video raises if used. +try: + import torchcodec +except Exception: + torchcodec = None + + +class MuseGlimmerVideoProcessor(BaseVideoProcessor): + """MuseGlimmer video preprocessing behind the standard HF + ``AutoVideoProcessor`` API. + + Exposes MuseGlimmer's training-faithful video handling as a first-class video + processor so you can call ``processor(videos="clip.mp4")`` and downstream + users can discover it via the standard HF ``AutoVideoProcessor`` API plus a + ``video_preprocessor_config.json``. It wraps: + + * torchcodec decode (matches training; other decoders diverge), + * uniform frame sampling to a whole multiple of ``patch_temporal``, + * real per-group PTS timestamps (rendered as ``Time: X.Xs`` by + MuseGlimmerProcessor), + * ``patch_temporal`` frame-grouping (frames cat on the channel axis -> + ``[patch_temporal * 3, H, W]``; the encoder detects video by channel count). + + It deliberately overrides ``preprocess`` instead of the BaseVideoProcessor fast + pipeline (group_videos_by_shape / smart_resize / single stacked tensor): MuseGlimmer + consumes a LIST of variable-size group tensors and needs the real per-group PTS, + neither of which the stacked fast path models. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + patch_size: int = 14, + downsample_factor: int = 2, + patch_temporal: int = 2, + max_video_frame_tokens: int = 144, + image_mean: float = 0.5, + image_std: float = 0.5, + video_num_frames: int = 96, + video_sampling_fps: float = 2.0, + **kwargs, + ): + super().__init__(**kwargs) + self.patch_size = patch_size + self.downsample_factor = downsample_factor + self.patch_temporal = patch_temporal + self.max_video_frame_tokens = max_video_frame_tokens + self.image_mean = image_mean + self.image_std = image_std + self.video_num_frames = video_num_frames + self.video_sampling_fps = video_sampling_fps + + def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor: + return T.functional.normalize( + T.functional.to_tensor(image), + [self.image_mean] * 3, + [self.image_std] * 3, + ) + + def decode_video(self, video_path: str) -> tuple[list[Image.Image], list[float]]: + """Sample frames + per-group PTS with torchcodec (the training decode path). + + ``timestamps[g]`` is the ACTUAL decoded PTS of the first frame in temporal + group ``g``; ``len(frames)`` is a whole multiple of ``patch_temporal``. + """ + if torchcodec is None: + raise RuntimeError( + "torchcodec is required for video decoding (it matches the training " + "decode path). See the Environment Setup steps in hf/README.md." + ) + pt = self.patch_temporal + reader = torchcodec.decoders.VideoDecoder(video_path) + total = len(reader) + assert reader.metadata.average_fps is not None, ( + f"Video has no FPS metadata: {video_path}" + ) + fps = reader.metadata.average_fps + assert self.video_sampling_fps and self.video_sampling_fps > 0, ( + f"video_sampling_fps must be positive, got {self.video_sampling_fps}" + ) + n = min( + int(total * self.video_sampling_fps / fps), self.video_num_frames, total + ) + n = max(pt, (n // pt) * pt) + n = min(n, total) + if n < pt: + raise ValueError( + f"Video has only {total} decodable frame(s) but needs at least " + f"{pt} (one temporal patch): {video_path}" + ) + indices = torch.linspace(0, total - 1, n).long().tolist() + frames: list[Image.Image] = [] + timestamps: list[float] = [] + for j, i in enumerate(indices): + fr = reader[i] + frames.append( + Image.fromarray(fr.data.permute(1, 2, 0).numpy()).convert("RGB") + ) + if j % pt == 0: + pts = getattr(fr, "pts_seconds", None) + timestamps.append(float(pts) if pts is not None else i / fps) + return frames, timestamps + + def compute_video_frame_size(self, img_w: int, img_h: int) -> tuple[int, int, int]: + """Pick (H, W, tokens_per_frame) for a video frame under the per-frame cap. + + Mirrors MuseGlimmerImageProcessor.compute_image_size but uses the (smaller) + per-frame video token budget; shares the same ``_grid_size`` grid logic as + modeling_muse_glimmer.MuseGlimmerVisionEncoder. + """ + ph = self.patch_size * self.downsample_factor + return _grid_size(img_w, img_h, ph, self.max_video_frame_tokens) + + def _group_frames( + self, frames: list[Image.Image] + ) -> tuple[list[torch.Tensor], int, int]: + """Resize/normalize frames and cat ``patch_temporal`` frames per group. + + Returns (group_tensors, n_groups, tokens_per_group). + """ + pt = self.patch_temporal + if pt <= 0: + raise ValueError(f"patch_temporal must be positive, got {pt}") + if not frames: + raise ValueError("video must contain at least one frame") + if padding := (-len(frames)) % pt: + frames = [*frames, *([frames[-1]] * padding)] + first = frames[0].convert("RGB") + target_h, target_w, n_tokens = self.compute_video_frame_size( + first.width, first.height + ) + groups: list[torch.Tensor] = [] + for i in range(0, len(frames), pt): + grp = [ + self._to_norm_tensor( + frames[i + j] + .convert("RGB") + .resize((target_w, target_h), Image.LANCZOS) + ) + for j in range(pt) + ] + groups.append(torch.cat(grp, dim=0)) + return groups, len(groups), n_tokens + + @staticmethod + def _normalize_videos(videos) -> list: + """Normalize the ``videos`` arg to a list of per-video items. + + A video item is a path (str/Path) or a list of PIL frames. Accepts a + single video (one path, or one list of frames) or a list of those. + """ + if isinstance(videos, (str, Path)): + return [videos] + if videos and not isinstance(videos[0], (list, tuple, str, Path)): + return [videos] # a single list of PIL frames + return list(videos or []) + + def preprocess_one( + self, + video: str | Path | list[Image.Image], + timestamps: list[float] | None = None, + ) -> tuple[list[torch.Tensor], int, int, list[float]]: + """Decode (if a path) + group ONE video. + + Returns (group_tensors, n_groups, tokens_per_group, group_timestamps). This + is the single per-video implementation shared by ``preprocess`` (the HF + AutoVideoProcessor batch API) and ``MuseGlimmerProcessor.__call__`` (prompt + building), so the two paths cannot drift. For a path, frames + real + per-group PTS come from ``decode_video``; for pre-decoded frames the given + ``timestamps`` are used verbatim (empty if not supplied). + """ + if isinstance(video, (str, Path)): + frames, ts = self.decode_video(str(video)) + else: + frames = [f.convert("RGB") for f in video] + ts = list(timestamps) if timestamps is not None else [] + groups, n_groups, tokens_per_group = self._group_frames(frames) + return groups, n_groups, tokens_per_group, ts + + def preprocess( + self, + videos, + video_timestamps: list[list[float]] | None = None, + return_tensors: str | None = "pt", + **kwargs, + ) -> BatchFeature: + """Preprocess video file path(s) or pre-decoded frame list(s). + + Accepts a single video (path str, or list of PIL frames) or a list of + those. Returns a BatchFeature carrying, in video order: + * ``pixel_values`` -- flat list of [pt*3, H, W] group tensors, + * ``video_num_groups`` -- groups per video, + * ``video_tokens_per_group`` -- patch-token count per group per video, + * ``video_timestamps`` -- per-group PTS per video, + i.e. exactly what MuseGlimmerProcessor needs to build the + ``<|vid_start|> ( Time: X.Xs <|video|>*P [sep] )* <|vid_end|>`` block. + """ + videos = self._normalize_videos(videos) + pixel_values: list[torch.Tensor] = [] + num_groups: list[int] = [] + tokens_per_group: list[int] = [] + out_ts: list[list[float]] = [] + for idx, v in enumerate(videos): + ts_in = video_timestamps[idx] if video_timestamps else None + groups, ng, tpg, ts = self.preprocess_one(v, ts_in) + pixel_values += groups + num_groups.append(ng) + tokens_per_group.append(tpg) + out_ts.append(ts) + batch = BatchFeature( + data={ + "video_num_groups": num_groups, + "video_tokens_per_group": tokens_per_group, + "video_timestamps": out_ts, + }, + tensor_type=None, + ) + # Variable-size list -- BatchFeature would fail to stack into one tensor. + batch["pixel_values"] = pixel_values + return batch + + +class MuseGlimmerProcessor(ProcessorMixin): + """Bundle MuseGlimmerImageProcessor + MuseGlimmerVideoProcessor + tokenizer; + expand media sentinels into spans. + + Images go through ``image_processor`` (MuseGlimmerImageProcessor) and videos through + ``video_processor`` (MuseGlimmerVideoProcessor) -- one preprocessing implementation + each, no overlap. + """ + + attributes = ["image_processor", "video_processor", "tokenizer"] + image_processor_class = "AutoImageProcessor" + video_processor_class = "AutoVideoProcessor" + tokenizer_class = "PreTrainedTokenizerFast" + + def __init__( + self, + image_processor=None, + video_processor=None, + tokenizer=None, + chat_template=None, + **kwargs, + ): + if image_processor is None: + image_processor = MuseGlimmerImageProcessor() + if video_processor is None: + video_processor = MuseGlimmerVideoProcessor() + super().__init__( + image_processor, + video_processor, + tokenizer, + chat_template=chat_template or MUSE_GLIMMER_MM_CHAT_TEMPLATE, + **kwargs, + ) + + def _sid(self, token: str) -> int: + return self.tokenizer.convert_tokens_to_ids(token) + + def _image_block(self, n_tokens: int) -> list[int]: + return ( + [self._sid("<|image_start|>")] + + [self._sid("<|patch|>")] * n_tokens + + [self._sid("<|image_end|>")] + ) + + def _video_block( + self, + n_groups: int, + tokens_per_group: int, + timestamps: list[float] | None = None, + ) -> list[int]: + """Per-group ``Time: X.Xs`` + <|video|>*P, separated/terminated. + + ``timestamps`` (one per temporal group, in seconds) should be the ACTUAL + decoded frame times -- training renders the real factored frame PTS, not a + uniform grid. Falls back to ``g*patch_temporal/fps`` only when timestamps + are not supplied (e.g. callers that don't track frame times); that + approximation can differ from training by the per-frame jitter. + """ + vid = self._sid("<|video|>") + sep = self._sid("<|vid_frame_separator|>") + pt = self.video_processor.patch_temporal + fps = self.video_processor.video_sampling_fps + block = [self._sid("<|vid_start|>")] + for g in range(n_groups): + ts = timestamps[g] if timestamps is not None else g * pt / fps + block += self.tokenizer.encode(f"Time: {ts:.1f}s", add_special_tokens=False) + block += [vid] * tokens_per_group + block.append(sep if g < n_groups - 1 else self._sid("<|vid_end|>")) + return block + + def __call__( + self, + text: str | list[str] | None = None, + images: list[Image.Image] | None = None, + videos: list[list[Image.Image] | str] | None = None, + video_timestamps: list[list[float]] | None = None, + return_tensors: str | None = "pt", + **kwargs, + ) -> BatchFeature: + if text is None: + raise ValueError( + "`text` is required (use apply_chat_template to build it)." + ) + if isinstance(text, (list, tuple)): + if len(text) != 1: + raise ValueError( + "MuseGlimmerProcessor supports a single text sample per call." + ) + text = text[0] + + images = list(images or []) + videos = list(videos or []) + image_sentinel = self._sid(IMAGE_SENTINEL) + video_sentinel = self._sid(VIDEO_SENTINEL) + + # Preprocess media up front so we can expand sentinels in document order. + # Images via image_processor; videos via video_processor (each video item + # is a path -- decoded here -- or a pre-decoded list of PIL frames). + prepped_images = [self.image_processor.preprocess_image(im) for im in images] + prepped_videos = [ + self.video_processor.preprocess_one( + v, video_timestamps[i] if video_timestamps else None + ) + for i, v in enumerate(videos) + ] + + ids = self.tokenizer.encode(text, add_special_tokens=False) + n_img = sum(1 for t in ids if t == image_sentinel) + n_vid = sum(1 for t in ids if t == video_sentinel) + if n_img != len(prepped_images): + raise ValueError( + f"{n_img} image sentinel(s) in text but " + f"{len(prepped_images)} image(s) given." + ) + if n_vid != len(prepped_videos): + raise ValueError( + f"{n_vid} video sentinel(s) in text but " + f"{len(prepped_videos)} video(s) given." + ) + + out_ids: list[int] = [] + pixel_values: list[torch.Tensor] = [] + img_i = vid_i = 0 + for tid in ids: + if tid == image_sentinel: + tensor, n_tokens = prepped_images[img_i] + img_i += 1 + out_ids += self._image_block(n_tokens) + pixel_values.append(tensor) + elif tid == video_sentinel: + groups, n_groups, tokens_per_group, ts = prepped_videos[vid_i] + vid_i += 1 + out_ids += self._video_block(n_groups, tokens_per_group, ts or None) + pixel_values += groups + else: + out_ids.append(tid) + + data: dict = { + "input_ids": [out_ids], + "attention_mask": [[1] * len(out_ids)], + } + batch = BatchFeature(data=data, tensor_type=return_tensors) + # Keep pixel_values as a list of variable-size tensors (MuseGlimmerModel + # consumes a list); BatchFeature would fail to stack them into one tensor. + if pixel_values: + batch["pixel_values"] = pixel_values + return batch + + +__all__ = [ + "MuseGlimmerImageProcessor", + "MuseGlimmerVideoProcessor", + "MuseGlimmerProcessor", + "MUSE_GLIMMER_MM_CHAT_TEMPLATE", +] diff --git a/vllm/transformers_utils/utils.py b/vllm/transformers_utils/utils.py index cd215421a981..674f6a05bf98 100644 --- a/vllm/transformers_utils/utils.py +++ b/vllm/transformers_utils/utils.py @@ -48,15 +48,14 @@ def modelscope_list_repo_files( api = HubApi() api.login(token) + # same as huggingface_hub.list_repo_files - files = [ + return [ file["Path"] for file in api.get_model_files( model_id=repo_id, revision=revision, recursive=True ) - if file["Type"] == "blob" ] - return files def _maybe_json_dict(path: str | PathLike) -> dict[str, str]: diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py new file mode 100644 index 000000000000..8e2d7e938b6d --- /dev/null +++ b/vllm/utils/b12x.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lazy accessors for the optional ``b12x`` package.""" + +import functools +import importlib +import importlib.util +from collections.abc import Callable, Hashable, Iterable +from dataclasses import dataclass, fields, is_dataclass +from types import ModuleType +from typing import Any + +import torch + + +@dataclass(frozen=True) +class B12xWarmupUnit: + name: str + key: Hashable + compile: Callable[[], None] + + +@functools.cache +def has_b12x() -> bool: + """Return whether the B12X package is installed.""" + return importlib.util.find_spec("b12x") is not None + + +@functools.cache +def _get_submodule(module_name: str) -> ModuleType | None: + if not has_b12x(): + return None + try: + return importlib.import_module(module_name) + except (ImportError, ModuleNotFoundError): + return None + + +def get_b12x_blockscaled() -> ModuleType | None: + return _get_submodule("b12x.gemm.blockscaled") + + +def get_b12x_intrinsics() -> ModuleType | None: + return _get_submodule("b12x._lib.intrinsics") + + +def get_b12x_mxfp8_linear() -> ModuleType | None: + return _get_submodule("b12x.gemm.mxfp8_linear") + + +def get_b12x_tensor_fp8_linear() -> ModuleType | None: + return _get_submodule("b12x.gemm.tensor_fp8_linear") + + +def b12x_warmup_token_counts( + *, + max_tokens: int, + cudagraph_capture_sizes: Iterable[int] = (), +) -> tuple[int, ...]: + # B12X deduplicates shapes that select the same internal kernel policy. + # Keep the complete serving shape set here rather than duplicating its + # policy-selection heuristics in vLLM. + counts = {1} + counts.update(int(size) for size in cudagraph_capture_sizes if int(size) > 0) + if int(max_tokens) > 0: + counts.add(int(max_tokens)) + return tuple(sorted(counts)) + + +def _same_packed_layout(current: Any, replacement: Any) -> bool: + if type(current) is not type(replacement): + return False + if isinstance(current, torch.Tensor): + return ( + current.shape == replacement.shape + and current.stride() == replacement.stride() + and current.dtype == replacement.dtype + and current.device == replacement.device + ) + if is_dataclass(current): + return all( + _same_packed_layout( + getattr(current, field.name), + getattr(replacement, field.name), + ) + for field in fields(current) + ) + return bool(current == replacement) + + +def _copy_packed_tensors(current: Any, replacement: Any) -> None: + if isinstance(current, torch.Tensor): + current.copy_(replacement) + elif is_dataclass(current): + for field in fields(current): + _copy_packed_tensors( + getattr(current, field.name), + getattr(replacement, field.name), + ) + + +@torch.no_grad() +def reuse_packed_weight_storage(current: Any, replacement: Any) -> Any: + """Reuse packed tensor addresses when a compatible weight is reloaded.""" + if current is None or not _same_packed_layout(current, replacement): + return replacement + _copy_packed_tensors(current, replacement) + return current diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 29adb2a6f3e9..4b78142c6f66 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -552,6 +552,29 @@ def fp8_fp4_mqa_logits( ) +def native_next_n_supported(next_n: int) -> bool: + """Whether the paged MQA logits kernel takes `next_n` Q rows per request. + + SM90 implements only {1, 2, 4}; SM100 and SM120 schedule any `next_n` via + multi-atom tiles. Unsupported values must be flattened to one row per query. + """ + if current_platform.is_device_capability_family(90): + return next_n in (1, 2, 4) + return True + + +def _paged_mqa_logits_schedule_slots(num_sms: int, next_n: int) -> int: + """Scheduler tasks the paged MQA logits kernel launches. + + SM90 `next_n=4` runs one task per 2-CTA multicast cluster rather than per + SM, and `fp8_fp4_paged_mqa_logits` asserts its metadata is sized to match. + """ + num_kv_multicast = ( + 2 if next_n == 4 and current_platform.is_device_capability_family(90) else 1 + ) + return num_sms // num_kv_multicast + + def get_paged_mqa_logits_metadata( context_lens: torch.Tensor, block_size: int, @@ -561,22 +584,24 @@ def get_paged_mqa_logits_metadata( """Build scheduling metadata for paged MQA logits. Args: - context_lens: Tensor of shape [B], dtype int32; effective context length - per batch element. + context_lens: Tensor of shape [B, next_n], dtype int32; effective + context length per Q row. block_size: KV-cache block size in tokens (e.g., 64). num_sms: Number of SMs available. 132 for Hopper indices: Optional request index for each varlen row. Returns: - Backend-specific tensor consumed by `fp8_fp4_paged_mqa_logits` to - schedule work across SMs. + Tensor of shape [slots + 1, 2] consumed by `fp8_fp4_paged_mqa_logits` + to schedule work across SMs. """ _lazy_init() if _get_paged_mqa_logits_metadata_impl is None: return _missing() + next_n = context_lens.shape[1] if context_lens.dim() == 2 else 1 + num_slots = _paged_mqa_logits_schedule_slots(num_sms, next_n) kwargs = {} if indices is None else {"indices": indices} return _get_paged_mqa_logits_metadata_impl( - context_lens, block_size, num_sms, **kwargs + context_lens, block_size, num_slots, **kwargs ) @@ -752,6 +777,7 @@ def should_use_deepgemm_for_fp8_linear( "fp8_fp4_mqa_logits", "fp8_fp4_paged_mqa_logits", "get_paged_mqa_logits_metadata", + "native_next_n_supported", "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", "is_deep_gemm_supported", diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 121621ba202c..4f7d2eecd780 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -126,6 +126,7 @@ def wrapper(*args, **kwargs): "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked" ) flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize") +flashinfer_mxfp4_quantize = _lazy_import_wrapper("flashinfer", "mxfp4_quantize") nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize") silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper( "flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize" @@ -236,6 +237,22 @@ def has_flashinfer_sparse_mla_sm120() -> bool: ) +@functools.cache +def has_flashinfer_sparse_mla_sm120_config(num_q_heads: int, top_k: int) -> bool: + """Return whether FlashInfer ships an SM120 DSV4 decode specialization. + + The public sparse MLA API predates some DSV4 shapes, so checking only that + the callable exists can select a package that later aborts or rejects a + valid vLLM configuration. Inspect FlashInfer's dispatch table until it + exposes a public capability query. + """ + if not has_flashinfer_sparse_mla_sm120(): + return False + mod = _get_submodule("flashinfer.mla._sparse_mla_sm120") + dispatch = getattr(mod, "_DECODE_DSV4_DISPATCH", None) if mod else None + return dispatch is not None and (int(num_q_heads), int(top_k)) in dispatch + + @functools.cache def has_flashinfer_cutedsl() -> bool: """Return ``True`` if FlashInfer cutedsl module is available.""" diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 9f6ea0fec401..03e301317f25 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -412,6 +412,21 @@ def _has_module(module_name: str) -> bool: return True +@cache +def _has_module_spec(module_name: str) -> bool: + """Return True if *module_name* is installed, without importing it. + + Unlike [`_has_module`][vllm.utils.import_utils._has_module], this only + resolves the import spec. It therefore does not pay the import cost of + heavyweight modules, at the price of not verifying that native + dependencies (shared libraries, etc.) are satisfied. The result is cached. + """ + try: + return importlib.util.find_spec(module_name) is not None + except Exception: + return False + + def has_deep_ep() -> bool: """Whether the optional `deep_ep` package is available.""" return _has_module("deep_ep") @@ -504,8 +519,13 @@ def has_triton_kernels() -> bool: @cache def has_tilelang() -> bool: - """Whether the optional `tilelang` package is available.""" - if not _has_module("tilelang"): + """Whether the optional `tilelang` package is available. + + Only the import spec is checked: importing `tilelang` is expensive, so + callers must import it lazily at their point of use rather than relying + on this function to have imported it already. + """ + if not _has_module_spec("tilelang"): return False # ROCm-only guard, imported lazily to avoid loading rocm on CUDA. from vllm.platforms import current_platform diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index cc823f704b65..bb0bb8df9a1e 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -30,6 +30,7 @@ from typing import Any, Literal, cast from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.triton_utils.importing import HAS_TRITON logger = init_logger(__name__) @@ -74,7 +75,12 @@ def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None: _setup_triton_autotuning_print() _setup_triton_jit_hook() _setup_cutedsl_jit_hook() - _setup_tilelang_jit_hook() + + # Refer to #51159. tilelang ships broken symbols + # on rocm. + # TODO: Remove the guard once tilelang upstream is fixed. + if not current_platform.is_rocm(): + _setup_tilelang_jit_hook() logger.info( "Kernel JIT monitor activated; monitored JIT compilations during " diff --git a/vllm/utils/serial_utils.py b/vllm/utils/serial_utils.py index 5fde5ac7105d..5bbeee7419ca 100644 --- a/vllm/utils/serial_utils.py +++ b/vllm/utils/serial_utils.py @@ -64,6 +64,13 @@ def tensor2base64(x: torch.Tensor) -> str: return pybase64.b64encode(binary_data).decode("utf-8") +def numpy2base64(array: np.ndarray) -> str: + """Encode a NumPy array using its `.npy` representation.""" + with io.BytesIO() as buffer: + np.save(buffer, array, allow_pickle=False) + return pybase64.b64encode(buffer.getbuffer()).decode("ascii") + + def tensor2binary( tensor: torch.Tensor, embed_dtype: "EmbedDType | MmMetadataDType", diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index e3f05df8dc69..daa217f8553d 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.utils import KVCacheLayoutType - from vllm.v1.kv_cache_interface import KVCacheSpec, KVQuantMode + from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, KVQuantMode from vllm.v1.kv_cache_interface import get_kv_quant_mode @@ -69,6 +69,13 @@ class AttentionBackend(ABC): def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(1)] + @classmethod + def get_supported_kernel_block_sizes_for_config( + cls, vllm_config: "VllmConfig" + ) -> list[int | MultipleOf]: + """Return kernel block sizes for a concrete engine configuration.""" + return cls.get_supported_kernel_block_sizes() + @staticmethod @abstractmethod def get_name() -> str: @@ -190,6 +197,20 @@ def supports_block_size(cls, block_size: int | None) -> bool: return True return False + @classmethod + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": + """Adjust the layer's KV cache spec for this backend's kernels. Used when the + kernels want KV packed in a specific way. + + NOTE: temporary compatibility API. Today the Attention layer builds the spec + from the model config and the backend only gets to adjust it post-hoc; the end + state is for the backend to build and return the spec directly, at which point + this hook goes away. + + (see: https://github.com/vllm-project/vllm/issues/42449) + """ + return spec + @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: supported_sizes = cls.get_supported_kernel_block_sizes() @@ -201,6 +222,13 @@ def get_preferred_block_size(cls, default_block_size: int) -> int: return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes) + @classmethod + def get_preferred_block_size_for_config( + cls, default_block_size: int, vllm_config: "VllmConfig" + ) -> int: + """Return the preferred block size for a concrete engine config.""" + return cls.get_preferred_block_size(default_block_size) + @classmethod def indexes_kv_by_block_stride(cls) -> bool: """Whether the backend reads KV pages by the runtime block stride. @@ -302,6 +330,11 @@ def supports_pcp(cls) -> bool: except NotImplementedError: return False + @classmethod + def supports_non_causal_dcp(cls) -> bool: + builder_cls = cls.get_builder_cls() + return bool(getattr(builder_cls, "supports_non_causal_multi_token_dcp", False)) + @classmethod def supports_attn_type(cls, attn_type: str) -> bool: """Check if backend supports a given attention type. @@ -350,6 +383,7 @@ def validate_configuration( use_kv_connector: bool = False, use_pcp: bool = False, use_adaptive_verification: bool = False, + use_dcp: bool = False, ) -> list[str]: invalid_reasons = [] if not cls.supports_head_size(head_size): @@ -386,6 +420,8 @@ def validate_configuration( invalid_reasons.append("sliding window not supported") if use_non_causal and not cls.supports_non_causal(): invalid_reasons.append("non-causal attention not supported") + if use_mla and use_non_causal and use_dcp and not cls.supports_non_causal_dcp(): + invalid_reasons.append("non-causal MLA attention with DCP not supported") if use_batch_invariant and not cls.supports_batch_invariance(): invalid_reasons.append("batch invariance not supported") if use_kv_connector and not cls.supports_kv_connector(): diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 3d09d789f3cc..a6cd254fea06 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -31,6 +31,7 @@ ) from vllm.v1.attention.backends.utils import ( KVCacheLayoutType, + get_num_attention_heads_from_layers, ) from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -151,13 +152,15 @@ def __init__( parallel_config = vllm_config.parallel_config self.num_kv_heads = kv_cache_spec.num_kv_heads - self.num_heads = vllm_config.model_config.get_num_attention_heads( - parallel_config - ) + # The scheduler metadata built here sizes a scratchpad from the query + # head count, so it must come from this group's layers: the model-wide + # count is wrong for models that vary it per layer (e.g. Laguna). + self.num_heads = get_num_attention_heads_from_layers( + vllm_config, layer_names + ) or vllm_config.model_config.get_num_attention_heads(parallel_config) self.head_dim = kv_cache_spec.head_size self.dtype = vllm_config.model_config.dtype - # Resolved from the layers on the first build(), once they exist. - self.window_size: int | None = None + self.window_size = self._group_sliding_window() self.block_size = vllm_config.cache_config.block_size self.kv_cache_dtype = vllm_config.cache_config.cache_dtype self.isa = _get_attn_isa( @@ -198,9 +201,6 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CPUAttentionMetadata: - if self.window_size is None: - self.window_size = self._group_sliding_window() - num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens max_query_len = common_attn_metadata.max_query_len diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 22f9e6a14d8f..907881e18773 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -72,7 +72,6 @@ def get_flash_attn_version( head_size: int | None = None, head_size_v: int | None = None, has_sinks: bool = False, - requires_local_attention: bool = False, ) -> int | None: if current_platform.is_xpu(): return 2 @@ -169,28 +168,26 @@ def get_flash_attn_version( ) fa_version = 2 - if ( - fa_version == 4 - and device_capability.major >= 10 - and head_size == 256 - and requires_local_attention - ): + # TODO: Restore the `requires_local_attention` restriction when FA4 + # head-dim 256 is re-enabled. + if fa_version == 4 and device_capability.major >= 10 and head_size == 256: logger.warning_once( - "FA4 on Blackwell does not support local attention with " - "head_size=256, defaulting to FA version 2." + "FA4 on Blackwell is temporarily disabled for head_size=256, " + "defaulting to FA version 2." ) fa_version = 2 # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict - # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA - # prefill). Development of symmetric 192, 384, and 512 support is being tracked - # in https://github.com/Dao-AILab/flash-attention/issues/2456 + # supported head dimensions to ≤128. The 192/128 MLA prefill case is + # supported; 256 is temporarily disabled until upstream supports the + # required features. Development of symmetric 192, 384, and 512 support + # is tracked in https://github.com/Dao-AILab/flash-attention/issues/2456 if ( fa_version == 4 and device_capability.major >= 10 and head_size is not None and head_size > 128 - and not (head_size == 256 or (head_size == 192 and head_size_v == 128)) + and not (head_size == 192 and head_size_v == 128) ): logger.warning_once( "FA4 on Blackwell does not support head_size=%d due to TMEM " @@ -241,9 +238,9 @@ def flash_attn_supports_kv_cache_dtype( head_size_v=head_size_v, has_sinks=has_sinks, ) - return (fa_version == 3 and current_platform.is_device_capability_family(90)) or ( - fa_version == 4 and current_platform.is_device_capability_family(100) - ) + return ( + fa_version in (3, 4) and current_platform.is_device_capability_family(90) + ) or (fa_version == 4 and current_platform.is_device_capability_family(100)) def flash_attn_supports_quant_query_input() -> bool: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index e7ef03d99cb8..d176bd5e928a 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -85,13 +85,58 @@ class FlashAttentionBackend(AttentionBackend): ] @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + def _get_sm90_fa4_fp8_kv_block_size( + vllm_config: VllmConfig | None = None, + ) -> int | None: + if vllm_config is None: + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None or vllm_config.model_config is None: + return None + + head_size = vllm_config.model_config.get_head_size() + if ( + current_platform.is_device_capability_family(90) + and vllm_config.cache_config.cache_dtype in ("fp8", "fp8_e4m3") + and head_size == 512 + and get_flash_attn_version(head_size=head_size) == 4 + ): + # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page. + return 64 + return None + + @classmethod + def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]: + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(): + # Sliding-window cache specs select the smallest advertised size. + # Report the kernel's exact page-size contract instead of the + # generic FlashAttention multiple-of-16 capability. + return [block_size] + return [MultipleOf(16)] + + @classmethod + def get_supported_kernel_block_sizes_for_config( + cls, vllm_config: VllmConfig + ) -> list[int | MultipleOf]: + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config): + return [block_size] return [MultipleOf(16)] forward_includes_kv_cache_update: bool = False @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(): + return max(default_block_size, block_size) + if current_platform.is_xpu(): + return max(default_block_size, 64) + return super().get_preferred_block_size(default_block_size) + + @classmethod + def get_preferred_block_size_for_config( + cls, default_block_size: int, vllm_config: VllmConfig + ) -> int: + if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config): + return max(default_block_size, block_size) if current_platform.is_xpu(): return max(default_block_size, 64) return super().get_preferred_block_size(default_block_size) @@ -860,7 +905,6 @@ def __init__( self.attn_type = attn_type self.vllm_flash_attn_version = get_flash_attn_version( requires_alibi=alibi_slopes is not None, - requires_local_attention=sliding_window is not None, head_size=head_size, has_sinks=sinks is not None, ) @@ -895,7 +939,17 @@ def __init__( "heads in the layer" ) - self.supports_quant_query_input = flash_attn_supports_quant_query_input() + # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes + # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q, + # K, and V to have the same FP8 dtype. + uses_sm90_fa4_fp8_kv_dequant = ( + self.vllm_flash_attn_version == 4 + and current_platform.is_device_capability_family(90) + and self.kv_cache_dtype in ("fp8", "fp8_e4m3") + ) + self.supports_quant_query_input = flash_attn_supports_quant_query_input() and ( + not uses_sm90_fa4_fp8_kv_dequant + ) vllm_config = get_current_vllm_config_or_none() dcp_a2a = ( diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 5ef7759b70c7..de51e2198ba2 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer with FlashInfer.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from enum import Enum from functools import partial from typing import ClassVar @@ -10,6 +10,7 @@ import numpy as np import torch from flashinfer import ( + BatchAttentionWithAttentionSinkWrapper, BatchDecodeWithPagedKVCacheWrapper, BatchPrefillWithPagedKVCacheWrapper, BatchPrefillWithRaggedKVCacheWrapper, @@ -45,10 +46,12 @@ supports_trtllm_attention, use_trtllm_attention, ) +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( PIN_MEMORY, canonicalize_singleton_dim_strides, + get_dtype_size, is_quantized_kv_cache, is_strictly_contiguous, nvfp4_kv_cache_full_dim, @@ -386,6 +389,21 @@ def run( class FlashInferBackend(AttentionBackend): + @classmethod + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": + """NVFP4 stores K and V as separate per-head slots of packed fp4 data + plus fp8 block scales.""" + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_nvfp4: + return spec + hs_k = nvfp4_kv_cache_full_dim(spec.head_size) + hs_v = nvfp4_kv_cache_full_dim(spec.head_size_v) + assert hs_k == hs_v, "nvfp4 with asymmetric K/V head sizes not yet supported" + return replace( + spec, + num_head_slots=2 * spec.num_kv_heads, + state_content_bytes=hs_k * get_dtype_size(spec.dtype), + ) + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ "auto", @@ -1113,11 +1131,27 @@ def _get_prefill_wrapper( "NVFP4 KV cache." ) if self._noncausal_prefill_wrapper is None: - self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( - self._get_workspace_buffer(), - get_kv_cache_layout(), - backend="auto", - ) + if self.has_sinks and current_platform.is_device_capability_family(120): + self._noncausal_prefill_wrapper = ( + BatchAttentionWithAttentionSinkWrapper( + self._get_workspace_buffer(), + get_kv_cache_layout(), + backend="auto", + q_data_type=self.q_data_type_prefill, + kv_data_type=self.kv_cache_dtype, + head_dim_qk=self.head_dim, + head_dim_vo=self.head_dim, + window_left=self.window_left, + ) + ) + else: + self._noncausal_prefill_wrapper = ( + BatchPrefillWithPagedKVCacheWrapper( + self._get_workspace_buffer(), + get_kv_cache_layout(), + backend="auto", + ) + ) return self._noncausal_prefill_wrapper if self._prefill_wrapper is None: @@ -1127,14 +1161,27 @@ def _get_prefill_wrapper( dcp_a2a=self.dcp_a2a, ) else: - # NVFP4 KV cache requires the trtllm-gen backend inside - # the wrapper; fa2/fa3 do not support nvfp4. - backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" - self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( - self._get_workspace_buffer(), - get_kv_cache_layout(), - backend=backend, - ) + if self.has_sinks and current_platform.is_device_capability_family(120): + assert not self.is_kvcache_nvfp4 + self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper( + self._get_workspace_buffer(), + get_kv_cache_layout(), + backend="auto", + q_data_type=self.q_data_type_prefill, + kv_data_type=self.kv_cache_dtype, + head_dim_qk=self.head_dim, + head_dim_vo=self.head_dim, + window_left=self.window_left, + ) + else: + # NVFP4 KV cache requires the trtllm-gen backend inside + # the wrapper; fa2/fa3 do not support nvfp4. + backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" + self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( + self._get_workspace_buffer(), + get_kv_cache_layout(), + backend=backend, + ) assert self._prefill_wrapper is not None return self._prefill_wrapper @@ -1362,13 +1409,15 @@ def build( # seq_lens_cpu is not needed since TRTLLM paths use GPU tensors # (block_tables, seq_lens) directly. needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm - seq_lens_cpu = common_attn_metadata.seq_lens_cpu if needs_seq_lens_cpu else None - seq_lens_np = seq_lens_cpu.numpy() if seq_lens_cpu is not None else None - num_blocks_np = ( - (seq_lens_np + (page_size - 1)) // page_size - if seq_lens_np is not None - else None - ) + if needs_seq_lens_cpu: + with gpu_sync_allowed(): + seq_lens_cpu = common_attn_metadata.seq_lens_cpu + seq_lens_np = seq_lens_cpu.numpy() + num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size + else: + seq_lens_cpu = None + seq_lens_np = None + num_blocks_np = None # Adjust seq_lens_cpu for DCP if self.use_dcp: @@ -2086,16 +2135,28 @@ def forward( else: out_prefill = output[num_decode_tokens:] - prefill_wrapper.run( - prefill_query, - kv_cache_for_fi, - q_scale=layer._q_scale_float, - k_scale=layer._k_scale_float, - v_scale=layer._v_scale_float, - out=out_prefill, - kv_cache_sf=kv_cache_sf, - sinks=self.sinks, - ) + if isinstance( + prefill_wrapper, BatchAttentionWithAttentionSinkWrapper + ): + assert self.sinks is not None + prefill_wrapper.run( + prefill_query, + kv_cache_for_fi, + self.sinks, + self.scale * layer._q_scale_float * layer._k_scale_float, + v_scale=layer._v_scale_float, + out=out_prefill, + ) + else: + prefill_wrapper.run( + prefill_query, + kv_cache_for_fi, + q_scale=layer._q_scale_float, + k_scale=layer._k_scale_float, + v_scale=layer._v_scale_float, + out=out_prefill, + kv_cache_sf=kv_cache_sf, + ) if needs_fp8_out_prefill: output[ diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py index 36b115f64444..47d331ba7bd4 100644 --- a/vllm/v1/attention/backends/mla/compressor_utils.py +++ b/vllm/v1/attention/backends/mla/compressor_utils.py @@ -3,6 +3,18 @@ import torch from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import cdiv + +_DSPARK_SWA_INDEX_ALIGNMENT = 64 + + +def get_dspark_swa_index_width( + window_size: int, + num_speculative_tokens: int, +) -> int: + """Return the padded width of non-causal DSpark SWA indices.""" + width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0) + return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT @triton.jit diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index bf92bf692bf0..79beb5332e2a 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import ClassVar +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar import torch from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla @@ -15,6 +16,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import ( MLACommonBackend, + MLACommonDecodeMetadata, MLACommonImpl, MLACommonMetadata, MLACommonMetadataBuilder, @@ -30,6 +32,10 @@ ) from vllm.v1.attention.backends.utils import KVCacheLayoutType +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import AttentionSpec + logger = init_logger(__name__) @@ -112,12 +118,56 @@ def _get_multi_ctas_kv_counter_buffer( return _fi_multi_ctas_kv_counter -class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): +@dataclass +class FlashInferMLADecodeMetadata(MLACommonDecodeMetadata): + flattened_block_table: torch.Tensor | None = None + flattened_seq_lens: torch.Tensor | None = None + query_len: int = 0 + + +@dataclass +class FlashInferMLAMetadata(MLACommonMetadata[FlashInferMLADecodeMetadata]): + pass + + +class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[FlashInferMLAMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM # Non-causal DSpark blocks are flattened to single-token rows in forward_mqa. supports_non_causal_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, + FlashInferMLAMetadata, + supports_dcp_with_varlen=True, + ) + + def _build_decode( + self, + block_table_tensor: torch.Tensor, + seq_lens_device: torch.Tensor, + max_seq_len: int, + query_start_loc_cpu: torch.Tensor, + query_start_loc_device: torch.Tensor, + num_decode_tokens: int, + dcp_tot_seq_lens_device: torch.Tensor | None, + ) -> FlashInferMLADecodeMetadata: + return FlashInferMLADecodeMetadata( + block_table=block_table_tensor, + seq_lens=seq_lens_device, + dcp_tot_seq_lens=dcp_tot_seq_lens_device, + ) + class FlashInferMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] @@ -193,7 +243,7 @@ def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": return "HND" -class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]): +class FlashInferMLAImpl(MLACommonImpl[FlashInferMLAMetadata]): can_return_lse_for_decode: bool = True # trtllm-gen MLA decode emits LSE in log2 (per flashinfer's own # reference at flashinfer/trace/templates/attention.py:81: @@ -264,7 +314,7 @@ def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], kv_c_and_k_pe_cache: torch.Tensor, - attn_metadata: MLACommonMetadata, + attn_metadata: FlashInferMLAMetadata, layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: assert kv_c_and_k_pe_cache.numel() > 0 @@ -276,16 +326,28 @@ def forward_mqa( block_table = attn_metadata.decode.block_table seq_lens = attn_metadata.decode.seq_lens + query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes if not attn_metadata.causal: # Non-causal DSpark block: flatten to single-token decode rows with # per-row context seq_lens (trtllm-gen has no causal flag and would # otherwise mask the block causally). - query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes q = q.unsqueeze(1) if query_len > 1: - block_table = block_table.repeat_interleave(query_len, dim=0) - seq_lens = seq_lens.repeat_interleave(query_len) + block_table, seq_lens = self._prepare_flattened_decode_metadata( + attn_metadata, + query_len, + causal=False, + ) + elif self.dcp_world_size > 1 and query_len > 1: + # Causal DCP block: flatten to single-token decode rows with + # per-row rank-local seq_lens for each query's visible prefix. + block_table, seq_lens = self._prepare_flattened_decode_metadata( + attn_metadata, + query_len, + causal=True, + ) + q = q.unsqueeze(1) # trtllm API requires extra dimension q_len_per_request for MTP elif attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0: logger.warning_once( @@ -358,3 +420,49 @@ def forward_mqa( o = o.view(-1, o.shape[-2], o.shape[-1]) return o, lse + + def _prepare_flattened_decode_metadata( + self, + attn_metadata: FlashInferMLAMetadata, + query_len: int, + *, + causal: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare flattened decode metadata once for all layers in the group.""" + decode = attn_metadata.decode + assert decode is not None + if decode.query_len: + assert decode.query_len == query_len + assert decode.flattened_block_table is not None + assert decode.flattened_seq_lens is not None + return decode.flattened_block_table, decode.flattened_seq_lens + + block_table = decode.block_table.repeat_interleave(query_len, dim=0) + if causal: + global_seq_lens = decode.dcp_tot_seq_lens + assert global_seq_lens is not None + offsets = torch.arange( + query_len - 1, + -1, + -1, + device=global_seq_lens.device, + dtype=global_seq_lens.dtype, + ) + per_query_global_lens = torch.clamp( + (global_seq_lens.unsqueeze(1) - offsets).reshape(-1), min=0 + ) + interleave = self.cp_kv_cache_interleave_size + dcp_span = self.dcp_world_size * interleave + remainder = torch.clamp( + per_query_global_lens % dcp_span - self.dcp_rank * interleave, + min=0, + max=interleave, + ) + seq_lens = per_query_global_lens // dcp_span * interleave + remainder + else: + seq_lens = decode.seq_lens.repeat_interleave(query_len) + + decode.flattened_block_table = block_table + decode.flattened_seq_lens = seq_lens + decode.query_len = query_len + return block_table, seq_lens diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 30951ca4c12c..1177d8b27afe 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -21,6 +21,7 @@ from vllm.utils.deep_gemm import ( get_paged_mqa_logits_metadata, has_deep_gemm, + native_next_n_supported, ) from vllm.utils.platform_utils import num_compute_units from vllm.v1.attention.backend import ( @@ -39,6 +40,28 @@ logger = init_logger(__name__) +# The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout) +# and mxfp4 is the opt-in Blackwell path. +DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4") + + +def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool: + """Whether the DeepSeek sparse indexer should use the MXFP4 K cache.""" + kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("fp8") + if kv_dtype not in DSA_INDEXER_KV_DTYPES: + raise ValueError( + f"indexer_kv_dtype={kv_dtype!r} is not supported by the DeepSeek " + f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})." + ) + use_fp4 = kv_dtype == "mxfp4" + if use_fp4 and not current_platform.is_device_capability_family(100): + raise ValueError( + "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs " + "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and " + "earlier architectures are not supported." + ) + return use_fp4 + @triton.jit def _prepare_uniform_decode_kernel( @@ -56,9 +79,12 @@ def _prepare_uniform_decode_kernel( req_id = idx // max_decode_len local_idx = idx % max_decode_len - # Compute number of KVs attended to by this token. + # Compute number of KVs attended to by this token. Padding requests have + # seq_len == 0, which would otherwise make the first token of each padded + # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels + # read these as uint32, turning -1 into ~4e9. seq_len = tl.load(seq_lens_ptr + req_id) - per_token_seq_len = seq_len - max_decode_len + local_idx + 1 + per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0) tl.store(decode_seq_lens_ptr + idx, per_token_seq_len) # Copy block table row. @@ -140,7 +166,7 @@ def get_name() -> str: @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [1, 64] if current_platform.is_rocm() else [64] + return [1, MultipleOf(16)] if current_platform.is_rocm() else [64] @classmethod def get_supported_head_sizes(cls) -> list[int]: @@ -467,6 +493,20 @@ def _supports_varlen_paged_mqa_logits() -> bool: ) +def _supports_native_decode(next_n: int) -> bool: + """Whether decode can pass `next_n` Q rows per request to the kernel + instead of flattening to one single-token row per query, which re-reads + the KV tile once per row. + """ + if not (current_platform.is_cuda() and has_deep_gemm()): + return next_n in (1, 2) + if current_platform.is_device_capability_family(100): + return True + if current_platform.is_device_capability_family(90): + return native_next_n_supported(next_n) + return next_n in (1, 2) + + class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # The indexer opts out of the shared reorder-threshold vote (see __init__), # so this is None; its own split uses self.decode_threshold. @@ -508,34 +548,16 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: if self.vllm_config.speculative_config else 0 ) - self.use_fp4_indexer_cache = ( - self.vllm_config.attention_config.use_fp4_indexer_cache - ) - - assert ( - current_platform.is_device_capability_family(100) - or not self.use_fp4_indexer_cache - ), ( - "use_fp4_indexer_cache requires Blackwell datacenter GPUs " - "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and " - "earlier architectures are not supported." - ) + self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config) next_n = self.num_speculative_tokens + 1 self.decode_threshold = next_n self.reorder_batch_threshold = None - # NOTE: SM100 datacenter GPUs support any next_n natively via the - # multi-atom paged MQA logits kernels (FP8 and FP4 indexer - # caches). Outside the SM100 family the FP8 - # paged MQA logits kernel only supports next_n in (1, 2) - # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there. - self.use_flattening = not current_platform.is_device_capability_family( - 100 - ) and next_n not in (1, 2) + self.use_flattening = not _supports_native_decode(next_n) self.supports_varlen = _supports_varlen_paged_mqa_logits() logger.info_once( "DSA indexer decode path: use_flattening=%s supports_varlen=%s " - "(next_n=%d, use_fp4_indexer_cache=%s)", + "(next_n=%d, use_fp4_cache=%s)", self.use_flattening, self.supports_varlen, next_n, @@ -585,7 +607,8 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: device=self.device, ) - # See: DeepGMM/csrc/apis/attention.hpp + # See: DeepGMM/csrc/apis/attention.hpp. Sized for one slot per SM; + # build() narrows it to whatever the kernel actually schedules. self.scheduler_metadata_buffer = torch.empty( (self.num_sms + 1, 2), dtype=torch.int32, device=self.device ) @@ -739,12 +762,15 @@ def _prepare_decode_tensors( seq_lens_buffer = self.decode_seq_lens_buffer[ : num_decodes * max_decode_len ].view(num_decodes, max_decode_len) + # Clamp at 0: padding requests have seq_len == 0, which would + # otherwise make token 0 negative (next_n=2 gives 0-2+1+0 = -1). + # Downstream kernels read these as uint32, turning -1 into ~4e9. seq_lens_buffer[:] = ( seq_lens.unsqueeze(1) - max_decode_len + 1 + self.offsets_buffer[:max_decode_len] - ) + ).clamp_(min=0) seq_lens = seq_lens_buffer return seq_lens, block_table, decode_lens, num_decodes, requires_padding @@ -924,9 +950,16 @@ def build( max_decode_len = int(decode_lens_cpu.max().item()) next_n = 1 + self.num_speculative_tokens + # The kernel sees max_decode_len Q rows, not the configured next_n, + # so legality is per-step: on SM90 a uniformly 3-deep batch has no + # native kernel. max_decode_len <= 1 always has one. + step_next_n_ok = max_decode_len <= 1 or _supports_native_decode( + max_decode_len + ) use_native = ( not (self.use_flattening or self.supports_varlen) and max_decode_len <= next_n + and step_next_n_ok ) global_seq_lens_for_decode = self._prepare_global_decode_seq_lens( @@ -992,20 +1025,23 @@ def build( seq_lens = seq_lens.unsqueeze(-1) # DeepGEMM is required for the paged MQA logits on CUDA devices + schedule_metadata = self.scheduler_metadata_buffer if current_platform.is_cuda() and has_deep_gemm(): - self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata( + metadata = get_paged_mqa_logits_metadata( seq_lens, self.kv_cache_spec.storage_block_size, self.num_sms, indices=decode_indices, ) + schedule_metadata = self.scheduler_metadata_buffer[: metadata.shape[0]] + schedule_metadata[:] = metadata decode_metadata = DeepSeekV32IndexerDecodeMetadata( block_table=block_table, seq_lens=seq_lens, decode_lens=decode_lens, requires_padding=requires_padding, - schedule_metadata=self.scheduler_metadata_buffer, + schedule_metadata=schedule_metadata, indices=decode_indices, global_seq_lens=global_seq_lens_for_decode, ) diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index ccbdf728903a..2f832a1357a9 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -777,7 +777,11 @@ def build( attn_metadata.reduce_indptr = self._mla_reduce_indptr attn_metadata.reduce_final_map = self._mla_reduce_final_map attn_metadata.reduce_partial_map = self._mla_reduce_partial_map - if self._fp8_prefill_enabled and attn_metadata.prefill is not None: + if ( + self._fp8_prefill_enabled + and attn_metadata.prefill is not None + and attn_metadata.prefill.chunked_context is None + ): self._build_fp8_prefill_ps_metadata(attn_metadata, common_attn_metadata) return attn_metadata diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index d5db42cb7f96..3f10d6105349 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py @@ -275,7 +275,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend): @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [1, 64] + return [1, MultipleOf(16)] @staticmethod def get_name() -> str: diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index da0c125c62df..d858eb677cba 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -22,6 +22,9 @@ CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.mla.compressor_utils import ( + get_dspark_swa_index_width, +) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata from vllm.v1.kv_cache_interface import ( @@ -97,6 +100,8 @@ 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, model_version="deepseek_v4", @@ -172,8 +177,10 @@ class DeepseekSparseSWAMetadata: is_valid_token: torch.Tensor | None = None # [num_tokens] token_to_req_indices: torch.Tensor | None = None # [num_tokens] - decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, window_size] + decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, width] decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens] + # window_size (causal) or noncausal_index_width (DSpark non-causal). + decode_swa_width: int = 0 # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill). prefill_swa_indices: torch.Tensor | None = ( None # [num_prefill_tokens, 1, window_size] @@ -485,11 +492,14 @@ def __init__(self, *args, **kwargs): # DSpark draft: the block is non-causal (every query attends to the # trailing window of context PLUS all query tokens, including future ones), # so its per-token index list is wider than `window_size`. The kernel pads - # the q-head count to B_TOPK (64/128), which requires the index width to be - # a multiple of 128. + # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical + # SWA window remains unchanged when the padded matrix is built. self.is_dspark = spec_config is not None and spec_config.use_dspark() self.noncausal_index_width = ( - cdiv(self.window_size + self.num_speculative_tokens, 128) * 128 + get_dspark_swa_index_width( + self.window_size, + self.num_speculative_tokens, + ) if self.is_dspark else 0 ) @@ -534,6 +544,9 @@ def build( is_valid_token.copy_(slot_mapping >= 0) non_causal = not common_attn_metadata.causal + decode_swa_width = ( + self.noncausal_index_width if non_causal else self.window_size + ) decode_swa_indices = self.decode_swa_indices if num_decode_tokens > 0: self.decode_swa_lens[num_decode_tokens:] = 0 @@ -632,6 +645,7 @@ def build( token_to_req_indices=token_to_req_indices, decode_swa_indices=decode_swa_indices[:num_decode_tokens], decode_swa_lens=self.decode_swa_lens[:num_decode_tokens], + decode_swa_width=decode_swa_width, prefill_swa_indices=( self.prefill_swa_indices[:num_prefill_tokens] if num_prefill_tokens > 0 diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 8ab81c438f20..57a628b82b53 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -61,6 +61,7 @@ class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): # The kernel accepts an explicit causal mask, so a non-causal DSpark # block can remain fused instead of being flattened to single tokens. supports_non_causal_multi_token_decode: ClassVar[bool] = True + supports_non_causal_multi_token_dcp: ClassVar[bool] = True def __init__( self, diff --git a/vllm/v1/attention/backends/recoverssm_metadata.py b/vllm/v1/attention/backends/recoverssm_metadata.py new file mode 100644 index 000000000000..7e4cca68f8b3 --- /dev/null +++ b/vllm/v1/attention/backends/recoverssm_metadata.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import abc +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class RecoverSSMPostprocessMetadata: + """Metadata used during postprocessing for align-mode prefix caching.""" + + num_spec_decodes: int + request_indices: torch.Tensor | None + block_table: torch.Tensor + num_computed_tokens: torch.Tensor + block_size: int + + +class RecoverSSMMetadata(abc.ABC): + @abc.abstractmethod + def commit_recoverssm_state( + self, num_accepted_tokens: torch.Tensor + ) -> RecoverSSMPostprocessMetadata | None: + raise NotImplementedError diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 7814a8c0f8de..4b1c167e7f4f 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """High-Performance Triton-only Attention layer.""" -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import ClassVar import torch @@ -19,7 +19,7 @@ from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.math_utils import next_power_of_2 -from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.utils.torch_utils import get_dtype_size, is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -274,6 +274,20 @@ def update_draft_decode_metadata(self, _metadata: TritonAttentionMetadata) -> No class TritonAttentionBackend(AttentionBackend): + @classmethod + def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": + """Per-token-head modes pack inline fp32 scales after each head's + data, so the content is (data + one scale) per K/V side.""" + mode = spec.kv_quant_mode + if spec.state_content_bytes is not None or not mode.is_per_token_head: + return spec + hs_k, hs_v = spec.head_size, spec.head_size_v + if mode == KVQuantMode.INT4_PER_TOKEN_HEAD: + hs_k, hs_v = hs_k // 2, hs_v // 2 + scale_bytes = get_dtype_size(torch.float32) + content = (hs_k + hs_v) * get_dtype_size(spec.dtype) + 2 * scale_bytes + return replace(spec, state_content_bytes=content) + supported_dtypes: ClassVar[list[torch.dtype]] = [ torch.float16, torch.bfloat16, diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py index ab9e375b0ae6..5d743be886ae 100644 --- a/vllm/v1/attention/backends/turboquant_attn.py +++ b/vllm/v1/attention/backends/turboquant_attn.py @@ -19,7 +19,7 @@ import contextlib import functools import math -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, ClassVar import torch @@ -137,6 +137,21 @@ class TurboQuantAttentionBackend(AttentionBackend): "turboquant_3bit_nc", ] + @classmethod + def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec: + """TurboQuant packs K+V into one slot per head.""" + if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_turboquant: + return spec + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + + # KVQuantMode member names mirror the preset strings. + tq = TurboQuantConfig.from_cache_dtype( + spec.kv_quant_mode.name.lower(), spec.head_size + ) + return replace(spec, state_content_bytes=tq.slot_size_aligned) + @staticmethod def get_name() -> str: return "TURBOQUANT" diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index ed494ce2dae4..5c9ab386d4e0 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -214,6 +214,87 @@ def _cp_gather_indexer_quant_cache_kernel( tl.store(dst_k_ptr + offset, val, mask=valid_block) +@triton.jit(do_not_specialize=["num_batches"]) +def _cp_gather_indexer_quant_cache_gfx950_kernel( + kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B] + # [n_blks, blk_size, head_dim] + kv_cache_scale_ptr, # [n_blks, blk_size] + k_fp8_ptr, # [num_tokens, head_dim] + k_scale_ptr, # [num_tokens] + block_table_ptr, # [batch_size, block_table_stride] + cu_seqlen_ptr, # [batch_size + 1] + token_to_seq_ptr, # [num_tokens] + block_size, + block_table_stride, + kv_cache_stride, + kv_cache_scale_stride, + LAYOUT: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_TILE_SIZE: tl.constexpr, + HEAD_TILE_SIZE: tl.constexpr, + num_batches, + BLOCK_TABLE_WIDTH: tl.constexpr, + NUM_BLOCKS: tl.constexpr, +): + tid = tl.program_id(0) + offset = tl.arange(0, HEAD_DIM) + batch_id = tl.load(token_to_seq_ptr + tid) + valid_batch = (batch_id >= 0) & (batch_id < num_batches) + safe_batch_id = tl.where(valid_batch, batch_id, 0) + batch_start = tl.load(cu_seqlen_ptr + safe_batch_id, mask=valid_batch, other=0) + batch_end = tl.load(cu_seqlen_ptr + safe_batch_id + 1, mask=valid_batch, other=0) + batch_offset = tid - batch_start + valid_token = valid_batch & (tid >= batch_start) & (tid < batch_end) + if not valid_token: + return + block_table_id = batch_offset // block_size + block_offset = batch_offset % block_size + valid_block_table = ( + valid_token + & (block_table_id >= 0) + & (block_table_id < BLOCK_TABLE_WIDTH) + & (block_offset >= 0) + & (block_offset < block_size) + ) + safe_block_table_id = tl.where(valid_block_table, block_table_id, 0) + block_table_offset = safe_batch_id * block_table_stride + safe_block_table_id + block_id = tl.load( + block_table_ptr + block_table_offset, mask=valid_block_table, other=-1 + ) + valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS) + # The packed KV layout makes per-block strides large + # enough that block_id * stride can exceed 32-bit range. + safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64) + safe_block_offset = tl.where(valid_block, block_offset, 0) + tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE + if LAYOUT == "SHUFFLE": + src_cache_offset = ( + safe_block_id * kv_cache_stride + + (safe_block_offset // BLOCK_TILE_SIZE) * HEAD_DIM * BLOCK_TILE_SIZE + + tiled_block_offset * HEAD_TILE_SIZE + ) + else: + src_cache_offset = ( + safe_block_id * kv_cache_stride + safe_block_offset * HEAD_DIM + ) + src_scale_offset = safe_block_id * kv_cache_scale_stride + safe_block_offset + dst_offset = tid * HEAD_DIM + src_scale_ptr = kv_cache_scale_ptr + src_scale_offset + src_cache_ptr = kv_cache_ptr + src_cache_offset + dst_k_ptr = k_fp8_ptr + dst_offset + scale_val = tl.load(src_scale_ptr, mask=valid_block, other=0.0) + tl.store(k_scale_ptr + tid, scale_val) + if LAYOUT == "SHUFFLE": + tiled_src_offset = ( + offset // HEAD_TILE_SIZE * HEAD_TILE_SIZE * BLOCK_TILE_SIZE + + offset % HEAD_TILE_SIZE + ) + else: + tiled_src_offset = offset + val = tl.load(src_cache_ptr + tiled_src_offset) + tl.store(dst_k_ptr + offset, val, mask=valid_block) + + def cp_gather_indexer_k_quant_cache_triton( k_cache: torch.Tensor, # [num_blocks, block_size, head_dim + 4] k_fp8: torch.Tensor, @@ -237,7 +318,7 @@ def cp_gather_indexer_k_quant_cache_triton( grid = (num_tokens,) k_fp8_scale = k_fp8_scale.view(torch.float32) layout = "NORMAL" if block_size == 1 else "SHUFFLE" - _cp_gather_indexer_quant_cache_kernel[grid]( + kernel_args = ( k_cache_value, k_cache_scale, k_fp8, @@ -253,11 +334,22 @@ def cp_gather_indexer_k_quant_cache_triton( head_dim, block_tile_size, head_tile_size, - num_tokens, - cu_seqlen.shape[0] - 1, - block_table.shape[1], - num_blocks, ) + if _ON_GFX950: + _cp_gather_indexer_quant_cache_gfx950_kernel[grid]( + *kernel_args, + cu_seqlen.shape[0] - 1, + block_table.shape[1], + num_blocks, + ) + else: + _cp_gather_indexer_quant_cache_kernel[grid]( + *kernel_args, + num_tokens, + cu_seqlen.shape[0] - 1, + block_table.shape[1], + num_blocks, + ) # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_attention.py#L156 @@ -563,21 +655,14 @@ def rocm_fp8_mqa_logits( Logits tensor of shape [M, N], dtype `torch.float32`. """ - # TODO(ganyi): Temporarily workaround, will remove the module check and reference - # path after aiter merge this kernel into main from vllm._aiter_ops import rocm_aiter_ops k_fp8, scale = kv - # Temporarily route gfx942 to the vendored ROCm/aiter#3257 workaround. - # Remove this branch once vLLM bumps AITER to a version that includes - # ROCm/aiter#3257. if _ON_GFX942 and rocm_aiter_ops.is_enabled(): - from vllm.v1.attention.ops.triton_fp8_mqa_logits import ( - fp8_mqa_logits_gfx942, - ) + from aiter.ops.flydsl import flydsl_fp8_mqa_logits - return fp8_mqa_logits_gfx942( + return flydsl_fp8_mqa_logits( q, k_fp8, scale, weights, cu_seqlen_ks, cu_seqlen_ke ) @@ -1218,6 +1303,90 @@ def _sparse_attn_prefill_ragged_kernel( ) +@triton.jit +def _decode_e8m0_scales_triton(encoded_scales): + scale_bits = encoded_scales.to(tl.int32) << 23 + scale_bits = tl.where(encoded_scales == 0, 1 << 22, scale_bits) + return scale_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + CHUNK_START: tl.constexpr, + CHUNK_SIZE: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, +): + offsets = CHUNK_START + tl.arange(0, CHUNK_SIZE) + x_uint8 = tl.load( + token_data_ptr[:, None] + offsets[None, :], + mask=valid[:, None], + other=0, + ) + scale_offsets = CHUNK_START // 64 + tl.arange(0, CHUNK_SIZE // 64) + encoded_scales = tl.load( + token_scale_ptr[:, None] + scale_offsets[None, :], + mask=valid[:, None], + other=127, + ) + scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, CHUNK_SIZE // 64, 64)) + scales = tl.reshape(scales, (BLOCK_K, CHUNK_SIZE)) + if IS_FNUZ: + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32) + else: + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + value = (x_f32 * scales).to(tl.bfloat16) + zero = tl.zeros((BLOCK_K, CHUNK_SIZE), dtype=tl.bfloat16) + return tl.where(valid[:, None], value, zero) + + +@triton.jit +def _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, +): + tail_offsets = 384 + tl.arange(0, 128) + nope_mask = tail_offsets < NOPE_DIM + x_uint8 = tl.load( + token_data_ptr[:, None] + tail_offsets[None, :], + mask=valid[:, None] & nope_mask[None, :], + other=0, + ) + scale_offsets = 6 + tl.arange(0, 2) + scale_mask = scale_offsets < NOPE_DIM // 64 + encoded_scales = tl.load( + token_scale_ptr[:, None] + scale_offsets[None, :], + mask=valid[:, None] & scale_mask[None, :], + other=127, + ) + scales = _decode_e8m0_scales_triton(encoded_scales) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, 2, 64)) + scales = tl.reshape(scales, (BLOCK_K, 128)) + if IS_FNUZ: + x_f32 = x_uint8.to(tl.float8e4b8, bitcast=True).to(tl.bfloat16).to(tl.float32) + else: + x_f32 = x_uint8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + nope = (x_f32 * scales).to(tl.bfloat16) + + rope_ptr = (token_data_ptr + NOPE_DIM).to(tl.pointer_type(tl.bfloat16)) + rope = tl.load( + rope_ptr[:, None] + (tail_offsets[None, :] - NOPE_DIM), + mask=valid[:, None] & ~nope_mask[None, :], + other=0.0, + ) + value = tl.where(nope_mask[None, :], nope, rope) + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + return tl.where(valid[:, None], value, zero) + + @triton.jit def _sparse_attn_decode_ragged_kernel( q_ptr, @@ -1692,6 +1861,407 @@ def _sparse_attn_decode_partial_kernel( ) +@triton.jit +def _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + cache_ptr, + slot, + valid, + cache_stride0, + scale: tl.constexpr, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + BLOCK_SIZE: tl.constexpr, + NOPE_DIM: tl.constexpr, + BLOCK_K: tl.constexpr, + IS_FNUZ: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, +): + safe_slot = tl.where(valid, slot, 0) + block_idx = safe_slot // BLOCK_SIZE + pos_in_block = safe_slot % BLOCK_SIZE + cache_block_ptr = cache_ptr + block_idx.to(tl.int64) * cache_stride0 + token_data_ptr = cache_block_ptr + pos_in_block * 576 + token_scale_ptr = cache_block_ptr + BLOCK_SIZE * 576 + pos_in_block * 8 + k_nope_0a = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 0, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_0b = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 128, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_nope_1 = _load_fp8_ds_mla_gfx950_nope_exact_chunk( + token_data_ptr, + token_scale_ptr, + valid, + 256, + 128, + BLOCK_K, + IS_FNUZ, + ) + k_tail = _load_fp8_ds_mla_gfx950_tail128( + token_data_ptr, + token_scale_ptr, + valid, + NOPE_DIM, + BLOCK_K, + IS_FNUZ, + ) + if not TRUST_EXTRA_CACHE_NAN_FREE: + zero = tl.zeros((BLOCK_K, 128), dtype=tl.bfloat16) + k_nope_0a = tl.where(k_nope_0a == k_nope_0a, k_nope_0a, zero) + k_nope_0b = tl.where(k_nope_0b == k_nope_0b, k_nope_0b, zero) + k_nope_1 = tl.where(k_nope_1 == k_nope_1, k_nope_1, zero) + k_tail = tl.where(k_tail == k_tail, k_tail, zero) + k_nope_0 = tl.cat(k_nope_0a, k_nope_0b, dim=1) + k_tail_256 = tl.cat(k_nope_1, k_tail, dim=1) + k_combined = tl.cat(k_nope_0, k_tail_256, dim=1) + + scores = tl.dot(q_combined, tl.trans(k_combined)) + scores *= scale * 1.4426950408889634 + scores = tl.where( + head_mask[:, None] & valid[None, :], + scores, + -3.4028234663852886e38, + ) + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0) + l_new = l_i * alpha + tl.sum(p, axis=1) + p_bf16 = p.to(k_nope_0a.dtype) + acc_nope_0a = acc_nope_0a * alpha[:, None] + tl.dot(p_bf16, k_nope_0a) + acc_nope_0b = acc_nope_0b * alpha[:, None] + tl.dot(p_bf16, k_nope_0b) + acc_nope_1 = acc_nope_1 * alpha[:, None] + tl.dot(p_bf16, k_nope_1) + acc_tail = acc_tail * alpha[:, None] + tl.dot(p_bf16, k_tail) + return ( + m_new, + l_new, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) + + +@triton.jit +def _sparse_attn_decode_gfx950_partial_kernel( + q_ptr, + main_cache_ptr, + main_indices_ptr, + main_indptr_ptr, + extra_cache_ptr, + extra_indices_ptr, + extra_indptr_ptr, + part_m_ptr, + part_l_ptr, + part_acc_ptr, + q_stride0: tl.constexpr, + q_stride1: tl.constexpr, + main_cache_stride0: tl.constexpr, + extra_cache_stride0: tl.constexpr, + main_num_rows, + extra_num_rows, + MAIN_BLOCK_SIZE: tl.constexpr, + EXTRA_BLOCK_SIZE: tl.constexpr, + scale: tl.constexpr, + num_heads: tl.constexpr, + HAS_EXTRA: tl.constexpr, + NOPE_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + IS_FNUZ_MAIN: tl.constexpr, + IS_FNUZ_EXTRA: tl.constexpr, + TRUST_EXTRA_CACHE_NAN_FREE: tl.constexpr, + ADAPTIVE_SPLITS: tl.constexpr, + ONE_WAVE_SPLITS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + query_idx = tl.program_id(0) + split_id = tl.program_id(1) + pid_h = tl.program_id(2) + + tl.static_assert(NOPE_DIM == 448) + tl.static_assert(ROPE_DIM == 64) + + head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + if num_heads % BLOCK_H == 0: + head_mask = tl.full((BLOCK_H,), True, tl.int1) + else: + head_mask = head_offsets < num_heads + neg_large = -3.4028234663852886e38 + + if ADAPTIVE_SPLITS: + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + if HAS_EXTRA: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + else: + extra_start = 0 + extra_len = 0 + split4_span: tl.constexpr = 4 * BLOCK_K + split4_iters = (main_len + split4_span - 1) // split4_span + split4_iters += (extra_len + split4_span - 1) // split4_span + use_four_splits = split4_iters <= 3 + work_splits = NUM_SPLITS + if ONE_WAVE_SPLITS > 4 and ONE_WAVE_SPLITS < NUM_SPLITS: + one_wave_span: tl.constexpr = ONE_WAVE_SPLITS * BLOCK_K + one_wave_iters = (main_len + one_wave_span - 1) // one_wave_span + one_wave_iters += (extra_len + one_wave_span - 1) // one_wave_span + work_splits = tl.where(one_wave_iters <= 3, ONE_WAVE_SPLITS, work_splits) + work_splits = tl.where(use_four_splits, 4, work_splits) + if split_id >= work_splits: + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets + tl.store(part_m_ptr + pm_base, neg_large, mask=head_mask) + tl.store(part_l_ptr + pm_base, 0.0, mask=head_mask) + return + else: + work_splits = NUM_SPLITS + + nope_offsets_0a = tl.arange(0, 128) + nope_offsets_0b = 128 + tl.arange(0, 128) + nope_offsets_0 = tl.arange(0, 256) + tail_offsets = 256 + tl.arange(0, 256) + nope_offsets_1 = 256 + tl.arange(0, 128) + tail_offsets_128 = 384 + tl.arange(0, 128) + + q_row_ptr = q_ptr + query_idx * q_stride0 + head_offsets[:, None] * q_stride1 + q_nope_0 = tl.load( + q_row_ptr + nope_offsets_0[None, :], + mask=head_mask[:, None], + other=0.0, + ) + q_tail = tl.load( + q_row_ptr + tail_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) + q_combined = tl.cat(q_nope_0, q_tail, dim=1) + + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc_nope_0a = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_nope_0b = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_nope_1 = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + acc_tail = tl.zeros((BLOCK_H, 128), dtype=tl.float32) + k_offsets = tl.arange(0, BLOCK_K) + + if not ADAPTIVE_SPLITS: + main_start = tl.load(main_indptr_ptr + query_idx) + main_end = tl.load(main_indptr_ptr + query_idx + 1) + main_len = main_end - main_start + main_chunk = (main_len + work_splits - 1) // work_splits + main_lo = split_id * main_chunk + main_hi = tl.minimum(main_lo + main_chunk, main_len) + + for k_start in tl.range( + main_lo, + main_hi, + BLOCK_K, + num_stages=NUM_STAGES, + ): + k_pos = k_start + k_offsets + in_range = k_pos < main_hi + slot = tl.load(main_indices_ptr + main_start + k_pos, mask=in_range, other=-1) + valid = in_range & (slot >= 0) & (slot < main_num_rows) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + main_cache_ptr, + slot, + valid, + main_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + MAIN_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_MAIN, + False, + ) + + if HAS_EXTRA: + if not ADAPTIVE_SPLITS: + extra_start = tl.load(extra_indptr_ptr + query_idx) + extra_end = tl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + extra_chunk = (extra_len + work_splits - 1) // work_splits + extra_lo = split_id * extra_chunk + extra_hi = tl.minimum(extra_lo + extra_chunk, extra_len) + + outer_block_k: tl.constexpr = 2 * BLOCK_K + outer_k_offsets = tl.arange(0, outer_block_k) + extra_hi_full = ( + extra_lo + ((extra_hi - extra_lo) // outer_block_k) * outer_block_k + ) + for k_start in tl.range( + extra_lo, + extra_hi_full, + outer_block_k, + num_stages=NUM_STAGES, + ): + slot = tl.load(extra_indices_ptr + extra_start + k_start + outer_k_offsets) + valid = (slot >= 0) & (slot < extra_num_rows) + slot_pairs = tl.trans(tl.reshape(slot, (2, BLOCK_K))) + valid_pairs = tl.trans(tl.reshape(valid, (2, BLOCK_K))) + slot_lo, slot_hi = tl.split(slot_pairs) + valid_lo, valid_hi = tl.split(valid_pairs) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot_lo, + valid_lo, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot_hi, + valid_hi, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + for tail_idx in tl.static_range(2): + tail_start = extra_hi_full + tail_idx * BLOCK_K + if tail_start < extra_hi: + k_pos = tail_start + k_offsets + in_range = k_pos < extra_hi + slot = tl.load( + extra_indices_ptr + extra_start + k_pos, + mask=in_range, + other=-1, + ) + valid = in_range & (slot >= 0) & (slot < extra_num_rows) + ( + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + ) = _sparse_attn_decode_gfx950_partial_loaded_tile( + q_combined, + extra_cache_ptr, + slot, + valid, + extra_cache_stride0, + scale, + head_mask, + m_i, + l_i, + acc_nope_0a, + acc_nope_0b, + acc_nope_1, + acc_tail, + EXTRA_BLOCK_SIZE, + NOPE_DIM, + BLOCK_K, + IS_FNUZ_EXTRA, + TRUST_EXTRA_CACHE_NAN_FREE, + ) + + pm_base = (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets + m_store = tl.where(l_i > 0.0, m_i * 0.6931471805599453, neg_large) + tl.store(part_m_ptr + pm_base, m_store, mask=head_mask) + tl.store(part_l_ptr + pm_base, l_i, mask=head_mask) + acc_base = part_acc_ptr + ( + (query_idx * NUM_SPLITS + split_id) * num_heads + head_offsets[:, None] + ) * (NOPE_DIM + ROPE_DIM) + tl.store( + acc_base + nope_offsets_0a[None, :], + acc_nope_0a, + mask=head_mask[:, None], + ) + tl.store( + acc_base + nope_offsets_0b[None, :], + acc_nope_0b, + mask=head_mask[:, None], + ) + tl.store( + acc_base + nope_offsets_1[None, :], + acc_nope_1, + mask=head_mask[:, None], + ) + tl.store( + acc_base + tail_offsets_128[None, :], + acc_tail, + mask=head_mask[:, None], + ) + + @triton.jit def _sparse_attn_decode_reduce_kernel( part_m_ptr, @@ -1708,6 +2278,7 @@ def _sparse_attn_decode_reduce_kernel( pa_stride_h, num_heads, HAS_ATTN_SINK: tl.constexpr, + ADAPTIVE_SPLITS: tl.constexpr, COMB_DIM: tl.constexpr, BLOCK_H: tl.constexpr, NUM_SPLITS: tl.constexpr, @@ -1776,17 +2347,27 @@ def _sparse_attn_decode_reduce_kernel( other=neg_large, ) w_s = tl.exp(m_s - m_final) + if ADAPTIVE_SPLITS: + active_split = m_s > neg_large + w_s = tl.where(head_mask & active_split, w_s, 0.0) acc_base = ( part_acc_ptr + query_idx * pa_stride0 + s * pa_stride_s + head_offsets[:, None] * pa_stride_h ) - acc_s = tl.load( - acc_base + comb_offsets[None, :], - mask=head_mask[:, None], - other=0.0, - ) + if ADAPTIVE_SPLITS: + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None] & active_split[:, None], + other=0.0, + ) + else: + acc_s = tl.load( + acc_base + comb_offsets[None, :], + mask=head_mask[:, None], + other=0.0, + ) acc += w_s[:, None] * acc_s out = tl.where(l_final[:, None] > 0.0, acc / denom[:, None], 0.0) @@ -1993,6 +2574,49 @@ def _decode_num_splits( return best_splits +def _decode_gfx950_num_splits( + num_queries: int, + heads_blocks: int, + avg_main_len: float = 0.0, + avg_extra_len: float = 0.0, + block_k: int = 32, +) -> int: + base = max(1, num_queries * heads_blocks) + cu = max(1, _decode_cu_count()) + target_workgroups = 2 * cu + num_splits = min( + 32, + max( + 1, + math.ceil(target_workgroups / base), + ), + ) + if ( + base >= 16 + and num_splits > 4 + and _decode_partial_iters(avg_main_len, avg_extra_len, 4, block_k) <= 3 + ): + return 4 + if 16 <= base < 64: + one_wave_splits = max(1, cu // base) + one_wave_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, one_wave_splits, block_k + ) + target_waves = 1 if one_wave_iters <= 9 else 2 + num_splits = min(num_splits, max(1, target_waves * cu // base)) + if base >= 16 and num_splits > 1: + target_waves = (base * num_splits + cu - 1) // cu + target_iters = _decode_partial_iters( + avg_main_len, avg_extra_len, num_splits, block_k + ) + for splits in range(1, num_splits): + waves = (base * splits + cu - 1) // cu + iters = _decode_partial_iters(avg_main_len, avg_extra_len, splits, block_k) + if waves == target_waves and iters == target_iters: + return splits + return num_splits + + def _rocm_sparse_attn_decode_ragged_triton( q: torch.Tensor, main_cache: torch.Tensor, @@ -2005,6 +2629,9 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_cache: torch.Tensor | None = None, extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, + extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert main_cache.ndim == 3, ( @@ -2045,6 +2672,10 @@ def _rocm_sparse_attn_decode_ragged_triton( and extra_indices is not None and extra_indptr is not None ) + assert not extra_cache_nan_free or (_ON_GFX950 and has_extra), ( + "extra_cache_nan_free requires a gfx950 compressed cache with trusted " + "canonical-writer provenance" + ) if has_extra: assert extra_cache is not None assert extra_indices is not None @@ -2066,7 +2697,16 @@ def _rocm_sparse_attn_decode_ragged_triton( extra_indptr = torch.zeros(num_queries + 1, device=q.device, dtype=torch.int32) block_h = 16 - out = torch.empty_like(q, dtype=torch.bfloat16) + if out is None: + out = torch.empty_like(q, dtype=torch.bfloat16) + else: + assert out.shape == q.shape, f"expected out shape {q.shape}, got {out.shape}" + assert out.device == q.device, ( + f"expected out on device {q.device}, got {out.device}" + ) + assert out.dtype == torch.bfloat16, ( + f"expected out dtype {torch.bfloat16}, got {out.dtype}" + ) heads_blocks = triton.cdiv(num_heads, block_h) nope_block = triton.next_power_of_2(nope_head_dim) comb_dim = nope_head_dim + rope_head_dim @@ -2110,14 +2750,35 @@ def _rocm_sparse_attn_decode_ragged_triton( return out block_k = 32 # KV tokens walked per split-K iteration. Tuned on gfx950. - # Average per-query segment lengths, read sync-free from the ragged index - # sizes, let the split heuristic avoid over-splitting - # main_indices/extra_indices are flat [nnz] int32. - inv_q = 1.0 / max(1, num_queries) - avg_main_len = main_indices.numel() * inv_q - avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 - num_splits = _decode_num_splits( - num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + if _ON_GFX950: + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_gfx950_num_splits( + num_queries, + heads_blocks, + avg_main_len, + avg_extra_len, + block_k, + ) + else: + # Average per-query segment lengths, read sync-free from the ragged + # index sizes, let the split heuristic avoid over-splitting. + inv_q = 1.0 / max(1, num_queries) + avg_main_len = main_indices.numel() * inv_q + avg_extra_len = (extra_indices.numel() * inv_q) if has_extra else 0.0 + num_splits = _decode_num_splits( + num_queries, heads_blocks, avg_main_len, avg_extra_len, block_k + ) + + base_workgroups = num_queries * heads_blocks + adaptive_splits = ( + _ON_GFX950 and adaptive_splits and base_workgroups >= 16 and num_splits > 4 + ) + one_wave_splits = ( + max(1, _decode_cu_count() // base_workgroups) + if adaptive_splits and 16 <= base_workgroups < 64 + else num_splits ) part_m = torch.empty( @@ -2130,48 +2791,88 @@ def _rocm_sparse_attn_decode_ragged_triton( device=q.device, ) - _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( - q, - main_cache, - main_indices, - main_indptr, - extra_cache, - extra_indices, - extra_indptr, - part_m, - part_l, - part_acc, - q.stride(0), - q.stride(1), - main_cache.stride(0), - extra_cache.stride(0), - part_m.stride(0), - part_m.stride(1), - part_acc.stride(0), - part_acc.stride(1), - part_acc.stride(2), - main_cache.shape[0] * main_cache.shape[1], - extra_cache.shape[0] * extra_cache.shape[1], - main_cache.shape[1], - extra_cache.shape[1], - scale, - num_heads, - HAS_EXTRA=has_extra, - NOPE_DIM=nope_head_dim, - NOPE_BLOCK=nope_block, - ROPE_DIM=rope_head_dim, - # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). - # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). - # Reading both with a single IS_FNUZ would decode one of them with the - # wrong FNUZ/OCP scale ratio (~1.87×). - IS_FNUZ_MAIN=is_fnuz, - IS_FNUZ_EXTRA=False, - BLOCK_H=block_h, - BLOCK_K=block_k, - NUM_SPLITS=num_splits, - NUM_STAGES=1, - num_warps=4, - ) + if _ON_GFX950: + _sparse_attn_decode_gfx950_partial_kernel[ + (num_queries, num_splits, heads_blocks) + ]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + ROPE_DIM=rope_head_dim, + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + TRUST_EXTRA_CACHE_NAN_FREE=extra_cache_nan_free, + ADAPTIVE_SPLITS=adaptive_splits, + ONE_WAVE_SPLITS=one_wave_splits, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + waves_per_eu=0, + ) + else: + _sparse_attn_decode_partial_kernel[(num_queries, num_splits, heads_blocks)]( + q, + main_cache, + main_indices, + main_indptr, + extra_cache, + extra_indices, + extra_indptr, + part_m, + part_l, + part_acc, + q.stride(0), + q.stride(1), + main_cache.stride(0), + extra_cache.stride(0), + part_m.stride(0), + part_m.stride(1), + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + main_cache.shape[0] * main_cache.shape[1], + extra_cache.shape[0] * extra_cache.shape[1], + main_cache.shape[1], + extra_cache.shape[1], + scale, + num_heads, + HAS_EXTRA=has_extra, + NOPE_DIM=nope_head_dim, + NOPE_BLOCK=nope_block, + ROPE_DIM=rope_head_dim, + # main_cache = swa_k_cache (C++ encoder, FNUZ on gfx942 / OCP on gfx950). + # extra_cache = compressed kv_cache (Triton encoder, OCP everywhere). + # Reading both with a single IS_FNUZ would decode one of them with the + # wrong FNUZ/OCP scale ratio (~1.87×). + IS_FNUZ_MAIN=is_fnuz, + IS_FNUZ_EXTRA=False, + BLOCK_H=block_h, + BLOCK_K=block_k, + NUM_SPLITS=num_splits, + NUM_STAGES=1, + num_warps=4, + ) _sparse_attn_decode_reduce_kernel[(num_queries, num_heads)]( part_m, @@ -2188,6 +2889,7 @@ def _rocm_sparse_attn_decode_ragged_triton( part_acc.stride(2), num_heads, HAS_ATTN_SINK=has_attn_sink, + ADAPTIVE_SPLITS=adaptive_splits, COMB_DIM=comb_dim, BLOCK_H=1, NUM_SPLITS=num_splits, @@ -2213,6 +2915,9 @@ def _rocm_sparse_attn_decode_triton( main_ragged_indptr: torch.Tensor | None = None, extra_ragged_indices: torch.Tensor | None = None, extra_ragged_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, + extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> torch.Tensor: if main_ragged_indices is None or main_ragged_indptr is None: main_ragged_indices, main_ragged_indptr = build_ragged_indices_from_dense( @@ -2248,6 +2953,9 @@ def _rocm_sparse_attn_decode_triton( extra_cache=extra_cache, extra_indices=extra_ragged_indices, extra_indptr=extra_ragged_indptr, + out=out, + extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) @@ -2319,6 +3027,8 @@ def rocm_sparse_attn_decode( nope_head_dim: int, rope_head_dim: int, output: torch.Tensor, + extra_cache_nan_free: bool = False, + adaptive_splits: bool = False, ) -> None: assert swa_k_cache.dtype == torch.uint8, ( "ROCm Triton sparse decode expects uint8 fp8_ds_mla SWA cache, " @@ -2348,6 +3058,7 @@ def rocm_sparse_attn_decode( if topk_indices is not None: extra_indices = topk_indices.reshape(topk_indices.shape[0], -1) + direct_out = output if _ON_GFX950 and output.dtype == torch.bfloat16 else None attn_out = _rocm_sparse_attn_decode_triton( q=q, main_cache=swa_k_cache, @@ -2364,5 +3075,9 @@ def rocm_sparse_attn_decode( main_ragged_indptr=swa_ragged_indptr, extra_ragged_indices=topk_ragged_indices, extra_ragged_indptr=topk_ragged_indptr, + out=direct_out, + extra_cache_nan_free=extra_cache_nan_free, + adaptive_splits=adaptive_splits, ) - output.copy_(attn_out.to(output.dtype)) + if direct_out is None: + output.copy_(attn_out.to(output.dtype)) diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index 86a0b6ee9d31..b7aeb8c6a769 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -20,6 +20,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import direct_register_custom_op @@ -50,7 +51,12 @@ def flash_attn_maxseqlen_wrapper( cu_seqlens = torch.arange( 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device ) - max_seqlen = q_len if max_seqlen is None else max_seqlen.item() + if max_seqlen is None: + max_seqlen = q_len + else: + # `flash_attn_varlen_func` needs a Python int for kernel launch bounds. + with gpu_sync_allowed(): + max_seqlen = max_seqlen.item() q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]) output = flash_attn_varlen_func( @@ -159,7 +165,12 @@ def triton_attn_wrapper( cu_seqlens = torch.arange( 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device ) - max_seqlen = q_len if max_seqlen is None else max_seqlen.item() + if max_seqlen is None: + max_seqlen = q_len + else: + # `context_attention_fwd` needs a Python int. + with gpu_sync_allowed(): + max_seqlen = max_seqlen.item() q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]) output = torch.empty_like(q) @@ -261,7 +272,9 @@ def torch_sdpa_wrapper( outputs = [] - lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + # `torch.split` needs Python int sizes. + with gpu_sync_allowed(): + lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() q_chunks = torch.split(q, lens, dim=1) k_chunks = torch.split(k, lens, dim=1) v_chunks = torch.split(v, lens, dim=1) @@ -337,7 +350,10 @@ def flashinfer_wrapper( batch_offsets_qko = cu_seqlens[:cu_seqlength].view(-1, 1, 1, 1) batch_offsets_v = cu_seqlens[cu_seqlength:].view(-1, 1, 1, 1) sequence_lengths = sequence_lengths.view(-1, 1, 1, 1) - max_seqlen = max_seqlen.item() + # `cudnn_batch_prefill_with_kv_cache` needs Python ints for the + # max-token-per-seq bounds. + with gpu_sync_allowed(): + max_seqlen = max_seqlen.item() output, _ = cudnn_batch_prefill_with_kv_cache( q, diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 00f2aed5c506..b10cb3497737 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -38,6 +38,7 @@ class AttentionSelectorConfig(NamedTuple): use_kv_connector: bool = False use_pcp: bool = False use_adaptive_verification: bool = False + use_dcp: bool = False def __repr__(self): return ( @@ -56,7 +57,8 @@ def __repr__(self): f"use_batch_invariant={self.use_batch_invariant}, " f"use_kv_connector={self.use_kv_connector}, " f"use_adaptive_verification={self.use_adaptive_verification}, " - f"use_pcp={self.use_pcp})" + f"use_pcp={self.use_pcp}, " + f"use_dcp={self.use_dcp})" ) @@ -168,6 +170,7 @@ def get_attn_backend( use_kv_connector=use_kv_connector, use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1, use_adaptive_verification=use_adaptive_verification, + use_dcp=vllm_config.parallel_config.decode_context_parallel_size > 1, ) # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over diff --git a/vllm/v1/core/encoder_cache_manager.py b/vllm/v1/core/encoder_cache_manager.py index 02133ff5b888..8d2a81a11b13 100644 --- a/vllm/v1/core/encoder_cache_manager.py +++ b/vllm/v1/core/encoder_cache_manager.py @@ -269,7 +269,9 @@ def get_freed_mm_hashes(self) -> list[str]: encoder outputs can be removed from their caches. The internal list is cleared after this call. """ - freed = self.freed + # An entry evicted early in the scheduling pass can be allocated again + # later in the same pass. Keep its worker-side tensor in that case. + freed = [mm_hash for mm_hash in self.freed if mm_hash not in self.cached] self.freed = [] return freed diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 321bbb0a76ac..0829ec0d2be8 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -4,9 +4,8 @@ from collections.abc import Sequence from typing import NamedTuple -from vllm import envs from vllm.logger import init_logger -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( @@ -45,16 +44,18 @@ def _validate_prefix_cache_retention_interval( isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec)) for g in kv_cache_config.kv_cache_groups ): + if retention_interval == 0: + return raise ValueError( - "VLLM_PREFIX_CACHE_RETENTION_INTERVAL is set but this model has " + "prefix_cache_retention_interval is set but this model has " "no sliding-window or Mamba KV cache group, so retention has no " - "effect. Unset it (it only applies to sliding-window and Mamba " + "effect. Set it to 0 (it only applies to sliding-window and Mamba " "attention)." ) if retention_interval < 0 or retention_interval % scheduler_block_size != 0: raise ValueError( - f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL ({retention_interval}) " + f"prefix_cache_retention_interval ({retention_interval}) " "must be non-negative and a multiple of scheduler_block_size " f"({scheduler_block_size})." ) @@ -80,6 +81,7 @@ def __init__( scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, + num_prefill_lookahead: int = 0, ): self.kv_cache_config = kv_cache_config self.max_model_len = max_model_len @@ -91,6 +93,7 @@ def __init__( for g in kv_cache_config.kv_cache_groups ) self.scheduler_block_size = scheduler_block_size + self.num_reprefillable_tokens = max(0, num_prefill_lookahead - 1) self.block_pool = BlockPool( num_gpu_blocks=kv_cache_config.num_blocks, @@ -108,6 +111,28 @@ def __init__( if use_eagle and not self.eagle_group_ids: self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups))) + # During chunked prefill with EAGLE, the single next prefill lookahead + # token past the chunk boundary is combined with the final hidden state + # and written to the KV cache. Therefore, the final chunk token must be + # excluded from prefix cache hits to prevent requests from acquiring the + # KV cache slot polluted with the next prefill token, which may or may not + # be present after the matching prefix. The last-block drop handles this + # edge case. During multi-module MTP, the issue generalizes to a prefill + # lookahead of num_speculative_tokens, so the dropped tail must be large + # enough to contain them. Hits land on scheduler-block boundaries (see + # `_cache_hit_alignment_tokens`), so the excluded tail is + # scheduler_block_size, not the group's own block size. + if ( + enable_caching + and self.eagle_group_ids + and scheduler_block_size < num_prefill_lookahead + ): + raise ValueError( + f"Multi-module MTP with prefix caching requires scheduler_block_size" + f" (={scheduler_block_size}) >= num_speculative_tokens" + f" (={num_prefill_lookahead})." + ) + self.single_type_managers = tuple( get_manager_for_kv_cache_spec( kv_cache_spec=kv_cache_group.kv_cache_spec, @@ -127,7 +152,7 @@ def __init__( # A positive retention interval must be a multiple of the base hit granularity # (``scheduler_block_size``) to land on real cache-hit boundaries. # 0 = keep only the latest replay boundary; None = dense; - self.retention_interval = envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL + self.retention_interval = kv_cache_config.prefix_cache_retention_interval _validate_prefix_cache_retention_interval( self.retention_interval, self.scheduler_block_size, kv_cache_config ) @@ -286,9 +311,14 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: (including tokens that are already cached). """ for manager in self.single_type_managers: + # Only cache tokens with finalized KV. The last num_reprefillable_tokens + # tokens can be re-prefilled during multi-module MTP. + num_tokens_to_cache = max( + 0, num_computed_tokens - self.num_reprefillable_tokens + ) manager.cache_blocks( request, - num_computed_tokens, + num_tokens_to_cache, retention_interval=self.retention_interval, ) @@ -407,6 +437,7 @@ def __init__( scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, + num_prefill_lookahead: int = 0, ): super().__init__( kv_cache_config, @@ -420,6 +451,7 @@ def __init__( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) self.num_single_type_manager = len(self.single_type_managers) @@ -457,6 +489,7 @@ def __init__( scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, + num_prefill_lookahead: int = 0, ): super().__init__( kv_cache_config, @@ -470,6 +503,7 @@ def __init__( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec self.block_size = self.kv_cache_spec.block_size @@ -542,6 +576,7 @@ def __init__( scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, + num_prefill_lookahead: int = 0, ): super().__init__( kv_cache_config, @@ -555,6 +590,7 @@ def __init__( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) # hash_block_size: the block size used to compute block hashes. # The actual block size usually equals hash_block_size, but in cases where @@ -669,28 +705,37 @@ def verify_and_split_kv_cache_groups(self) -> None: for gid in group.group_ids: self.single_type_managers[gid].use_eagle = True - def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: + def _align_cacheable(self, num_tokens: int) -> int: + """Largest prefix of ``num_tokens`` a future cache hit could match. + + Hits are ``scheduler_block_size``-aligned (see + ``find_longest_cache_hit``) unless fine-grained partial hash hits are + enabled, in which case no rounding applies -- rounding even to + ``hash_block_size`` would re-register a privatized Mamba tail. + """ if self.enable_partial_hash_hits: - aligned_num_computed_tokens = num_computed_tokens - else: - # Cache hits in this coordinator are always a multiple of - # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``). - # Within an aligned region, SWA groups may only consult a subset of - # blocks per ``scheduler_block_size``-segment so the unused blocks - # also stay out of the prefix-cache hash map. - aligned_num_computed_tokens = ( - num_computed_tokens - // self.scheduler_block_size - * self.scheduler_block_size - ) + return num_tokens + return round_down(num_tokens, self.scheduler_block_size) + + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: + cached_num_computed_tokens = self._align_cacheable(num_computed_tokens) for manager in self.single_type_managers: - num_tokens_to_cache = aligned_num_computed_tokens + num_tokens_to_cache = cached_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop # it, so make that lookahead block eligible to be cached. - if manager.use_eagle and aligned_num_computed_tokens > 0: + if manager.use_eagle and cached_num_computed_tokens > 0: + # Only cache tokens with finalized KV. The last + # num_reprefillable_tokens tokens can be re-prefilled during + # multi-module MTP. + num_finalized_computed_tokens = max( + 0, num_computed_tokens - self.num_reprefillable_tokens + ) + cached_num_finalized_computed_tokens = self._align_cacheable( + num_finalized_computed_tokens + ) num_tokens_to_cache = min( - num_computed_tokens, - aligned_num_computed_tokens + manager.block_size, + num_finalized_computed_tokens, + cached_num_finalized_computed_tokens + manager.block_size, ) # The manager already knows the fine hit granularity # (``scheduler_block_size``); retention is passed separately so it @@ -880,6 +925,7 @@ def get_kv_cache_coordinator( scheduler_block_size: int, hash_block_size: int, metrics_collector: KVCacheMetricsCollector | None = None, + num_prefill_lookahead: int = 0, ) -> KVCacheCoordinator: if not enable_caching: return KVCacheCoordinatorNoPrefixCache( @@ -893,6 +939,7 @@ def get_kv_cache_coordinator( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) if len(kv_cache_config.kv_cache_groups) == 1: return UnitaryKVCacheCoordinator( @@ -907,6 +954,7 @@ def get_kv_cache_coordinator( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) return HybridKVCacheCoordinator( kv_cache_config, @@ -920,4 +968,5 @@ def get_kv_cache_coordinator( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index ca1fb73420a2..d1af91b65993 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -125,6 +125,7 @@ def __init__( max_in_flight_tokens: int | None = None, enable_caching: bool = True, use_eagle: bool = False, + num_prefill_lookahead: int = 0, log_stats: bool = False, enable_kv_cache_events: bool = False, dcp_world_size: int = 1, @@ -161,6 +162,7 @@ def __init__( scheduler_block_size=scheduler_block_size, hash_block_size=hash_block_size, metrics_collector=self.metrics_collector, + num_prefill_lookahead=num_prefill_lookahead, ) self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) self.block_pool = self.coordinator.block_pool @@ -241,7 +243,7 @@ def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int - ``shared_prefix_boundary``: the block-aligned token position of a shared prefix that a sparse-retention group (Mamba / sliding window) has not cached yet (Marconi-style APC), or 0 if none. - Pinned so ``VLLM_PREFIX_CACHE_RETENTION_INTERVAL`` does not drop + Pinned so sparse prefix-cache retention does not drop the junction and defeat cross-request reuse. """ # We skip finding the prefix cache hit when prefix caching is diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 57d6600b368e..d6e401d1184e 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1347,6 +1347,9 @@ def get_kv_cache_config_from_groups( num_blocks=1, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), ) # Determine how model runners should initialize the KV cache tensors. @@ -1406,6 +1409,9 @@ def get_kv_cache_config_from_groups( num_blocks=num_blocks, kv_cache_tensors=kv_cache_tensors, kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), ) @@ -1468,6 +1474,9 @@ def promoted_page_size_padded(spec: AttentionSpec, block_size: int) -> int | Non promoted_specs[layer_name] = replace_as( spec, target_cls, + # Promoted specs allocate blocks for all tokens and never free + # below the window, so the trailing-edge extension is moot. + drop=("extra_retained_tokens",), block_size=block_size, page_size_padded=promoted_page_size_padded(spec, block_size), ) @@ -2128,6 +2137,21 @@ def get_kv_cache_configs( # Check if the KV cache specs are registered correctly. # This is to prevent that some layers are initialized with unregistered specs. KVCacheSpecRegistry.check_kv_cache_spec_registry(merged_kv_cache_specs) + + # 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 + ) + for layer_name, layer_spec in merged_kv_cache_specs.items(): + if isinstance(layer_spec, SlidingWindowSpec): + merged_kv_cache_specs[layer_name] = replace( + layer_spec, extra_retained_tokens=extra_retained_tokens + ) + # Get global KV cache groups. This also handles spec unification for # hybrid models when disable_hybrid_kv_cache_manager is enabled. # After this call, merged_kv_cache_specs may be modified in-place. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 9d14d00edc56..e4a21328660a 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -247,6 +247,13 @@ def __init__( self.use_eagle = 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.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: if speculative_config.num_speculative_tokens_per_batch_size: @@ -256,6 +263,12 @@ 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 + ) # Create the KV cache manager. if hash_block_size is None: @@ -267,6 +280,7 @@ def __init__( max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, use_eagle=self.use_eagle, + num_prefill_lookahead=self.num_prefill_lookahead, log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, @@ -438,6 +452,27 @@ def _get_local_prefix_cache_hit( ) return blocks, num_local, shared_prefix_boundary, False + def _reserve_prefill_lookahead( + self, + request: Request, + num_computed_tokens: int, + num_new_tokens: int, + ) -> int: + """Never end a prefill chunk within num_prefill_lookahead of the + prefill end. + + At a chunked-prefill boundary, the multi-module MTP drafter consumes + the next num_prefill_lookahead known prefill tokens as draft inputs. A + boundary closer to the end than that would make it fall back to + sampled drafts, permanently polluting the trailing modules' KV caches. + Either finish the prefill or leave at least num_prefill_lookahead for + the next chunk. No-op for eagle-family drafters (lookahead 1). + """ + remaining = request.num_tokens - num_computed_tokens - num_new_tokens + if 0 < remaining < self.num_prefill_lookahead: + num_new_tokens -= self.num_prefill_lookahead - remaining + return max(num_new_tokens, 0) + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: @@ -561,9 +596,15 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: request.num_computed_tokens, num_new_tokens, encoder_compute_budget, - shift_computed_tokens=1 if self.use_eagle else 0, + shift_computed_tokens=self.num_prefill_lookahead, ) + # Multi-module MTP: avoid ending a prefill chunk within + # num_prefill_lookahead of the prefill end. + num_new_tokens = self._reserve_prefill_lookahead( + request, request.num_computed_tokens, num_new_tokens + ) + if num_new_tokens == 0: # The request cannot be scheduled because one of the following # reasons: @@ -576,6 +617,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # 3. The encoder cache is exhausted. # 4. Insufficient budget for a block-aligned chunk in hybrid # models with mamba cache mode \"align\". + # 5. Insufficient budget to keep a multi-module MTP prefill + # chunk out of the prefill-lookahead window. # NOTE(woosuk): Here, by doing `continue` instead of `break`, # we do not strictly follow the FCFS scheduling policy and # allow the lower-priority requests to be scheduled. @@ -946,11 +989,18 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_computed_tokens, num_new_tokens, encoder_compute_budget, - shift_computed_tokens=1 if self.use_eagle else 0, + shift_computed_tokens=self.num_prefill_lookahead, ) - if num_new_tokens == 0: - # The request cannot be scheduled. - break + + # Multi-module MTP: avoid ending a prefill chunk within + # num_prefill_lookahead of the prefill end. + num_new_tokens = self._reserve_prefill_lookahead( + request, num_computed_tokens, num_new_tokens + ) + + if num_new_tokens == 0: + # The request cannot be scheduled. + break # During async KV load, no forward pass is run yet. # Allocate speculative lookahead slots later to avoid @@ -2157,9 +2207,9 @@ def _free_encoder_inputs(self, request: Request) -> None: return # Defer the free by the drafter's look-ahead so an entry stays - # referenced until the drafter's +1 read has also passed it, mirroring - # the shift the encoder scheduling path applies. - spec_lookahead = 1 if self.use_eagle else 0 + # referenced until the drafter's read-ahead has also passed it, + # mirroring the shift the encoder scheduling path applies. + spec_lookahead = self.num_prefill_lookahead # Here, we use list(set) to avoid modifying the set while iterating # over it. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index f5f854a5dfd9..75acb7f4973e 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -28,7 +28,6 @@ SinkFullAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, - TQFullAttentionSpec, ) from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import Request @@ -561,6 +560,11 @@ def find_longest_cache_hit( return an empty list. If eagle is enabled, drop the last matched block to force recompute the last block to get the required hidden states for eagle drafting head. + For multi-module MTP, this recompute also rewrites the dropped block's + draft-layer KVs, which depend on up to num_speculative_tokens - 1 + tokens past the matched prefix (i.e. on the cache writer's + continuation, which the block hash does not cover); the coordinator + asserts the block size covers that window. Need to be customized for each attention type. Args: @@ -877,6 +881,10 @@ class SlidingWindowManager(SingleTypeKVCacheManager): def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None: super().__init__(kv_cache_spec, **kwargs) self.sliding_window = kv_cache_spec.sliding_window + # Extra trailing tokens to retain below the window (never attended) so a + # multi-module MTP store-side lag can still reconstruct the window from + # cached blocks. + self.extra_retained_tokens = kv_cache_spec.extra_retained_tokens @classmethod def _contiguous_blocks_for_hit( @@ -1072,13 +1080,22 @@ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: attention computation since they are outside the sliding window. Thus, get_num_skipped_tokens(7) == 4. + The trailing edge of the window is extended by ``extra_retained_tokens`` + so that those extra trailing tokens' blocks are retained (but not + attended). This is needed for multi-module spec decoding which can + re-prefill the last num_spec_prefill_tokens - 1 tokens from the end + of the sequence, and thus needs to delay freeing/caching of blocks. + Args: num_computed_tokens: The number of tokens that have been computed. Returns: The number of tokens that will be skipped for attention computation. """ - return max(0, num_computed_tokens - self.sliding_window + 1) + return max( + 0, + num_computed_tokens - self.sliding_window + 1 - self.extra_retained_tokens, + ) def get_num_common_prefix_blocks(self, running_request_id: str) -> int: """ @@ -1911,11 +1928,6 @@ def register_all_kvcache_specs(vllm_config): ) # FullAttentionSpec subclasses — grouped with FullAttentionSpec - KVCacheSpecRegistry.register( - TQFullAttentionSpec, - FullAttentionManager, - uniform_type_base_spec=FullAttentionSpec, - ) KVCacheSpecRegistry.register( MLAAttentionSpec, FullAttentionManager, uniform_type_base_spec=FullAttentionSpec ) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index d89c9a3f5dc8..d70778eb51e3 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -92,6 +92,9 @@ class EngineCoreReadyResponse: kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None kv_events_config: KVEventsConfig | None = None + weight_transfer_backend: str | None = None + enable_sleep_mode: bool = False + supports_draft_weight_updates: bool = False class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 956c800c7f74..295c66b360ed 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -173,6 +173,8 @@ def __init__( ) if self.scheduler.connector is not None: # type: ignore self.model_executor.init_kv_output_aggregator(self.scheduler.connector) # type: ignore + if self.scheduler.ec_connector is not None: # type: ignore + self.model_executor.init_ec_output_aggregator() mm_registry = MULTIMODAL_REGISTRY self.mm_receiver_cache = mm_registry.engine_receiver_cache_from_config( @@ -1635,6 +1637,15 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: max_num_batched_tokens=scheduler_config.max_num_batched_tokens, instance_id=self.vllm_config.instance_id, kv_events_config=self.scheduler.get_kv_event_publisher_config(), + weight_transfer_backend=( + self.vllm_config.weight_transfer_config.backend + if self.vllm_config.weight_transfer_config is not None + else None + ), + enable_sleep_mode=self.vllm_config.model_config.enable_sleep_mode, + supports_draft_weight_updates=( + self.model_executor.supports_draft_weight_updates() + ), ) def process_input_sockets( diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index 686044a4031b..ac870c931a32 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -483,6 +483,7 @@ def _validate_model_input( if prompt_ids and tokenizer is not None: max_input_id = max(prompt_ids, default=0) + min_input_id = min(prompt_ids, default=0) # NOTE: tokenizer.max_token_id is the tokenizer’s vocab size while # self.model_config.get_vocab_size() is the model’s vocab size. @@ -495,6 +496,15 @@ def _validate_model_input( # Here we take the max of the two to determine if a token id is # truly out-of-vocabulary. model_vocab_size = model_config.get_vocab_size() + # A negative id is out of vocabulary just like an over-large one, + # but is not caught by the upper-bound check below. Reject it here + # so it is not used as an embedding index downstream. This + # validation path is shared by generate, embedding and pooling + # requests, so the check covers all three. + if min_input_id < 0: + raise VLLMValidationError( + f"Token id {min_input_id} is out of vocabulary" + ) if max_input_id > max(tokenizer.max_token_id, model_vocab_size - 1): raise VLLMValidationError( f"Token id {max_input_id} is out of vocabulary" diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index 404acd50de99..6fff5f032d10 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -9,6 +9,7 @@ import vllm.envs as envs from vllm.config import VllmConfig +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, @@ -110,6 +111,7 @@ def __init__( self.is_sleeping = False self.sleeping_tags: set[str] = set() self.kv_output_aggregator: KVOutputAggregator | None = None + self.ec_output_aggregator: ECOutputAggregator | None = None @abstractmethod def _init_executor(self) -> None: @@ -283,12 +285,21 @@ def init_kv_output_aggregator(self, connector: "KVConnectorBase") -> None: connector, self.parallel_config.world_size ) + def init_ec_output_aggregator(self) -> None: + self.ec_output_aggregator = ECOutputAggregator() + @cached_property # Avoid unnecessary RPC calls def supported_tasks(self) -> tuple[SupportedTask, ...]: output: list[tuple[SupportedTask, ...]] output = self.collective_rpc("get_supported_tasks") return output[0] + def supports_draft_weight_updates(self) -> bool: + worker_support: list[bool] = self.collective_rpc( + "supports_draft_weight_updates" + ) + return all(worker_support) + def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, "lora_id must be greater than 0." return all(self.collective_rpc("add_lora", args=(lora_request,))) diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index afc333723d0f..8f0f638deb88 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -29,6 +29,7 @@ from vllm.config import VllmConfig from vllm.distributed import destroy_distributed_environment, destroy_model_parallel from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue +from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator from vllm.distributed.parallel_state import ( get_dcp_group, @@ -129,7 +130,10 @@ def _init_executor(self) -> None: f"_parallel_size ({pcp_size}). " ) - set_multiprocessing_worker_envs(self.local_world_size) + num_local_procs = self.local_world_size * max( + 1, self.parallel_config.data_parallel_size_local + ) + set_multiprocessing_worker_envs(num_local_procs) if aiter_requires_tcp_store(): distributed_init_method = get_distributed_init_method( @@ -343,6 +347,7 @@ def execute_model( # type: ignore[override] non_block=non_block, timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, kv_output_aggregator=self.kv_output_aggregator, + ec_output_aggregator=self.ec_output_aggregator, ) def sample_tokens( # type: ignore[override] @@ -355,6 +360,7 @@ def sample_tokens( # type: ignore[override] non_block=non_block, timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS, kv_output_aggregator=self.kv_output_aggregator, + ec_output_aggregator=self.ec_output_aggregator, ) def execute_dummy_batch(self) -> None: @@ -375,9 +381,10 @@ def collective_rpc( # type: ignore[override] non_block: bool = False, unique_reply_rank: int | None = None, kv_output_aggregator: KVOutputAggregator | None = None, + ec_output_aggregator: ECOutputAggregator | None = None, ) -> Any: - """Returns single result if unique_reply_rank and/or kv_output_aggregator - is provided, otherwise list.""" + """Returns single result if unique_reply_rank and/or an output + aggregator is provided, otherwise list.""" assert self.rpc_broadcast_mq is not None, ( "collective_rpc should not be called on follower node" ) @@ -387,11 +394,21 @@ def collective_rpc( # type: ignore[override] deadline = None if timeout is None else time.monotonic() + timeout kwargs = kwargs or {} - if kv_output_aggregator is not None: + aggregators = [a for a in (kv_output_aggregator, ec_output_aggregator) if a] + aggregate: Callable[[Any], Any] + if aggregators: output_rank = None - aggregate: Callable[[Any], Any] = partial( - kv_output_aggregator.aggregate, output_rank=unique_reply_rank or 0 - ) + + def _aggregate(outputs: Any) -> Any: + # Each aggregator merges its own connector's output onto + # outputs[output_rank] in place and returns it, so chaining is safe. + rank = unique_reply_rank or 0 + result = outputs[rank] + for a in aggregators: + result = a.aggregate(outputs, output_rank=rank) + return result + + aggregate = _aggregate else: output_rank = unique_reply_rank aggregate = lambda x: x diff --git a/vllm/v1/executor/ray_executor.py b/vllm/v1/executor/ray_executor.py index 39749ffc257e..986e4d9bbb03 100644 --- a/vllm/v1/executor/ray_executor.py +++ b/vllm/v1/executor/ray_executor.py @@ -89,6 +89,19 @@ def _init_executor(self) -> None: # KV connector setup self.has_connector = self.vllm_config.kv_transfer_config is not None + if ( + self.vllm_config.ec_transfer_config is not None + and self.parallel_config.world_size > 1 + ): + raise NotImplementedError( + "EC connector worker metadata is not supported with the " + "legacy Ray executor when world_size > 1: only the output " + "of a single worker is fetched, silently dropping the " + "other workers' EC connector state. Set " + "VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 to use RayExecutorV2, " + "or use the multiprocessing executor instead." + ) + self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and ( self.vllm_config.ec_transfer_config is None or self.vllm_config.ec_transfer_config.is_ec_consumer diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index dcc8298781b5..cbead6d3c885 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -5,6 +5,7 @@ import copy from collections import Counter +from collections.abc import Collection from dataclasses import dataclass, fields, replace from enum import Enum, IntEnum from math import prod @@ -15,7 +16,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up -from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim +from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry @@ -45,7 +46,11 @@ class KVQuantMode(IntEnum): FP8_PER_TOKEN_HEAD = 3 # per-token-head dynamic scales for fp8 INT4_PER_TOKEN_HEAD = 4 # packed 2×int4/byte, RHT + asymmetric zp NVFP4 = 5 # packed fp4 data + fp8 block scales - TURBOQUANT = 6 # Hadamard-rotated Lloyd-Max quant, packed K+V per slot + # Hadamard-rotated Lloyd-Max quant, packed K+V per slot. + TURBOQUANT_K8V4 = 6 + TURBOQUANT_4BIT_NC = 7 + TURBOQUANT_K3V4_NC = 8 + TURBOQUANT_3BIT_NC = 9 @property def is_per_token_head(self) -> bool: @@ -63,8 +68,13 @@ def is_nvfp4(self) -> bool: @property def is_turboquant(self) -> bool: - """True for turboquant quantization mode.""" - return self == KVQuantMode.TURBOQUANT + """True for any turboquant quantization mode.""" + return self in ( + KVQuantMode.TURBOQUANT_K8V4, + KVQuantMode.TURBOQUANT_4BIT_NC, + KVQuantMode.TURBOQUANT_K3V4_NC, + KVQuantMode.TURBOQUANT_3BIT_NC, + ) def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: @@ -78,7 +88,7 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode: if kv_cache_dtype.startswith("nvfp4"): return KVQuantMode.NVFP4 if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("turboquant_"): - return KVQuantMode.TURBOQUANT + return KVQuantMode[kv_cache_dtype.upper()] if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"): return KVQuantMode.FP8_PER_TENSOR return KVQuantMode.NONE @@ -88,14 +98,24 @@ def is_quantized_kv_cache(kv_cache_dtype: str) -> bool: return get_kv_quant_mode(kv_cache_dtype) != KVQuantMode.NONE -def replace_as(spec: KVCacheSpec, target_cls: type[_SpecT], **changes) -> _SpecT: +def replace_as( + spec: KVCacheSpec, + target_cls: type[_SpecT], + *, + drop: Collection[str] = (), + **changes, +) -> _SpecT: """``dataclasses.replace``, but rebuilding *spec* as *target_cls* e.g. ``SlidingWindowSpec`` -> ``FullAttentionSpec`` - Every field of *spec* must exist on *target_cls*; fields only *target_cls* has keep - their default values. + Every field of *spec* must exist on *target_cls* unless named in *drop*; + fields only *target_cls* has keep their default values. """ - kwargs = {f.name: getattr(spec, f.name) for f in fields(spec) if f.init} + kwargs = { + f.name: getattr(spec, f.name) + for f in fields(spec) + if f.init and f.name not in drop + } kwargs.update(changes) return target_cls(**kwargs) @@ -199,21 +219,37 @@ class AttentionSpec(KVCacheSpec): num_kv_heads: int head_size: int dtype: torch.dtype + head_size_v: int = None # type: ignore[assignment] kv_quant_mode: KVQuantMode = KVQuantMode.NONE page_size_padded: int | None = None indexes_kv_by_block_stride: bool = False + num_head_slots: int | None = None + """H of the logical ``[B, H, N, C]`` page when packing diverges from one + slot per KV head. None means one slot per KV head. Published by the backend. + """ + state_content_bytes: int | None = None + """C in bytes when packed; None means dense K/V content.""" + + def __post_init__(self): + if self.head_size_v is None: + object.__setattr__(self, "head_size_v", self.head_size) + + @property + def num_heads(self) -> int: + if self.num_head_slots is not None: + return self.num_head_slots + return self.num_kv_heads + + @property + def state_content_size_bytes(self) -> int: + """Bytes per (head slot, stored state) cell of the page.""" + if self.state_content_bytes is not None: + return self.state_content_bytes + return (self.head_size + self.head_size_v) * get_dtype_size(self.dtype) @property def unpadded_page_size_bytes(self) -> int: - unpadded = self.real_page_size_bytes - # Per-token-head scales are stored in separate tensors managed - # by the attention backend, but the memory is carved from the - # raw KV cache allocation so it must be budgeted here. - if self.kv_quant_mode.is_per_token_head: - unpadded += ( - 2 * self.block_size * self.num_kv_heads * get_dtype_size(torch.float32) - ) - return unpadded + return self.num_heads * self.storage_block_size * self.state_content_size_bytes @property def page_size_bytes(self) -> int: @@ -224,20 +260,10 @@ def page_size_bytes(self) -> int: @property def real_page_size_bytes(self) -> int: - if self.kv_quant_mode.is_nvfp4: - # Packed layout: fp4 data + fp8 block scales per head. - head_dim = nvfp4_kv_cache_full_dim(self.head_size) - elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD: - head_dim = self.head_size // 2 - else: - head_dim = self.head_size - return ( - 2 - * self.block_size - * self.num_kv_heads - * head_dim - * get_dtype_size(self.dtype) - ) + """Alias of ``unpadded_page_size_bytes`` + TODO(lucas): follow up with TPU backend to see if we can remove this property. + """ + return self.unpadded_page_size_bytes def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config @@ -256,8 +282,6 @@ class FullAttentionSpec(AttentionSpec): In this case, we use FullAttentionSpec and record the sliding window size. """ - head_size_v: int = None # type: ignore[assignment] - sliding_window: int | None = None """ Default to None for not using sliding window attention. @@ -273,10 +297,6 @@ class FullAttentionSpec(AttentionSpec): cache layout itself. """ - def __post_init__(self): - if self.head_size_v is None: - object.__setattr__(self, "head_size_v", self.head_size) - def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size @@ -326,6 +346,8 @@ def merge(cls, specs: list[Self]) -> Self: kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride, + num_head_slots=specs[0].num_head_slots, + state_content_bytes=specs[0].state_content_bytes, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), # If any layer in the group is non-causal, treat the group as @@ -346,23 +368,6 @@ def merge(cls, specs: list[Self]) -> Self: ) return merged_spec - @property - def real_page_size_bytes(self) -> int: - if self.kv_quant_mode.is_nvfp4: - # Packed layout per head: fp4 data + fp8 block scales. - # fp4 data: head_size//2 bytes (2 fp4 values per byte) - # fp8 block scale: head_size//16 bytes (1 scale per 16 elements) - last_dim = nvfp4_kv_cache_full_dim( - self.head_size - ) + nvfp4_kv_cache_full_dim(self.head_size_v) - elif self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD: - last_dim = self.head_size // 2 + self.head_size_v // 2 - else: - last_dim = self.head_size + self.head_size_v - return ( - self.block_size * self.num_kv_heads * last_dim * get_dtype_size(self.dtype) - ) - def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): if spec.alignment is None: @@ -373,32 +378,6 @@ def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): object.__setattr__(spec, "page_size_padded", padded_page_size) -@dataclass(frozen=True, kw_only=True) -class TQFullAttentionSpec(FullAttentionSpec): - """FullAttentionSpec with TQ-aware page size. - - Python equivalent of the C++ TQ4FullAttentionSpec. Overrides - real_page_size_bytes to use TQ slot bytes instead of the raw - head_size * dtype formula. - """ - - tq_slot_size: int = 0 - - @property - def real_page_size_bytes(self) -> int: - if self.tq_slot_size > 0: - return self.block_size * self.num_kv_heads * self.tq_slot_size - return super().real_page_size_bytes - - @classmethod - def merge(cls, specs: list[Self]) -> Self: - merged = super().merge(specs) - assert all(s.tq_slot_size == specs[0].tq_slot_size for s in specs), ( - "All TQ layers in the same KV cache group must use the same tq_slot_size." - ) - return replace(merged, tq_slot_size=specs[0].tq_slot_size) - - @dataclass(frozen=True, kw_only=True) class MLAAttentionSpec(FullAttentionSpec): # TODO(Lucas/Chen): less hacky way to do this @@ -409,6 +388,8 @@ class MLAAttentionSpec(FullAttentionSpec): model_version: str | None = None # Marks draft groups that flatten a non-causal query block into decode rows. non_causal_multi_token_decode: bool = False + # MLA stores a single latent vector per state; there is no separate V. + head_size_v: int = 0 def __post_init__(self): super().__post_init__() @@ -418,27 +399,6 @@ def __post_init__(self): def storage_block_size(self) -> int: return self.block_size // self.compress_ratio - @property - def real_page_size_bytes(self) -> int: - if self.cache_dtype_str == "fp8_ds_mla": - if self.model_version == "deepseek_v4": - # DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token. - # head_size stays semantic (512); bytes are determined here. - return self.storage_block_size * 584 - # V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 + - # qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py. - return self.block_size * 656 - if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD: - head_dim = self.head_size // 2 - else: - head_dim = self.head_size - return ( - self.storage_block_size - * self.num_kv_heads - * head_dim - * get_dtype_size(self.dtype) - ) - @classmethod def merge(cls, specs: list[Self]) -> Self: assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), ( @@ -465,6 +425,8 @@ def merge(cls, specs: list[Self]) -> Self: dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, + num_head_slots=specs[0].num_head_slots, + state_content_bytes=specs[0].state_content_bytes, indexes_kv_by_block_stride=block_stride_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), compress_ratio=compress_ratio_set.pop(), @@ -522,6 +484,8 @@ def merge(cls, specs: list[RSWASpec]) -> RSWASpec: kv_quant_mode=base.kv_quant_mode, page_size_padded=base.page_size_padded, indexes_kv_by_block_stride=base.indexes_kv_by_block_stride, + num_head_slots=base.num_head_slots, + state_content_bytes=base.state_content_bytes, sliding_window=base.sliding_window, attention_chunk_size=base.attention_chunk_size, non_causal=base.non_causal, @@ -572,31 +536,12 @@ def is_uniform_with_collection( @dataclass(frozen=True, kw_only=True) class SlidingWindowSpec(AttentionSpec): sliding_window: int - head_size_v: int = None # type: ignore[assignment] - - def __post_init__(self): - if self.head_size_v is None: - object.__setattr__(self, "head_size_v", self.head_size) - - @property - def real_page_size_bytes(self) -> int: - # Mirror ``FullAttentionSpec.real_page_size_bytes`` for NVFP4 KV cache. - if self.kv_quant_mode.is_nvfp4: - last_dim = nvfp4_kv_cache_full_dim( - self.head_size - ) + nvfp4_kv_cache_full_dim(self.head_size_v) - return ( - self.block_size - * self.num_kv_heads - * last_dim - * get_dtype_size(self.dtype) - ) - return ( - self.block_size - * self.num_kv_heads - * (self.head_size + self.head_size_v) - * get_dtype_size(self.dtype) - ) + # The trailing edge of the window is extended by ``extra_retained_tokens`` + # so that those extra trailing tokens' blocks are retained (but not + # attended). This is needed for multi-module spec decoding which can + # re-prefill the last num_spec_prefill_tokens - 1 tokens from the end + # of the sequence, and thus needs to delay freeing/caching of blocks. + extra_retained_tokens: int = 0 def max_admission_blocks_per_request( self, max_in_flight_tokens: int, max_model_len: int @@ -614,8 +559,13 @@ def max_admission_blocks_per_request( """ # During chunked prefill, we hold KV for the last `sliding_window-1` # computed tokens plus the in-flight tokens (frees happen on the - # processed-token basis); never more than `max_model_len`. - num_tokens = min(self.sliding_window - 1 + max_in_flight_tokens, max_model_len) + # processed-token basis); never more than `max_model_len`. An additional + # `extra_retained_tokens` trailing tokens are kept alive below the + # window for multi-module spec decoding, and must be accounted here too. + num_tokens = min( + self.sliding_window - 1 + self.extra_retained_tokens + max_in_flight_tokens, + max_model_len, + ) # +1 because the sliding window may not start from the beginning of # the block. E.g. block size 4 and num_token 4 needs two blocks # [XXCD][EF] to store the 6-token window [CDEF]. @@ -650,31 +600,20 @@ class SlidingWindowMLASpec(SlidingWindowSpec): alignment: int | None = None # Default to None for no padding. compress_ratio: int = 1 model_version: str | None = None + # MLA stores a single latent vector per state; there is no separate V. + head_size_v: int = 0 def __post_init__(self): + assert self.model_version in (None, "deepseek_v4"), ( + f"Unsupported model version: {self.model_version}" + ) + super().__post_init__() _apply_alignment_padding(self) @property def storage_block_size(self) -> int: return self.block_size // self.compress_ratio - @property - def real_page_size_bytes(self) -> int: - if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla": - # DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B - # per token. FlashInfer's contiguous bf16/fp8 cache falls through to - # the element-size formula below. - return self.storage_block_size * 584 - assert self.model_version in (None, "deepseek_v4"), ( - f"Unsupported model version: {self.model_version}" - ) - return ( - self.storage_block_size - * self.num_kv_heads - * self.head_size - * get_dtype_size(self.dtype) - ) - @classmethod def merge(cls, specs: list[Self]) -> Self: assert all(isinstance(spec, SlidingWindowMLASpec) for spec in specs), ( @@ -686,16 +625,18 @@ def merge(cls, specs: list[Self]) -> Self: model_version_set = set(spec.model_version for spec in specs) sliding_window_set = set(spec.sliding_window for spec in specs) block_stride_set = set(spec.indexes_kv_by_block_stride for spec in specs) + extra_retained_set = set(spec.extra_retained_tokens for spec in specs) assert ( len(cache_dtype_str_set) == 1 and len(compress_ratio_set) == 1 and len(model_version_set) == 1 and len(sliding_window_set) == 1 and len(block_stride_set) == 1 + and len(extra_retained_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, compress ratio, model version, sliding " - "window size, and KV block stride indexing." + "window size, KV block stride indexing, and retained token count." ) return cls( block_size=specs[0].block_size, @@ -703,8 +644,11 @@ def merge(cls, specs: list[Self]) -> Self: head_size=specs[0].head_size, dtype=specs[0].dtype, page_size_padded=specs[0].page_size_padded, + num_head_slots=specs[0].num_head_slots, + state_content_bytes=specs[0].state_content_bytes, indexes_kv_by_block_stride=block_stride_set.pop(), sliding_window=sliding_window_set.pop(), + extra_retained_tokens=extra_retained_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), compress_ratio=compress_ratio_set.pop(), model_version=model_version_set.pop(), @@ -828,6 +772,8 @@ def merge(cls, specs: list[Self]) -> Self: kv_quant_mode=specs[0].kv_quant_mode, page_size_padded=specs[0].page_size_padded, indexes_kv_by_block_stride=specs[0].indexes_kv_by_block_stride, + num_head_slots=specs[0].num_head_slots, + state_content_bytes=specs[0].state_content_bytes, sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), non_causal=any(spec.non_causal for spec in specs), @@ -1024,6 +970,8 @@ class KVCacheConfig: For models with multiple types of attention, there will be multiple groups, see `_get_kv_cache_config_uniform_page_size` for more details. """ + prefix_cache_retention_interval: int | None = None + """Resolved retention policy for local prefix-cache checkpoints.""" @property def has_mamba_layers(self) -> bool: diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 0bbee7667527..80b909afbc8d 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -14,6 +14,7 @@ from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: + from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata from vllm.distributed.kv_events import KVConnectorKVEvents from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorWorkerMetadata, @@ -23,6 +24,7 @@ KVConnectorStats = object KVConnectorWorkerMetadata = object KVConnectorKVEvents = object + ECConnectorWorkerMetadata = object class LogprobsLists(NamedTuple): @@ -292,6 +294,14 @@ class ECConnectorOutput: # [mm_hash] finished_sending: set[str] | None = None finished_recving: set[str] | None = None + ec_connector_worker_meta: ECConnectorWorkerMetadata | None = None + + def is_empty(self): + return ( + not self.finished_sending + and not self.finished_recving + and not self.ec_connector_worker_meta + ) # ModelRunnerOutput is serialized and sent to the scheduler process. @@ -362,6 +372,32 @@ def with_kv_conn_output_only( output.kv_connector_output = kv_connector_output return output + @staticmethod + def with_ec_conn_output_only( + ec_connector_output: ECConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return an otherwise-empty output carrying `ec_connector_output`.""" + return ModelRunnerOutput.with_ec_conn_output( + EMPTY_MODEL_RUNNER_OUTPUT, ec_connector_output + ) + + @staticmethod + def with_ec_conn_output( + output: "ModelRunnerOutput", + ec_connector_output: ECConnectorOutput | None, + ) -> "ModelRunnerOutput": + """Return `output` carrying `ec_connector_output`. + + The shared empty output is copied rather than written to, so callers + must use the return value. + """ + if ec_connector_output is None or ec_connector_output.is_empty(): + return output + if output is EMPTY_MODEL_RUNNER_OUTPUT: + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.ec_connector_output = ec_connector_output + return output + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 7d593e69cd7f..d8f8c8956d25 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -12,6 +12,7 @@ import torch from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import num_compute_units @@ -919,12 +920,13 @@ def apply_top_k_top_p_triton( # Cache lookup table entries on each device. tables = _TRITON_TABLE_CACHE.get(logits.device) if tables is None: - normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE) - percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE) - _TRITON_TABLE_CACHE[logits.device] = ( - normal_cdf_to_sigma_table, - percentile_to_std_table, - ) + with gpu_sync_allowed(): + normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE) + percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE) + _TRITON_TABLE_CACHE[logits.device] = ( + normal_cdf_to_sigma_table, + percentile_to_std_table, + ) else: normal_cdf_to_sigma_table, percentile_to_std_table = tables diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index ac8512f22c8f..3efaf33d980f 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -6,6 +6,7 @@ import torch.nn as nn from vllm.config.model import LogprobsMode +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata @@ -342,9 +343,10 @@ def gather_logprobs( # of the compiled batched_count_greater_than. mark_unbacked makes # the size fully symbolic so dynamo doesn't specialize when # batch_size transitions from 1 to >=2. - torch._dynamo.decorators.mark_unbacked(logprobs, 0) - torch._dynamo.decorators.mark_unbacked(token_logprobs, 0) - token_ranks = batched_count_greater_than(logprobs, token_logprobs) + with gpu_sync_allowed(first_only=True): + torch._dynamo.decorators.mark_unbacked(logprobs, 0) + torch._dynamo.decorators.mark_unbacked(token_logprobs, 0) + token_ranks = batched_count_greater_than(logprobs, token_logprobs) # Concatenate together with the topk. indices = torch.cat((token_ids, topk_indices), dim=1) diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index efac0111779e..43880bf656e5 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -100,8 +100,8 @@ def sync_batch(self, batch_update: BatchUpdate | None) -> None: for i1, i2, direction in batch_update.moved: if direction == MoveDirectionality.SWAP: - state1 = self._state.get(i1) - state2 = self._state.get(i2) + state1 = self._state.pop(i1, None) + state2 = self._state.pop(i2, None) if state1 is not None: self._state[i2] = state1 if state2 is not None: diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 2e6839fef557..1e4fcae68cf9 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -4,7 +4,7 @@ import contextlib from collections.abc import Iterable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any from vllm.config import VllmConfig @@ -190,7 +190,6 @@ def _derive_cpu_config( """Derive a CPU KVCacheConfig from the GPU config. Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio.""" # Import here to avoid potential circular imports - from vllm.v1.kv_cache_interface import KVCacheConfig as KVCacheConfigCls from vllm.v1.kv_cache_interface import KVCacheTensor assert len(gpu_config.kv_cache_tensors) > 0 @@ -215,10 +214,10 @@ def _derive_cpu_config( for t in gpu_config.kv_cache_tensors ] - return KVCacheConfigCls( + return replace( + gpu_config, num_blocks=num_cpu_blocks, kv_cache_tensors=cpu_tensors, - kv_cache_groups=gpu_config.kv_cache_groups, ) @staticmethod diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 626dd36dea9b..e773ecf30596 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -80,6 +80,21 @@ def __init__( self.draft_model_config.hf_config ) + @override + def load_model(self, target_model: torch.nn.Module) -> None: + from vllm.model_executor.models.qwen3_dflash import ( + dflash_target_rope_is_neox_style, + ) + + draft_model_config = self.speculative_config.draft_model_config + assert draft_model_config is not None + # The drafter must rotate Q/K the way its target does. Take that from the + # built target before super() constructs the draft. + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style + super().load_model(target_model) + @override def _create_draft_vllm_config(self) -> VllmConfig: base = super()._create_draft_vllm_config() diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 30ecbfa065a6..310576cd99a8 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -10,6 +10,7 @@ import torch from transformers import MistralCommonBackend +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader @@ -269,7 +270,7 @@ def _process_schema( begin: str = s["begin"] trig = next((t for t in triggers if begin.startswith(t)), None) if trig is None: - raise ValueError( + raise VLLMValidationError( f"Trigger {begin} not found in triggers {triggers}" ) tags.append( @@ -281,7 +282,9 @@ def _process_schema( ) ) if not tags: - raise ValueError("No structural tags found in the grammar spec.") + raise VLLMValidationError( + "No structural tags found in the grammar spec." + ) return llguidance.StructTag.to_grammar(tags) else: logger.error( @@ -300,7 +303,10 @@ def validate_guidance_grammar( if sampling_params.structured_outputs is None: return tp, grm = get_structured_output_key(sampling_params.structured_outputs) - guidance_grm = serialize_guidance_grammar(tp, grm) + try: + guidance_grm = serialize_guidance_grammar(tp, grm) + except (ValueError, KeyError, TypeError) as e: + raise VLLMValidationError(f"Invalid grammar specification: {e}") from e err = llguidance.LLMatcher.validate_grammar(guidance_grm, tokenizer) if err: - raise ValueError(f"Grammar error: {err}") + raise VLLMValidationError(f"Grammar error: {err}") diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 898aaa136e40..36eb5eb159e1 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -9,6 +9,7 @@ import torch from transformers import PreTrainedTokenizerBase +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader from vllm.utils.torch_utils import PIN_MEMORY @@ -166,7 +167,7 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams so_params.regex, ) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to compile regex for lm-format-enforcer: {err}" ) from err return @@ -176,19 +177,19 @@ def validate_structured_output_request_lm_format_enforcer(params: SamplingParams # make sure schema is valid json json.loads(so_params.json) except json.JSONDecodeError as e: - raise ValueError("Invalid JSON grammar specification.") from e + raise VLLMValidationError("Invalid JSON grammar specification.") from e else: try: json.dumps(so_params.json) except Exception as e: - raise ValueError( + raise VLLMValidationError( f"Error serializing structured outputs jsonschema: {e}" ) from e return elif so_params.choice: return elif so_params.grammar: - raise ValueError( + raise VLLMValidationError( "LM Format Enforcer structured outputs backend " "does not support grammar specifications" ) diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index e66ef6361a70..dd460f96ec91 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -13,6 +13,7 @@ import torch from regex import escape as regex_escape +from vllm.exceptions import VLLMValidationError from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader from vllm.utils.torch_utils import PIN_MEMORY @@ -186,22 +187,27 @@ def validate_structured_output_request_outlines(params: SamplingParams): json.loads(so_params.json) schema = so_params.json except json.JSONDecodeError as e: - raise ValueError("Invalid JSON grammar specification.") from e + raise VLLMValidationError("Invalid JSON grammar specification.") from e else: try: schema = json.dumps(so_params.json) except Exception as e: - raise ValueError( + raise VLLMValidationError( f"Error serializing structured outputs jsonschema: {e}" ) from e - pattern = json_schema.build_regex_from_schema(schema) + try: + pattern = json_schema.build_regex_from_schema(schema) + except Exception as e: + raise VLLMValidationError( + f"Failed to transform json schema into a regex: {e}" + ) from e validate_regex_is_buildable(pattern) elif so_params.choice: choices = [regex_escape(str(choice)) for choice in so_params.choice] regex = "(" + "|".join(choices) + ")" validate_regex_is_buildable(regex) elif so_params.grammar: - raise ValueError( + raise VLLMValidationError( "Outlines structured outputs backend " "does not support grammar specifications" ) @@ -315,19 +321,19 @@ def validate_regex_is_buildable(pattern: str) -> None: parsed = sre_parse.parse(pattern) except sre_constants.error as e: - raise ValueError(f"Error parsing regex: {e}") from e + raise VLLMValidationError(f"Error parsing regex: {e}") from e try: _check_unsupported(parsed) except ValueError as e: - raise ValueError( + raise VLLMValidationError( f"Regex uses unsupported feature for structured outputs: {e}. " "Only basic matching constructs are supported—lookarounds, " "backreferences, and unicode boundaries are not." ) from e if _prefix_needs_context(parsed): - raise ValueError( + raise VLLMValidationError( "Regex does not have a anchored universal start state" "This means that the Regex uses anchors (^) or look-arounds " "in a way which requires context before any token is matched." diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 58b726bf72a4..258b1dff32f1 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -8,6 +8,7 @@ import torch import vllm.envs +from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.sampling_params import SamplingParams from vllm.utils.import_utils import LazyLoader @@ -276,7 +277,7 @@ def check_object(obj: dict[str, Any]) -> bool: def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: """Validate that the request is supported by structured output. - Raises ValueError if the request is not supported. + Raises VLLMValidationError if the request is not supported. """ if sampling_params.structured_outputs is None: return @@ -284,13 +285,21 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: so_params = sampling_params.structured_outputs if so_params.regex: + # A NUL byte is never meaningful in a regex pattern and is not handled + # by xgrammar's native regex converter. Reject it here, before the + # pattern reaches that native code; the try/except below does not cover + # this case. + if "\x00" in so_params.regex: + raise ValueError( + "structured_outputs.regex must not contain a NUL character ('\\x00')" + ) try: compile_regex_with_timeout( xgr.Grammar.from_regex, so_params.regex, ) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to transform regex into a grammar: {err}" ) from err @@ -299,7 +308,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: try: xgr.Grammar.from_ebnf(choice_grammar) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to transform choices into a grammar: {err}" ) from err so_params.choice = None @@ -311,19 +320,19 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: try: schema = json.loads(so_params.json) except json.JSONDecodeError as e: - raise ValueError("Invalid JSON grammar specification.") from e + raise VLLMValidationError("Invalid JSON grammar specification.") from e else: schema = so_params.json if has_xgrammar_unsupported_json_features(schema): - raise ValueError( + raise VLLMValidationError( "The provided JSON schema contains features not supported by xgrammar." ) try: xgr.Grammar.from_json_schema(schema) except Exception as err: - raise ValueError( + raise VLLMValidationError( f"Failed to transform json schema into a grammar: {err}" ) from err return @@ -334,7 +343,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: try: so_params.grammar = convert_lark_to_ebnf(so_params.grammar) except ValueError as e: - raise ValueError( + raise VLLMValidationError( "Failed to convert the grammar from Lark to EBNF. " ) from e @@ -343,7 +352,7 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: # parse the grammar, but we aren't compiling it. xgr.Grammar.from_ebnf(so_params.grammar) except Exception as e: - raise ValueError("Invalid grammar specification.") from e + raise VLLMValidationError("Invalid grammar specification.") from e return if so_params.structural_tag: @@ -364,4 +373,4 @@ def validate_xgrammar_grammar(sampling_params: SamplingParams) -> None: else: xgr.Grammar.from_structural_tag(so_params.structural_tag) except Exception as e: - raise ValueError("Invalid structural tag specification.") from e + raise VLLMValidationError("Invalid structural tag specification.") from e diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index 85228364afcf..5bf5df3b3985 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -2,13 +2,22 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +from dataclasses import dataclass from enum import Enum +from typing import Any import numpy as np import torch from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger +from vllm.model_executor.warmup.jit_warmup import ( + VllmJitKernel, +) +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonWarmupTensor, + triton_scalar_specialization_rep, +) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -134,6 +143,16 @@ def __init__( self.dcp_rank = 0 self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size self.slot_mapping_mode = slot_mapping_mode + if self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT: + _COMPUTE_SLOT_MAPPING_KERNEL.register_warmup( + kv_cache_block_size=self.kv_cache_block_size, + blocks_per_kv_block=self.blocks_per_kv_block, + total_cp_world_size=self.dcp_world_size, + total_cp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + block_table_stride=self.block_table.gpu.stride(0), + block_size=self.block_size, + ) def append_row( self, @@ -192,7 +211,8 @@ def compute_slot_mapping( return assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT - _compute_slot_mapping_kernel[(num_reqs + 1,)]( + _COMPUTE_SLOT_MAPPING_KERNEL( + num_reqs, num_tokens, self.max_num_batched_tokens, query_start_loc, @@ -201,13 +221,11 @@ def compute_slot_mapping( self.block_table.gpu.stride(0), self.block_size, self.slot_mapping.gpu, - KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size, - BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block, - TOTAL_CP_WORLD_SIZE=self.dcp_world_size, - TOTAL_CP_RANK=self.dcp_rank, - CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size, - PAD_ID=PAD_SLOT_ID, - BLOCK_SIZE=1024, + self.kv_cache_block_size, + self.blocks_per_kv_block, + self.dcp_world_size, + self.dcp_rank, + self.cp_kv_cache_interleave_size, ) def commit_block_table(self, num_reqs: int) -> None: @@ -376,67 +394,136 @@ def __getitem__(self, idx: int) -> "BlockTable": return self.block_tables[idx] -@triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) -def _compute_slot_mapping_kernel( - num_tokens, - max_num_tokens, - query_start_loc_ptr, # [num_reqs + 1], int32 - positions_ptr, # [num_tokens], int64 - block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) - block_table_stride, # max_num_blocks_per_req - block_size, - slot_mapping_ptr, # [max_num_tokens], int64 - KV_CACHE_BLOCK_SIZE: tl.constexpr, - BLOCKS_PER_KV_BLOCK: tl.constexpr, - TOTAL_CP_WORLD_SIZE: tl.constexpr, - TOTAL_CP_RANK: tl.constexpr, - CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, - PAD_ID: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - req_idx = tl.program_id(0) - - if req_idx == tl.num_programs(0) - 1: - # Pad remaining slots for CUDA graph compatibility. - for i in range(num_tokens, max_num_tokens, BLOCK_SIZE): +class ComputeSlotMappingKernel(VllmJitKernel["ComputeSlotMappingKernel.CompileKey"]): + triton_block_size = 1024 + + @dataclass(frozen=True) + class CompileKey: + kv_cache_block_size: int + blocks_per_kv_block: int + total_cp_world_size: int + total_cp_rank: int + cp_kv_cache_interleave_size: int + block_table_stride: int + block_size: int + + @staticmethod + @triton.jit(do_not_specialize=["num_tokens", "max_num_tokens"]) + def kernel( + num_tokens, + max_num_tokens, + query_start_loc_ptr, # [num_reqs + 1], int32 + positions_ptr, # [num_tokens], int64 + block_table_ptr, # [max_num_reqs, max_num_blocks_per_req], int32 (flat) + block_table_stride, # max_num_blocks_per_req + block_size, + slot_mapping_ptr, # [max_num_tokens], int64 + KV_CACHE_BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_KV_BLOCK: tl.constexpr, + TOTAL_CP_WORLD_SIZE: tl.constexpr, + TOTAL_CP_RANK: tl.constexpr, + CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr, + PAD_ID: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + req_idx = tl.program_id(0) + + if req_idx == tl.num_programs(0) - 1: + # Pad remaining slots for CUDA graph compatibility. + for i in range(num_tokens, max_num_tokens, BLOCK_SIZE): + offsets = i + tl.arange(0, BLOCK_SIZE) + tl.store( + slot_mapping_ptr + offsets, + PAD_ID, + mask=offsets < max_num_tokens, + ) + return + + start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64) + end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64) + + virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE + row_offset = req_idx * block_table_stride + for i in range(start_idx, end_idx, BLOCK_SIZE): offsets = i + tl.arange(0, BLOCK_SIZE) - tl.store( - slot_mapping_ptr + offsets, - PAD_ID, - mask=offsets < max_num_tokens, + mask = offsets < end_idx + pos = tl.load(positions_ptr + offsets, mask=mask, other=0) + virtual_block_indices = pos // virtual_block_size + virtual_block_offsets = pos - virtual_block_indices * virtual_block_size + is_local = ( + virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE + ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK + local_block_offsets = ( + virtual_block_offsets + // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) + ) * CP_KV_CACHE_INTERLEAVE_SIZE + ( + virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE + ) + + block_indices = ( + virtual_block_indices * BLOCKS_PER_KV_BLOCK + + local_block_offsets // block_size ) - return - - start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64) - end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64) - - virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE - row_offset = req_idx * block_table_stride - for i in range(start_idx, end_idx, BLOCK_SIZE): - offsets = i + tl.arange(0, BLOCK_SIZE) - mask = offsets < end_idx - pos = tl.load(positions_ptr + offsets, mask=mask, other=0) - virtual_block_indices = pos // virtual_block_size - virtual_block_offsets = pos - virtual_block_indices * virtual_block_size - is_local = ( - virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE - ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK - local_block_offsets = ( - virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) - ) * CP_KV_CACHE_INTERLEAVE_SIZE + ( - virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE + block_numbers = tl.load( + block_table_ptr + row_offset + block_indices, + mask=mask & is_local, + other=0, + ).to(tl.int64) + slot_offsets = local_block_offsets % block_size + slot_ids = block_numbers * block_size + slot_offsets + slot_ids = tl.where(is_local, slot_ids, PAD_ID) + tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) + + def dispatch( # type: ignore[override] + self, + *, + block_table_stride: int, + block_size: int, + **compile_key_fields: int, + ) -> CompileKey: + return self.CompileKey( + **compile_key_fields, + block_table_stride=triton_scalar_specialization_rep(block_table_stride), + block_size=triton_scalar_specialization_rep(block_size), ) - block_indices = ( - virtual_block_indices * BLOCKS_PER_KV_BLOCK - + local_block_offsets // block_size + def get_warmup_keys(self, **dispatch_kwargs: int) -> list[CompileKey]: + return self._trace_dispatch(self.dispatch)(**dispatch_kwargs) + + def compile(self, compile_key: CompileKey) -> None: + warmup = getattr(self.kernel, "warmup", None) + assert warmup is not None + int32_ptr = TritonWarmupTensor(torch.int32) + int64_ptr = TritonWarmupTensor(torch.int64) + warmup( + 2, # arbitrary, num_tokens in do_not_specialize + 2, # arbitrary, max_num_tokens in do_not_specialize + int32_ptr, + int64_ptr, + int32_ptr, + compile_key.block_table_stride, + compile_key.block_size, + int64_ptr, + KV_CACHE_BLOCK_SIZE=compile_key.kv_cache_block_size, + BLOCKS_PER_KV_BLOCK=compile_key.blocks_per_kv_block, + TOTAL_CP_WORLD_SIZE=compile_key.total_cp_world_size, + TOTAL_CP_RANK=compile_key.total_cp_rank, + CP_KV_CACHE_INTERLEAVE_SIZE=compile_key.cp_kv_cache_interleave_size, + PAD_ID=PAD_SLOT_ID, + BLOCK_SIZE=self.triton_block_size, + grid=(2,), ) - block_numbers = tl.load( - block_table_ptr + row_offset + block_indices, - mask=mask & is_local, - other=0, - ).to(tl.int64) - slot_offsets = local_block_offsets % block_size - slot_ids = block_numbers * block_size + slot_offsets - slot_ids = tl.where(is_local, slot_ids, PAD_ID) - tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask) + + def __call__( + self, + num_reqs: int, + *args: Any, + ) -> None: + self.kernel[(num_reqs + 1,)]( + *args, + PAD_ID=PAD_SLOT_ID, + BLOCK_SIZE=self.triton_block_size, + ) + + +_COMPUTE_SLOT_MAPPING_KERNEL = ComputeSlotMappingKernel() diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 67cffa58454a..3e7f0d4c1e88 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -74,7 +74,7 @@ def _postprocess_triton(self) -> None: import vllm.v1.worker.block_table - vllm.v1.worker.block_table._compute_slot_mapping_kernel = ( + vllm.v1.worker.block_table._COMPUTE_SLOT_MAPPING_KERNEL.kernel = ( cpu_tl.compute_slot_mapping_kernel ) diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index e7c6d81a9929..a7476bffa41c 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext + import torch import torch.distributed as dist from vllm.config import ParallelConfig from vllm.distributed.parallel_state import get_dp_group from vllm.logger import init_logger +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.worker.ubatch_utils import ( check_ubatch_thresholds, is_last_ubatch_empty, @@ -136,27 +139,33 @@ def _synchronize_dp_ranks( parallel_config=parallel_config, ) - # Synchronize cudagraph_mode across ranks first (take min). - # This is needed before DP padding decision since we use the synced - # cudagraph mode to determine whether DP padding is needed. - synced_cudagraph_mode = _post_process_cudagraph_mode(tensor) - - # Check conditions for microbatching - should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches) - - # DP padding is needed when cudagraph is enabled (synced across ranks) - # or when ubatching/DBO is active (ubatching requires uniform batch - # sizes across DP ranks currently). - # Use the synced runtime cudagraph mode rather than the compilation config - # so we can avoid padding when cudagraph is not enabled for this step. - should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch - - # Pad all DP ranks up to the maximum token count across ranks if - # should_dp_pad is True - num_tokens_after_padding = _post_process_dp_padding( - tensor, - should_dp_pad, - ) + # Only the NCCL path leaves `tensor` on device. With Gloo -- the default + # under async scheduling -- the all-reduce runs on CPU, so the reads below + # are host-side and the check should stay armed. + with ( + nullcontext() + if parallel_config.disable_nccl_for_dp_synchronization + else gpu_sync_allowed() + ): + # Synchronize cudagraph_mode across ranks first (take min). + # This is needed before DP padding decision since we use the synced + # cudagraph mode to determine whether DP padding is needed. + synced_cudagraph_mode = _post_process_cudagraph_mode(tensor) + + # Check conditions for microbatching + should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches) + + # DP padding is needed when cudagraph is enabled (synced across ranks) + # or when ubatching/DBO is active (ubatching requires uniform batch + # sizes across DP ranks currently). + # Use the synced runtime cudagraph mode rather than the compilation + # config so we can avoid padding when cudagraph is not enabled for + # this step. + should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch + + # Pad all DP ranks up to the maximum token count across ranks if + # should_dp_pad is True + num_tokens_after_padding = _post_process_dp_padding(tensor, should_dp_pad) return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 0b3543da1eb2..04a54a39a0ac 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -19,6 +19,7 @@ ) from vllm.model_executor.models.utils import scatter_output_slices from vllm.model_executor.models.vision import get_load_balance_assignment +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphConfig, EncoderItemSpec, @@ -203,6 +204,10 @@ def supports_modality(self, modality: str) -> bool: """Check if a modality is supported by this manager.""" return modality in self.config.modalities + def is_captured(self) -> bool: + """Return whether a CUDA graph pool is active.""" + return self.graph_pool is not None + def clear(self) -> None: """Release captured encoder CUDA graphs and the manager-local pool.""" for graph_set in self.budget_graphs.values(): @@ -294,7 +299,19 @@ def _find_smallest_fitting_budget_given_tokens( def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]: """Get item specs from the model.""" - return self.model.get_encoder_cudagraph_item_specs(mm_kwargs) + # Implementations read per-item grid/patch counts off device tensors + # to size the cudagraph buffers, so the D2H is inherent here. + with gpu_sync_allowed(): + return self.model.get_encoder_cudagraph_item_specs(mm_kwargs) + + def _select_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + """Select the mm kwargs for `indices` from the model.""" + # Same as `_get_item_specs`: implementations re-read the per-item + # grid/patch counts to slice the batch, so the D2H is inherent. + with gpu_sync_allowed(): + return self.model.select_encoder_cudagraph_items(mm_kwargs, indices) def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]: """Get per-item output token counts as plain ints.""" @@ -411,9 +428,7 @@ def append_current_batch() -> None: outputs_by_orig_idx: dict[int, torch.Tensor] = {} for batch_indices, path_budgets in batches: - batch_mm_kwargs = self.model.select_encoder_cudagraph_items( - mm_kwargs, batch_indices - ) + batch_mm_kwargs = self._select_items(mm_kwargs, batch_indices) graph_outputs: dict[str, torch.Tensor] = {} all_eager = True @@ -487,11 +502,9 @@ def _dp_shard( ] if len(local_indices) > 0: - local_mm_kwargs = self.model.select_encoder_cudagraph_items( - mm_kwargs, local_indices - ) + local_mm_kwargs = self._select_items(mm_kwargs, local_indices) else: - local_mm_kwargs = self.model.select_encoder_cudagraph_items(mm_kwargs, []) + local_mm_kwargs = self._select_items(mm_kwargs, []) max_output_tokens_per_rank = ( max( diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index d2977870da1e..77d56ad88298 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -75,6 +75,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: with set_current_vllm_config(vllm_config): indexes = backend.indexes_kv_by_block_stride() spec = replace(spec, indexes_kv_by_block_stride=indexes) + spec = backend.customize_spec(spec) kv_cache_spec[layer_name] = spec return kv_cache_spec diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py index 6dd8fd34743e..77010990b8e3 100644 --- a/vllm/v1/worker/gpu/cp_utils.py +++ b/vllm/v1/worker/gpu/cp_utils.py @@ -59,3 +59,24 @@ def _dcp_local_seq_lens_kernel( # For [num_reqs, max_num_reqs), pad with 0 local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0) tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs) + + +@triton.jit +def cp_local_slot( + positions, + block_numbers, + block_size, + cp_rank, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, + PAD_ID: tl.constexpr, +): + """Return rank-local KV slots, or PAD_ID for positions not owned by this rank.""" + block_offsets = positions % (block_size * CP_SIZE) + if CP_SIZE == 1: + return block_numbers * block_size + block_offsets + is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank + rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE) + remainder = block_offsets % CP_INTERLEAVE + local_offsets = rounds * CP_INTERLEAVE + remainder + return tl.where(is_local, block_numbers * block_size + local_offsets, PAD_ID) diff --git a/vllm/v1/worker/gpu/ec_connector.py b/vllm/v1/worker/gpu/ec_connector.py index 825763b82001..5dc8d92359a9 100644 --- a/vllm/v1/worker/gpu/ec_connector.py +++ b/vllm/v1/worker/gpu/ec_connector.py @@ -9,7 +9,11 @@ from vllm.config import VllmConfig from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase -from vllm.v1.outputs import ECConnectorOutput +from vllm.v1.outputs import ( + EMPTY_MODEL_RUNNER_OUTPUT, + ECConnectorOutput, + ModelRunnerOutput, +) if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -25,6 +29,12 @@ def maybe_get_output( ) -> Generator[ECConnectorOutput | None, None, None]: yield None + def no_forward( + self, + scheduler_output: "SchedulerOutput", + ) -> ModelRunnerOutput: + return EMPTY_MODEL_RUNNER_OUTPUT + class ActiveECConnector(ECConnector): def __init__( @@ -33,9 +43,11 @@ def __init__( encoder_cache: dict[str, torch.Tensor], ) -> None: self.encoder_cache = encoder_cache - self.save_new_caches = vllm_config.is_ec_producer_only self.ec_connector = get_ec_transfer() assert isinstance(self.ec_connector, ECConnectorBase) + # Every producer offloads freshly computed encoder outputs, including + # an ec_both node that also reloads them. + self.save_new_caches = self.ec_connector.is_producer @contextmanager def maybe_get_output( @@ -65,8 +77,19 @@ def maybe_get_output( output.finished_sending, output.finished_recving = ( ec_connector.get_finished(scheduler_output.finished_req_ids) ) + output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta() ec_connector.clear_connector_metadata() + def no_forward( + self, + scheduler_output: "SchedulerOutput", + ) -> ModelRunnerOutput: + # EC send/recv even if no work to do. + with self.maybe_get_output(scheduler_output) as ec_connector_output: + pass + + return ModelRunnerOutput.with_ec_conn_output_only(ec_connector_output) + NO_OP_EC_CONNECTOR = ECConnector() diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index 92bd43534264..a1c79a8eb7b7 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -4,6 +4,7 @@ import time from collections.abc import Collection from contextlib import contextmanager +from typing import TYPE_CHECKING import numpy as np import torch @@ -24,6 +25,9 @@ sanity_check_mm_encoder_outputs, ) +if TYPE_CHECKING: + from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager + logger = init_logger(__name__) @@ -36,6 +40,7 @@ def __init__( encoder_cache: EncoderCache, dtype: torch.dtype, device: torch.device, + cudagraph_manager: "EncoderCudaGraphManager | None" = None, enable_timing: bool = False, ): self.model = model @@ -45,6 +50,7 @@ def __init__( self.dtype = dtype self.device = device self.is_realtime = supports_realtime(model) + self.cudagraph_manager = cudagraph_manager self.enable_timing = enable_timing self.encoder_timing_registry: dict[str, EncoderTimingStats] = {} self._timing_lock = threading.Lock() @@ -53,6 +59,25 @@ def __init__( max_num_tokens, hidden_size, dtype=dtype, device=device ) + def has_cudagraph(self) -> bool: + return self.cudagraph_manager is not None + + @torch.inference_mode() + def capture(self) -> None: + manager = self.cudagraph_manager + assert manager is not None + + from vllm.distributed.parallel_state import graph_capture + from vllm.platforms import current_platform + + with graph_capture(device=self.device): + manager.capture(graph_pool=current_platform.graph_pool_handle()) + torch.accelerator.synchronize() + + def clear(self) -> None: + if self.cudagraph_manager is not None: + self.cudagraph_manager.clear() + def prepare_mm_inputs( self, scheduled_encoder_inputs: dict[str, list[int]] ) -> tuple[list[str], list[tuple[str, MultiModalKwargsItem]]]: @@ -118,7 +143,19 @@ def execute_mm_encoder( for modality, num_items, mm_kwargs_batch in group_and_batch_mm_kwargs( mm_kwargs, device=self.device, pin_memory=PIN_MEMORY ): - batch_outputs = self.model.embed_multimodal(**mm_kwargs_batch) + cg_manager = self.cudagraph_manager + cudagraph_output = ( + cg_manager.execute(mm_kwargs_batch) + if cg_manager is not None + and cg_manager.is_captured() + and cg_manager.supports_modality(modality) + else None + ) + batch_outputs = ( + cudagraph_output + if cudagraph_output is not None + else self.model.embed_multimodal(**mm_kwargs_batch) + ) sanity_check_mm_encoder_outputs(batch_outputs, expected_num_items=num_items) encoder_outputs.extend(batch_outputs) return encoder_outputs diff --git a/vllm/v1/worker/gpu/mm/lora.py b/vllm/v1/worker/gpu/mm/lora.py index 492914a74b6c..819f3b1cb36a 100644 --- a/vllm/v1/worker/gpu/mm/lora.py +++ b/vllm/v1/worker/gpu/mm/lora.py @@ -29,6 +29,7 @@ def set_active_mm_loras( token_lora_mapping: list[int] = [] lora_requests = set() encoder_token_counts: list[int] = [] + connector_token_counts: list[int] = [] # iterate through images for req_id, encoder_input_ids in scheduled_encoder_inputs.items(): @@ -42,10 +43,17 @@ def set_active_mm_loras( # iterate through visual tokens for mm_input_id in encoder_input_ids: pos_info = mm_features[mm_input_id].mm_position - num_tokens = model.get_num_mm_encoder_tokens(pos_info.get_num_embeds()) + + tower_tokens, connector_tokens = model.get_mm_lora_token_counts( + modality=mm_features[mm_input_id].modality, + mm_kwargs=mm_features[mm_input_id].data, + num_mm_embeds=pos_info.get_num_embeds(), + ) + prompt_lora_mapping.append(lora_id) - token_lora_mapping.extend([lora_id] * num_tokens) - encoder_token_counts.append(num_tokens) + token_lora_mapping.extend([lora_id] * tower_tokens) + encoder_token_counts.append(tower_tokens) + connector_token_counts.append(connector_tokens) if lora_id > 0: lora_request = lora_state.lora_requests.get(req_id) @@ -69,19 +77,13 @@ def set_active_mm_loras( if ( mm_mapping is None or not mm_mapping.connector - or not hasattr(model, "get_num_mm_connector_tokens") + or not all(count is not None for count in connector_token_counts) ): return connector_token_mapping = np.repeat( np.array(prompt_lora_mapping, dtype=np.int32), - np.array( - [ - model.get_num_mm_connector_tokens(num_tokens) - for num_tokens in encoder_token_counts - ], - dtype=np.int32, - ), + np.array(connector_token_counts, dtype=np.int32), ) lora_manager.set_active_adapters( lora_requests, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4b53defa4744..aff0c09cab58 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -51,6 +51,7 @@ get_offloader, set_offloader, ) +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import ( MultiModalBudget, @@ -65,6 +66,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import ( DraftTokenIds, + ECConnectorOutput, ModelRunnerOutput, RoutedExpertsTensors, make_empty_encoder_model_runner_output, @@ -149,6 +151,7 @@ copy_kv_cache_blocks_inplace, get_uniform_decode_token_count, ) +from vllm.v1.worker.workspace import use_workspace_lane logger = init_logger(__name__) @@ -164,7 +167,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.parallel_config = vllm_config.parallel_config self.scheduler_config = vllm_config.scheduler_config self.speculative_config = vllm_config.speculative_config + self._draft_workspace_lane = int( + self.speculative_config is not None and self.speculative_config.use_dspark() + ) self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) self.device = device self.dtype = self.model_config.dtype @@ -364,15 +371,16 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: assert self.speculative_config is not None set_eagle3_aux_hidden_state_layers(self.model, self.speculative_config) if isinstance(self.speculator, DraftModelSpeculator): - self.speculator.load_model(self.model) - eplb_models_added = self.eplb.maybe_register_speculator( - self.speculator, self.speculative_config, load_dummy_weights - ) + with use_workspace_lane(self._draft_workspace_lane): + self.speculator.load_model(self.model) + eplb_models_added = self.eplb.maybe_register_speculator( + self.speculator, self.speculative_config, load_dummy_weights + ) time_after_load = time.perf_counter() self.model_memory_usage = m.consumed_memory logger.info( - "Model loading took %s GiB and %.6f seconds", + "Model loading took %s GiB memory and %.6f seconds", format_gib(m.consumed_memory), time_after_load - time_before_load, ) @@ -419,6 +427,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: vocab_size=self.vocab_size, device=self.device, mask_stride=self.decode_query_len, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, ) if self.is_pooling_model and self.is_last_pp_rank: @@ -732,27 +741,28 @@ def _dummy_run( if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] - self.speculator.propose( - input_batch=input_batch, - attn_metadata=attn_metadata, - slot_mappings=slot_mappings_by_layer, - last_hidden_states=spec_hidden_states, - aux_hidden_states=aux_hidden_states, - num_sampled=torch.ones( - input_batch.num_reqs, dtype=torch.int32, device=self.device - ), - num_rejected=torch.zeros( - input_batch.num_reqs, dtype=torch.int32, device=self.device - ), - last_sampled=self.req_states.last_sampled_tokens, - next_prefill_tokens=self.req_states.next_prefill_tokens, - temperature=self.sampler.sampling_states.temperature.gpu, - seeds=self.sampler.sampling_states.seeds.gpu, - dummy_run=True, - skip_attn_for_dummy_run=skip_attn, - mm_inputs=mm_inputs, - is_profile=is_profile, - ) + with use_workspace_lane(self._draft_workspace_lane): + self.speculator.propose( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings=slot_mappings_by_layer, + last_hidden_states=spec_hidden_states, + aux_hidden_states=aux_hidden_states, + num_sampled=torch.ones( + input_batch.num_reqs, dtype=torch.int32, device=self.device + ), + num_rejected=torch.zeros( + input_batch.num_reqs, dtype=torch.int32, device=self.device + ), + last_sampled=self.req_states.last_sampled_tokens, + next_prefill_tokens=self.req_states.next_prefill_tokens, + temperature=self.sampler.sampling_states.temperature.gpu, + seeds=self.sampler.sampling_states.seeds.gpu, + dummy_run=True, + skip_attn_for_dummy_run=skip_attn, + mm_inputs=mm_inputs, + is_profile=is_profile, + ) self.step_timing.drafter_end() assert hidden_states is not None # Last PP rank always has hidden_states @@ -842,10 +852,16 @@ def capture_model(self) -> int: return 0 assert self.cudagraph_manager is not None - if not self.cudagraph_manager.needs_capture(): + capture_encoder = ( + self.model_state.supports_mm_inputs + and self.model_state.encoder_runner.has_cudagraph() + ) + capture_decoder = self.cudagraph_manager.needs_capture() + if not capture_encoder and not capture_decoder: logger.warning( - "Skipping CUDA graph capture. To turn on CUDA graph capture, " - "ensure `cudagraph_mode` was not manually set to `NONE`" + "Skipping encoder and decoder CUDA graph capture. To enable " + "encoder capture, ensure `cudagraph_mm_encoder` is enabled; " + "to enable decoder capture, ensure `cudagraph_mode` is not `NONE`." ) return 0 @@ -857,27 +873,32 @@ def capture_model(self) -> int: start_free_gpu_memory = torch.accelerator.get_memory_info()[0] with self.maybe_setup_dummy_loras(self.lora_config): - self.cudagraph_manager.capture( - self.model, - self.model_state, - self.input_buffers, - self.intermediate_tensors, - self.block_tables, - self.attn_groups, - self.kv_cache_config, - has_lora=self.lora_config is not None, - use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, - lora_capture_hook=create_lora_capture_hook(self.lora_config, self), - ) - if self.speculator is not None: - self.speculator.capture() - if self.adaptive_verification is not None: - with self.step_timing.collect() as timings: - for batch in self.adaptive_verification.batches_to_profile( - self.cudagraph_manager.captured_token_counts() - ): - self._dummy_run(**batch) - self.adaptive_verification.set_initial_cost_curves(timings) + if capture_encoder: + self.model_state.encoder_runner.capture() + + if capture_decoder: + self.cudagraph_manager.capture( + self.model, + self.model_state, + self.input_buffers, + self.intermediate_tensors, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + has_lora=self.lora_config is not None, + use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, + lora_capture_hook=create_lora_capture_hook(self.lora_config, self), + ) + if self.speculator is not None: + with use_workspace_lane(self._draft_workspace_lane): + self.speculator.capture() + if self.adaptive_verification is not None: + with self.step_timing.collect() as timings: + for batch in self.adaptive_verification.batches_to_profile( + self.cudagraph_manager.captured_token_counts() + ): + self._dummy_run(**batch) + self.adaptive_verification.set_initial_cost_curves(timings) end_time = time.perf_counter() end_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -1376,6 +1397,15 @@ def postprocess_sampled( idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu ) + def _merge_ec_connector_no_forward( + self, scheduler_output: SchedulerOutput, output: ModelRunnerOutput + ) -> ModelRunnerOutput: + """Let the EC connector send/recv on a step with no work to run.""" + return ModelRunnerOutput.with_ec_conn_output( + output, + self.ec_connector.no_forward(scheduler_output).ec_connector_output, + ) + @torch.inference_mode() def execute_model( self, @@ -1397,7 +1427,9 @@ def execute_model( if scheduler_output.total_num_scheduled_tokens == 0: # No need to run the model. empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self._merge_ec_connector_no_forward( + scheduler_output, empty_output + ) # Get batch descriptor and sync across DP ranks. num_reqs = len(scheduler_output.num_scheduled_tokens) @@ -1438,7 +1470,7 @@ def execute_model( if batch_desc.num_tokens == 0: # All DP ranks have zero tokens to run. empty_output = self.kv_connector.no_forward(scheduler_output) - return empty_output + return self._merge_ec_connector_no_forward(scheduler_output, empty_output) if not dummy_run: # Common case. @@ -1555,9 +1587,10 @@ def execute_model( input_ids = None if self.is_encoder_only: - output = make_empty_encoder_model_runner_output(scheduler_output) - output.ec_connector_output = ec_connector_output - return output + return ModelRunnerOutput.with_ec_conn_output( + make_empty_encoder_model_runner_output(scheduler_output), + ec_connector_output, + ) model_inputs = { "input_ids": input_ids, @@ -1662,6 +1695,7 @@ def execute_model( hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, finished_req_ids=finished_req_ids, + ec_connector_output=ec_connector_output, routed_experts=routed_experts, ) @@ -1685,6 +1719,7 @@ def sample_tokens( hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output routed_experts = self.execute_model_state.routed_experts self.execute_model_state = None @@ -1704,7 +1739,9 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + # The first PP rank holds the encoder cache, so pass its EC output on. + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) # Last rank: sample tokens hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( @@ -1789,20 +1826,21 @@ def sample_tokens( if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] - draft_tokens = self.speculator.propose( - input_batch, - attn_metadata, - slot_mappings_by_layer, - spec_hidden_states, - aux_hidden_states, - num_sampled, - num_rejected, - self.req_states.last_sampled_tokens, - self.req_states.next_prefill_tokens, - self.sampler.sampling_states.temperature.gpu, - self.sampler.sampling_states.seeds.gpu, - mm_inputs=mm_inputs, - ) + with use_workspace_lane(self._draft_workspace_lane): + draft_tokens = self.speculator.propose( + input_batch, + attn_metadata, + slot_mappings_by_layer, + spec_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + self.req_states.last_sampled_tokens, + self.req_states.next_prefill_tokens, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, + mm_inputs=mm_inputs, + ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens if self.adaptive_verification is not None: self.adaptive_verification.record_confidences( @@ -1820,6 +1858,7 @@ def sample_tokens( # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) model_runner_output.kv_connector_output = kv_connector_output + model_runner_output.ec_connector_output = ec_connector_output return async_output @@ -1836,6 +1875,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: input_batch = self.execute_model_state.input_batch hidden_states = self.execute_model_state.hidden_states finished_req_ids = self.execute_model_state.finished_req_ids + ec_connector_output = self.execute_model_state.ec_connector_output self.execute_model_state = None # Post-step KV connector related operations. @@ -1843,7 +1883,8 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: if not self.is_last_pp_rank: self.postprocess_num_computed_tokens(input_batch) - return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + output = ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) + return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) assert self.pooling_runner is not None pooler_output, finished_mask = self.pooling_runner.pool( @@ -1855,6 +1896,7 @@ def pool(self) -> AsyncPoolingOutput | ModelRunnerOutput | None: req_ids=input_batch.req_ids, req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)}, kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output, ) async_output = AsyncPoolingOutput( model_runner_output=model_runner_output, @@ -1879,12 +1921,15 @@ def shutdown(self) -> None: """Release GPU tensors (model weights, KV caches, workspace) so that memory is reclaimable when running in the same process.""" torch.accelerator.synchronize() + self.cudagraph_manager = None if hasattr(self, "kv_caches"): self.kv_caches.clear() if hasattr(self, "attn_groups"): self.attn_groups.clear() if hasattr(self, "kv_cache_config"): del self.kv_cache_config + if hasattr(self, "model_state") and self.model_state.supports_mm_inputs: + self.model_state.encoder_runner.clear() free_before_shutdown(self.vllm_config) if hasattr(self, "model_state"): del self.model_state @@ -1941,6 +1986,7 @@ class ExecuteModelState(NamedTuple): hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None finished_req_ids: set[str] + ec_connector_output: ECConnectorOutput | None routed_experts: RoutedExpertsTensors | None diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 2372a01dd300..ed95d8285684 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -3,8 +3,9 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig -from vllm.model_executor.layers.attention import CrossAttention, EncoderOnlyAttention +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.model_executor.layers.attention import Attention, CrossAttention +from vllm.v1.attention.backend import AttentionType from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -27,13 +28,16 @@ def init_model_state( return EncoderDecoderModelState(vllm_config, model, encoder_cache, device) - # Encoder-only models (BERT/RoBERTa): non-causal self-attention, no KV cache. - if any(isinstance(m, EncoderOnlyAttention) for m in model.modules()): + # Encoder-only attention is non-causal and needs no KV cache. + if any( + layer.attn_type == AttentionType.ENCODER_ONLY + for layer in get_layers_from_vllm_config(vllm_config, Attention).values() + ): from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState return EncoderOnlyModelState(vllm_config, model, encoder_cache, device) - if vllm_config.model_config.is_hybrid: + if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free: from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState return MambaHybridModelState(vllm_config, model, encoder_cache, device) diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 722b709bb529..980ec1dd8e1d 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -1,17 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod -from typing import Any +from typing import Any, cast import torch import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.models.interfaces import ( + SupportsEncoderCudaGraph, + supports_encoder_cudagraph, +) from vllm.tasks import GenerationTask from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner @@ -59,6 +64,22 @@ def __init__( self.supports_mm_inputs = encoder_cache is not None if encoder_cache is not None: + enable_encoder_cuda_graph = ( + not self.model_config.enforce_eager + and vllm_config.compilation_config.cudagraph_mm_encoder + and supports_encoder_cudagraph(model) + ) + cudagraph_manager = ( + EncoderCudaGraphManager( + vllm_config=vllm_config, + device=device, + dtype=self.dtype, + model=cast(SupportsEncoderCudaGraph, model), + ) + if enable_encoder_cuda_graph + else None + ) + self.encoder_cache = encoder_cache observability_config = vllm_config.observability_config self.encoder_runner = EncoderRunner( @@ -68,6 +89,7 @@ def __init__( encoder_cache=encoder_cache, dtype=self.dtype, device=self.device, + cudagraph_manager=cudagraph_manager, enable_timing=bool( observability_config and observability_config.enable_mm_processor_stats diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 2715f790dcce..c61a56f9aabf 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -20,6 +20,7 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState from vllm.v1.worker.mamba_utils import ( MambaSpecDecodeGPUContext, preprocess_mamba_align_fused_kernel, @@ -80,6 +81,9 @@ def __init__( # kernel reusing the postprocess copy machinery, so the per-step src # columns and the running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" + self.recoverssm = ( + RecoverSSMState() if self.cache_config.use_kda_recoverssm else None + ) if self._align_mode: self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device @@ -267,7 +271,7 @@ def prepare_attn( num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, ) - return build_attn_metadata( + attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, num_tokens=num_tokens, @@ -285,6 +289,13 @@ def prepare_attn( for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, ) + if self.recoverssm is not None: + self.recoverssm.record_step( + attn_metadata, + attn_groups, + for_capture=for_capture, + ) + return attn_metadata def postprocess_state( self, @@ -295,21 +306,34 @@ def postprocess_state( # Chunked prefill does not sample a token, so num_sampled can be 0. # Mamba treats num_accepted_tokens=1 as the neutral non-spec value. num_reqs = idx_mapping.shape[0] - if not num_reqs: - return + if num_reqs: + if not isinstance(num_sampled, int): + # idx_mapping may contain -1 sentinels (filtered rows) under PP; the + # kernel skips them rather than scattering with a host-side gather. + _scatter_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + num_sampled, + self.num_accepted_tokens_gpu, + ) + else: + # Fill with single value. + _fill_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + self.num_accepted_tokens_gpu, + max(num_sampled, 1), + ) - if not isinstance(num_sampled, int): - # idx_mapping may contain -1 sentinels (filtered rows) under PP; the - # kernel skips them rather than scattering with a host-side gather. - _scatter_num_accepted_kernel[(num_reqs,)]( - idx_mapping, num_sampled, self.num_accepted_tokens_gpu - ) - else: - # Fill with single value. - _fill_num_accepted_kernel[(num_reqs,)]( - idx_mapping, self.num_accepted_tokens_gpu, max(num_sampled, 1) + if self.recoverssm is not None: + self.recoverssm.commit_step( + num_sampled, + idx_mapping, + state_indices=(self._mamba_state_idx_gpu if self._align_mode else None), + num_accepted_tokens=self.num_accepted_tokens_gpu, ) + if not num_reqs: + return + # Align: save the running state to the block-aligned position when # spec-decode acceptance leaves the sequence non-block-aligned (mirrors # the V1 align postprocess). num_computed_tokens already holds the diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py new file mode 100644 index 000000000000..cfbe1a4d45d1 --- /dev/null +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch + +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata +from vllm.v1.worker.utils import AttentionGroup + + +class RecoverSSMState: + """Coordinates RecoverSSM metadata between attention and postprocessing.""" + + def __init__(self) -> None: + self._step: tuple[RecoverSSMMetadata, ...] | None = None + + def record_step( + self, + attn_metadata: dict[str, Any], + attn_groups: list[list[AttentionGroup]], + *, + for_capture: bool, + ) -> None: + if for_capture: + self._step = None + return + + step: list[RecoverSSMMetadata] = [] + for group_list in attn_groups: + for group in group_list: + metadata = attn_metadata[group.layer_names[0]] + if isinstance(metadata, RecoverSSMMetadata): + step.append(metadata) + self._step = tuple(step) + + def commit_step( + self, + num_sampled: torch.Tensor | int, + idx_mapping: torch.Tensor, + *, + state_indices: torch.Tensor | None, + num_accepted_tokens: torch.Tensor, + ) -> None: + step = self._step + self._step = None + if isinstance(num_sampled, int) or step is None: + return + + for metadata in step: + postprocess_meta = metadata.commit_recoverssm_state(num_sampled) + if postprocess_meta is None: + continue + assert state_indices is not None + # RecoverSSM already restored the accepted state. Update its running + # column and reset the next-step copy bias to the neutral value. + _postprocess_recoverssm_align_kernel[(postprocess_meta.num_spec_decodes,)]( + idx_mapping, + num_sampled, + postprocess_meta.request_indices, + postprocess_meta.num_computed_tokens, + state_indices, + num_accepted_tokens, + MAMBA_BLOCK_SIZE=postprocess_meta.block_size, + BLOCK_TABLE_WIDTH=postprocess_meta.block_table.shape[1], + ) + + +@triton.heuristics( + {"HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None} +) +@triton.jit +def _postprocess_recoverssm_align_kernel( + idx_mapping_ptr, + num_sampled_ptr, + request_indices_ptr, + num_computed_ptr, + state_idx_ptr, + num_accepted_ptr, + HAS_REQUEST_INDICES: tl.constexpr, + MAMBA_BLOCK_SIZE: tl.constexpr, + BLOCK_TABLE_WIDTH: tl.constexpr, +): + spec_idx = tl.program_id(0) + batch_idx = spec_idx + if HAS_REQUEST_INDICES: + batch_idx = tl.load(request_indices_ptr + spec_idx) + req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + if req_state_idx < 0: + return + num_sampled = tl.load(num_sampled_ptr + batch_idx) + num_computed = tl.load(num_computed_ptr + batch_idx) + tl.store( + state_idx_ptr + req_state_idx, + tl.minimum( + (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, + BLOCK_TABLE_WIDTH - 1, + ), + ) + tl.store(num_accepted_ptr + req_state_idx, 1) diff --git a/vllm/v1/worker/gpu/sample/bad_words.py b/vllm/v1/worker/gpu/sample/bad_words.py index 768ff30a0f06..3fb9d148ecac 100644 --- a/vllm/v1/worker/gpu/sample/bad_words.py +++ b/vllm/v1/worker/gpu/sample/bad_words.py @@ -153,8 +153,10 @@ def _bad_words_kernel( from_spec_input = actual_pos >= output_len if from_spec_input: + # input_ids at local position 0 is the last committed token; + # draft tokens start at local position 1. spec_offset = actual_pos - output_len - actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset) + actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset + 1) else: actual = tl.load(output_base + actual_pos) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index bb7dc40b0211..26371f014e33 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -53,6 +53,7 @@ def __init__( self.bad_words_state = BadWordsState(req_states) self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device) self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config) + self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool) self.num_speculative_tokens = num_speculative_tokens self.return_sampling_mask = return_sampling_mask self.use_flashinfer = ( @@ -69,6 +70,22 @@ def add_request( self.logprob_token_ids_state.add_request(req_idx, sampling_params) self.thinking_budget_state.add_request(req_idx, sampling_params) + states = self.sampling_states + temperature = states.temperature.np[req_idx] + self.needs_logits_processing[req_idx] = ( + self.logit_bias_state.use_logit_bias[req_idx] + or self.penalties_state.use_penalty[req_idx] + or self.bad_words_state.num_bad_words.np[req_idx] > 0 + or ( + self.thinking_budget_state.enabled + and self.thinking_budget_state.use_thinking_budget[req_idx] + ) + or (temperature != 0.0 and temperature != 1.0) + or states.min_p.np[req_idx] != 0.0 + or states.top_k.np[req_idx] != states.vocab_size + or states.top_p.np[req_idx] != 1.0 + ) + def apply_staged_writes(self) -> None: self.sampling_states.apply_staged_writes() self.penalties_state.apply_staged_writes() @@ -172,7 +189,7 @@ def apply_sampling_params( expanded_local_pos: torch.Tensor, skip_top_k_top_p: bool = False, ) -> torch.Tensor: - if not self._requires_logits_processing(idx_mapping_np): + if not np.any(self.needs_logits_processing[idx_mapping_np]): return logits # Copy logits to a new FP32 tensor. @@ -228,24 +245,6 @@ def apply_sampling_params( logits, expanded_idx_mapping, idx_mapping_np ) - def _requires_logits_processing(self, idx_mapping_np: np.ndarray) -> bool: - if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]): - return True - if np.any(self.penalties_state.use_penalty[idx_mapping_np]): - return True - if np.any(self.bad_words_state.num_bad_words.np[idx_mapping_np] > 0): - return True - - states = self.sampling_states - temperatures = states.temperature.np[idx_mapping_np] - if np.any((temperatures != 0.0) & (temperatures != 1.0)): - return True - if np.any(states.min_p.np[idx_mapping_np] != 0.0): - return True - if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size): - return True - return bool(np.any(states.top_p.np[idx_mapping_np] != 1.0)) - def sample( self, logits: torch.Tensor, diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index c70f169f7be6..4229696f255c 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -26,6 +26,12 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): ) return Gemma4Speculator(vllm_config, device) + elif speculative_config.use_multi_module_mtp(): + from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import ( + MultiModuleMTPSpeculator, + ) + + return MultiModuleMTPSpeculator(vllm_config, device) elif speculative_config.method == "mtp": from vllm.v1.worker.gpu.spec_decode.mtp.speculator import MTPSpeculator diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index e59202c87988..c68b5497458b 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -11,6 +11,7 @@ build_slot_mappings_by_layer, ) from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.cudagraph_utils import ( AttentionState, BatchExecutionDescriptor, @@ -41,6 +42,17 @@ def _prepare_dflash_inputs_to_capture( attn_metadata = None if not skip_attn: query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + dcp_local_seq_lens = None + if block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + input_buffers.dcp_local_seq_lens, + input_buffers.seq_lens, + num_reqs, + block_tables.cp_size, + block_tables.cp_rank, + block_tables.cp_interleave, + ) + dcp_local_seq_lens = input_buffers.dcp_local_seq_lens attn_metadata = build_attn_metadata( attn_groups=attn_groups, num_reqs=num_reqs, @@ -49,6 +61,7 @@ def _prepare_dflash_inputs_to_capture( query_start_loc_cpu=query_start_loc_cpu, max_query_len=num_tokens // num_reqs, seq_lens=input_batch.seq_lens, + dcp_local_seq_lens=dcp_local_seq_lens, max_seq_len=max_model_len, block_tables=input_block_tables, slot_mappings=slot_mappings, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 97c284c03d4a..92837afd68eb 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -17,6 +17,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState @@ -84,8 +85,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.sample_pos = torch.zeros( max_num_sampled_tokens, dtype=torch.int64, device=device ) - self.sample_idx_mapping = torch.zeros( - max_num_sampled_tokens, dtype=torch.int32, device=device + # -1 marks an inert sampling row. CUDA graph capture can execute the + # full buffer before a real batch has populated it, so zero would make + # every padding row scatter into request slot 0. + self.sample_idx_mapping = torch.full( + (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device ) # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into # draft_logits[req, step, :]. @@ -255,7 +259,6 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) - num_sample = num_reqs * self.num_speculative_steps sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] # sample_pos is the predicted token's position Q; verification keys @@ -283,10 +286,21 @@ def _build_draft_attn_metadata( num_query_per_req: int | None = None, causal: bool | Mapping[int, bool] = False, query_start_loc_np: np.ndarray | None = None, + dcp_local_seq_lens: torch.Tensor | None = None, ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None assert num_query_per_req is None # Omitted for DFlash, read from self instead + if dcp_local_seq_lens is None and self.block_tables.cp_size > 1: + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + num_reqs, + self.block_tables.cp_size, + self.block_tables.cp_rank, + self.block_tables.cp_interleave, + ) + dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens return super()._build_draft_attn_metadata( num_reqs, num_reqs_padded, @@ -296,6 +310,7 @@ def _build_draft_attn_metadata( num_query_per_req=self.num_query_per_req, causal=causal, query_start_loc_np=query_start_loc_np, + dcp_local_seq_lens=dcp_local_seq_lens, ) @torch.inference_mode() @@ -391,6 +406,9 @@ def propose( seeds, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], + self.block_tables.cp_rank, + self.block_tables.cp_size, + self.block_tables.cp_interleave, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, @@ -505,8 +523,11 @@ def _prepare_dflash_inputs_kernel( max_num_reqs, max_num_tokens, max_model_len, + cp_rank, SAMPLE_FROM_ANCHOR: tl.constexpr, PAD_SLOT_ID: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_INTERLEAVE: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): req_idx = tl.program_id(0) @@ -520,6 +541,7 @@ def _prepare_dflash_inputs_kernel( num_rejected = tl.load(num_rejected_ptr + req_idx) valid_ctx_end = ctx_end - num_rejected + num_valid_ctx = valid_ctx_end - ctx_start num_sampled = tl.load(num_sampled_ptr + req_idx) if num_sampled > 0: @@ -533,20 +555,37 @@ def _prepare_dflash_inputs_kernel( j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) is_ctx = j < num_ctx - is_query = (j >= num_ctx) & (j < num_ctx + num_query_per_req) - query_off = j - num_ctx + is_valid_ctx = j < num_valid_ctx + is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req) + query_off = j - num_valid_ctx # --- Context positions / slots --- ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) - ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0) - ctx_block_num = ctx_pos // block_size + ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0) + ctx_block_num = ctx_pos // (block_size * CP_SIZE) ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) ctx_block_id = tl.load( block_table_ptr + req_idx * block_table_stride + ctx_block_num, - mask=is_ctx, + mask=is_valid_ctx, other=0, ).to(tl.int64) - ctx_slot = ctx_block_id * block_size + (ctx_pos % block_size) + # Block 0 is the null block. Old sliding-window context positions can map + # to it after eviction; rejected suffix rows are invalid context as well. + # Neither kind of row may write draft KV into physical block 0. + ctx_resident = is_valid_ctx & (ctx_block_id != 0) + local_ctx_slot = cp_local_slot( + ctx_pos, ctx_block_id, block_size, cp_rank, CP_SIZE, CP_INTERLEAVE, PAD_SLOT_ID + ) + ctx_slot = tl.where( + ctx_resident, + local_ctx_slot, + PAD_SLOT_ID, + ) + # Stored over the full [0, num_ctx) span while the loads above are masked to + # [0, num_valid_ctx): the rejected suffix rows in between get position 0 and + # PAD_SLOT_ID. That is intentional — those rows write no KV and their + # positions are never consumed, but the span must stay fully initialized so + # a replayed graph cannot observe a stale value from an earlier batch. tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) @@ -556,14 +595,30 @@ def _prepare_dflash_inputs_kernel( is_bonus = is_query & (query_off == 0) input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id) - q_block_num = query_pos // block_size + q_block_num = query_pos // (block_size * CP_SIZE) q_block_num = tl.minimum(q_block_num, block_table_stride - 1) q_block_id = tl.load( block_table_ptr + req_idx * block_table_stride + q_block_num, mask=is_query, other=0, ).to(tl.int64) - q_slot = q_block_id * block_size + (query_pos % block_size) + # A null block is never a writable cache slot. This can occur when a + # sliding-window block table contains evicted/global padding entries. + q_resident = is_query & (q_block_id != 0) + local_q_slot = cp_local_slot( + query_pos, + q_block_id, + block_size, + cp_rank, + CP_SIZE, + CP_INTERLEAVE, + PAD_SLOT_ID, + ) + q_slot = tl.where( + q_resident, + local_q_slot, + PAD_SLOT_ID, + ) tl.store(out_input_ids_ptr + query_idx, input_id, mask=is_query) clamped_query_pos = tl.minimum(query_pos, max_model_len - 1) @@ -657,6 +712,9 @@ def prepare_dflash_inputs( # [max_num_reqs, max_num_blocks] block_table: torch.Tensor, block_size: int, + cp_rank: int, + cp_size: int, + cp_interleave: int, parallel_drafting_token_id: int, num_query_per_req: int, num_speculative_steps: int, @@ -704,7 +762,10 @@ def prepare_dflash_inputs( max_num_reqs, max_num_tokens, max_model_len, + cp_rank, SAMPLE_FROM_ANCHOR=sample_from_anchor, PAD_SLOT_ID=PAD_SLOT_ID, + CP_SIZE=cp_size, + CP_INTERLEAVE=cp_interleave, BLOCK_SIZE=BLOCK_SIZE, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 478274a05671..b9e6ea02f466 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -13,11 +13,19 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag - from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + from vllm.model_executor.models.qwen3_dflash import ( + dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, + ) speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + # The drafter must rotate Q/K the way its target does. Take that from the + # built target before super() constructs the draft. + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style # Select an attention backend that supports the drafter's attention: mixing # a non-causal layer onto a causal-only backend would fail. draft_vllm_config = replace( @@ -46,7 +54,10 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if hasattr(target_model, "get_language_model") else target_model ) - target_inner = target_language_model.model + # MuseGlimmerForCausalLM marks its inner MuseGlimmerModel as the language + # model, so get_language_model() already returns the inner module and has + # no .model of its own. + target_inner = getattr(target_language_model, "model", target_language_model) draft_inner = dflash_model.model # Skip embedding sharing under PP — each rank owns its own embedding. diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 13c7644315d5..78ae392132ce 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -21,12 +21,18 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal from vllm.model_executor.models.utils import get_draft_quant_config + # None re-runs backend auto-selection for the draft, which can pick a + # different attention class than the target; fall back to the target's. + draft_attention_backend = ( + speculative_config.attention_backend or vllm_config.attention_config.backend + ) + draft_vllm_config = replace( vllm_config, attention_config=replace( vllm_config.attention_config, use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), - backend=speculative_config.attention_backend, + backend=draft_attention_backend, ), cache_config=( replace( diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index dfa2c680109d..34fc5f4bef39 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py @@ -23,6 +23,23 @@ logger = init_logger(__name__) +def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None: + """Copy target KV scales while preserving their tensor representation. + + Default attention scales are scalar buffers, while some quantization + methods replace them with length-one or per-head parameters. Re-register + cloned buffers on the draft layer so the shared KV cache is interpreted + with the target's values and shapes without aliasing target parameters. + """ + for scale_name in ("_k_scale", "_v_scale"): + target_scale = getattr(target_attn, scale_name) + attn.register_buffer(scale_name, target_scale.detach().clone()) + for scale_name in ("_k_scale_float", "_v_scale_float"): + setattr(attn, scale_name, getattr(target_attn, scale_name)) + for scale_name in ("_k_scale_cpu", "_v_scale_cpu"): + getattr(attn, scale_name).copy_(getattr(target_attn, scale_name)) + + class Gemma4Speculator(AutoRegressiveSpeculator): @property def advance_draft_positions(self) -> bool: @@ -76,7 +93,7 @@ def _setup_gemma4_kv_sharing( model: nn.Module, target_attn_layer_names: set[str], ) -> None: - """Wire draft layers to share KV with the target model. + """Wire draft layers to share KV and KV scales with the target model. Each draft decoder layer is mapped to the last non-KV-shared target layer of the same attention type (sliding or full). @@ -92,15 +109,19 @@ def _setup_gemma4_kv_sharing( target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0) num_non_shared = len(target_layer_types) - target_num_kv_shared - type_to_target_indices: dict[str, list[int]] = defaultdict(list) - for idx, lt in enumerate(target_layer_types[:num_non_shared]): - type_to_target_indices[lt].append(idx) - - target_prefix = "model.layers" + target_names_by_index: dict[int, str] = {} for name in target_attn_layer_names: - if ".layers." in name: - target_prefix = name.split(".layers.")[0] + ".layers" - break + _, separator, layer_suffix = name.partition(".layers.") + if not separator: + continue + layer_index, _, _ = layer_suffix.partition(".") + if layer_index.isdigit(): + target_names_by_index[int(layer_index)] = name + + type_to_target_names: dict[str, list[str]] = defaultdict(list) + for idx, lt in enumerate(target_layer_types[:num_non_shared]): + if target_name := target_names_by_index.get(idx): + type_to_target_names[lt].append(target_name) draft_layer_types = getattr(draft_text_config, "layer_types", []) for draft_idx, layer in enumerate(model.model.layers): @@ -115,7 +136,7 @@ def _setup_gemma4_kv_sharing( if draft_idx < len(draft_layer_types) else "full_attention" ) - candidates = type_to_target_indices.get(draft_layer_type, []) + candidates = type_to_target_names.get(draft_layer_type, []) if not candidates: logger.warning( "No target layer of type '%s' for draft layer %d", @@ -124,9 +145,19 @@ def _setup_gemma4_kv_sharing( ) continue - target_idx = candidates[-1] - target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn" + target_layer_name = candidates[-1] attn.kv_sharing_target_layer_name = target_layer_name + + # KV-cache sharing aliases the cache tensor during allocation, but + # the quantization scales live on the Attention modules themselves. + # The BF16 draft model has no quantization config, so its K/V scales + # otherwise remain at the default 1.0 while it reads the target's + # calibrated FP8 cache. + target_attn = self.vllm_config.compilation_config.static_forward_context[ + target_layer_name + ] + _copy_target_kv_scales(attn, target_attn) + logger.info( "Gemma4 MTP: draft layer %d (%s) -> %s", draft_idx, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 69e0cc109160..80de83109bf5 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -234,6 +234,7 @@ def _build_draft_attn_metadata( num_query_per_req: int = 1, causal: bool | Mapping[int, bool] = True, query_start_loc_np: np.ndarray | None = None, + dcp_local_seq_lens: torch.Tensor | None = None, ) -> dict[str, Any] | None: if query_start_loc_np is not None: # Non-uniform query layout (e.g. multi-module MTP's mixed @@ -278,6 +279,11 @@ def _build_draft_attn_metadata( query_start_loc_cpu=query_start_loc_cpu, max_query_len=max_query_len, seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], + dcp_local_seq_lens=( + None + if dcp_local_seq_lens is None + else dcp_local_seq_lens[:num_reqs_padded] + ), max_seq_len=self.draft_max_seq_len, block_tables=block_tables, slot_mappings=slot_mappings, diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 221163ae3652..8437b0aed9a4 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -10,6 +10,32 @@ from vllm.v1.worker.gpu.input_batch import InputBatch +def _build_grammar_mapping( + req_ids: list[str], + grammar_req_ids: list[str], + cu_num_logits_np: np.ndarray, + num_draft_tokens_per_req: np.ndarray | None, + num_bonus_tokens: int, + mask_stride: int, +) -> list[int]: + mapping: list[int] = [] + req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} + for grammar_req_id in grammar_req_ids: + req_idx = req_id_to_idx[grammar_req_id] + if num_draft_tokens_per_req is None: + num_positions = int( + cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx] + ) + else: + # Grammar masks follow the scheduled layout even when adaptive + # verification compacts the actual CPU logit offsets to bonus-only. + num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens + mapping.extend( + req_idx * mask_stride + position for position in range(num_positions) + ) + return mapping + + class StructuredOutputsWorker: def __init__( self, @@ -17,6 +43,7 @@ def __init__( vocab_size: int, device: torch.device, mask_stride: int, + num_bonus_tokens: int, ): self.logits_indices = torch.zeros( max_num_logits, dtype=torch.int32, device=device @@ -27,6 +54,7 @@ def __init__( self.device = device self.copy_stream = torch.cuda.Stream() self.mask_stride = mask_stride + self.num_bonus_tokens = num_bonus_tokens def apply_grammar_bitmask( self, @@ -45,21 +73,17 @@ def apply_grammar_bitmask( ) # Construct bitmask -> logits mapping - mapping: list[int] = [] - req_ids = input_batch.req_ids - cu_num_logits = input_batch.cu_num_logits_np.tolist() - req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} - for grammar_req_id in grammar_req_ids: - req_idx = req_id_to_idx[grammar_req_id] - logits_start_idx = cu_num_logits[req_idx] - logits_end_idx = cu_num_logits[req_idx + 1] - # Key by (request, position) rather than absolute logit index: - # adaptive verification finalizes per-request logit offsets on - # device, so the kernel resolves them from the GPU cu_num_logits. - mapping.extend( - req_idx * self.mask_stride + position - for position in range(logits_end_idx - logits_start_idx) - ) + # Key by (request, position) rather than absolute logit index: + # adaptive verification finalizes per-request logit offsets on + # device, so the kernel resolves them from the GPU cu_num_logits. + mapping = _build_grammar_mapping( + input_batch.req_ids, + grammar_req_ids, + input_batch.cu_num_logits_np, + input_batch.num_draft_tokens_per_req, + self.num_bonus_tokens, + self.mask_stride, + ) # Asynchronously copy the mapping to GPU. with torch.cuda.stream(self.copy_stream): diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index defaccacf675..a8d08a7bf293 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -100,6 +100,7 @@ get_offloader, set_offloader, ) +from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.inputs import ( @@ -120,6 +121,7 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask from vllm.tracing import instrument from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks @@ -515,6 +517,7 @@ def __init__( self.scheduler_config = vllm_config.scheduler_config self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) model_config = self.model_config cache_config = self.cache_config @@ -586,6 +589,9 @@ def __init__( # Async scheduling self.use_async_scheduling = self.scheduler_config.async_scheduling + # Async PP broadcast of sampled token ids, waited on in _prepare_input_ids. + self._pp_recv_work: torch.distributed.Work | None = None + # Sampler self.sampler = Sampler( logprobs_mode=self.model_config.logprobs_mode, @@ -746,34 +752,37 @@ def __init__( self._init_kernel_block_sizes = [placeholder_block_size] self._init_max_num_blocks = [placeholder_max_num_blocks] self._init_slot_mapping_modes = [SlotMappingMode.TOKEN_TO_KV_SLOT] - self.input_batch = InputBatch( - max_num_reqs=self.max_num_reqs, - # We need to use the encoder length for encoder-decoder - # because of KV cache for cross-attention. - max_model_len=max(self.max_model_len, self.max_encoder_len), - max_num_batched_tokens=self.max_num_tokens, - device=self.device, - vocab_size=self.model_config.get_vocab_size(), - block_sizes=[placeholder_block_size], - kernel_block_sizes=[placeholder_block_size], - max_num_blocks_per_req=[placeholder_max_num_blocks], - num_spec_tokens=self.num_spec_tokens, - logitsprocs=build_logitsprocs( - self.vllm_config, - self.device, - PIN_MEMORY, - self.is_pooling_model, - custom_logitsprocs, - ), - # We currently don't know whether a particular custom logits processor - # uses output token ids so we set this conservatively. Thinking-budget - # tracking is requested dynamically when a budgeted request is in the batch. - logitsprocs_need_output_token_ids=bool(custom_logitsprocs), - is_pooling_model=self.is_pooling_model, - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, - reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, - ) + # Capture warmup providers registered by the initial placeholder InputBatch + with self.jit_warmup_registry.activate(): + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + # We need to use the encoder length for encoder-decoder + # because of KV cache for cross-attention. + max_model_len=max(self.max_model_len, self.max_encoder_len), + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=[placeholder_block_size], + kernel_block_sizes=[placeholder_block_size], + max_num_blocks_per_req=[placeholder_max_num_blocks], + num_spec_tokens=self.num_spec_tokens, + logitsprocs=build_logitsprocs( + self.vllm_config, + self.device, + PIN_MEMORY, + self.is_pooling_model, + custom_logitsprocs, + ), + # We currently don't know whether a particular custom logits processor + # uses output token ids so we set this conservatively. Thinking-budget + # tracking is requested dynamically when a budgeted request is in the + # batch. + logitsprocs_need_output_token_ids=bool(custom_logitsprocs), + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + ) # Separate cuda stream for overlapping transfer of sampled token ids from # GPU to CPU when async scheduling is enabled. @@ -1843,6 +1852,11 @@ def _prepare_input_ids( (-1 for new requests). """ + # Sync the async PP broadcast before reading sampled tokens. + if self._pp_recv_work is not None: + self._pp_recv_work.wait() + self._pp_recv_work = None + if self.input_batch.prev_sampled_token_ids is None: # Normal scheduling case self.input_ids.copy_to_gpu(total_num_scheduled_tokens) @@ -3123,18 +3137,24 @@ def _execute_mm_encoder( token_lora_mapping = [] lora_requests = set() encoder_token_counts = [] + connector_token_counts = [] - for req_id, pos_info in mm_lora_refs: + for (req_id, pos_info), (modality, mm_item) in zip( + mm_lora_refs, + mm_kwargs, + ): req_idx = self.input_batch.req_id_to_index[req_id] lora_id = int(self.input_batch.request_lora_mapping[req_idx]) - # Prefer pos_info.get_num_embeds to count precise MM embedding tokens. - num_tokens = self.model.get_num_mm_encoder_tokens( # type: ignore[attr-defined] - pos_info.get_num_embeds() + tower_tokens, connector_tokens = self.model.get_mm_lora_token_counts( # type: ignore[attr-defined] + modality=modality, + mm_kwargs=mm_item, + num_mm_embeds=pos_info.get_num_embeds(), ) prompt_lora_mapping.append(lora_id) - token_lora_mapping.extend([lora_id] * num_tokens) - encoder_token_counts.append(num_tokens) + token_lora_mapping.extend([lora_id] * tower_tokens) + encoder_token_counts.append(tower_tokens) + connector_token_counts.append(connector_tokens) if lora_id > 0: lora_request = self.input_batch.lora_id_to_lora_request.get(lora_id) @@ -3162,16 +3182,11 @@ def _execute_mm_encoder( if ( mm_mapping is not None and mm_mapping.connector - and hasattr(self.model, "get_num_mm_connector_tokens") + and all(count is not None for count in connector_token_counts) ): - post_op_counts = [ - self.model.get_num_mm_connector_tokens(num_tokens) # type: ignore[attr-defined] - for num_tokens in encoder_token_counts - ] - connector_token_mapping = np.repeat( np.array(prompt_lora_mapping, dtype=np.int32), - np.array(post_op_counts, dtype=np.int32), + np.array(connector_token_counts, dtype=np.int32), ) connector_mapping = LoRAMapping( index_mapping=tuple(connector_token_mapping.tolist()), @@ -3839,26 +3854,28 @@ def _bookkeeping_sync( self.routed_experts_slot_mapping_device[:total], non_blocking=True, ) - - # Get the valid generated tokens. - max_gen_len = sampled_token_ids.shape[-1] - if max_gen_len == 1: - # No spec decode tokens. - valid_sampled_token_ids = self._to_list(sampled_token_ids) - # Mask out the sampled tokens that should not be sampled. - for i in discard_sampled_tokens_req_indices: - valid_sampled_token_ids[int(i)].clear() - - if logprobs_tensors is not None: - logprobs_lists = logprobs_tensors.tolists() - else: - # Includes spec decode tokens. - valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output( - sampled_token_ids, - self.input_batch.vocab_size, - discard_sampled_tokens_req_indices, - logprobs_tensors=logprobs_tensors, - ) + with gpu_sync_allowed(): + # Get the valid generated tokens. + max_gen_len = sampled_token_ids.shape[-1] + if max_gen_len == 1: + # No spec decode tokens. + valid_sampled_token_ids = self._to_list(sampled_token_ids) + # Mask out the sampled tokens that should not be sampled. + for i in discard_sampled_tokens_req_indices: + valid_sampled_token_ids[int(i)].clear() + + if logprobs_tensors is not None: + logprobs_lists = logprobs_tensors.tolists() + else: + # Includes spec decode tokens. + valid_sampled_token_ids, logprobs_lists = ( + RejectionSampler.parse_output( + sampled_token_ids, + self.input_batch.vocab_size, + discard_sampled_tokens_req_indices, + logprobs_tensors=logprobs_tensors, + ) + ) else: valid_sampled_token_ids = [] invalid_req_indices = discard_sampled_tokens_req_indices.tolist() @@ -4962,7 +4979,9 @@ def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None: recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device) # skip for chunked prefill. if not self._is_all_reqs_chunked_prefill(): - torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) + self._pp_recv_work = torch.distributed.broadcast( + recv, src=pp.last_rank, group=pp.device_group, async_op=True + ) self.input_batch.prev_sampled_token_ids = recv # construct `prev_req_id_to_index` here so `_prepare_input_ids` @@ -5832,7 +5851,10 @@ def _get_nans_in_logits(self, logits: torch.Tensor | None) -> dict[str, int]: on device instead; see`AsyncGPUModelRunnerOutput`. """ try: - counts = [] if logits is None else count_nans_per_row(logits).tolist() + # Reporting per-request NaN counts requires them on the host; this + # path is opt-in diagnostics, so the D2H is intended. + with gpu_sync_allowed(): + counts = [] if logits is None else count_nans_per_row(logits).tolist() num_nans_in_logits = nans_to_dict(counts, self.input_batch.req_id_to_index) if envs.VLLM_RAISE_ON_LOGIT_NANS: raise_if_nan_logits(num_nans_in_logits) @@ -7399,24 +7421,26 @@ def may_reinitialize_input_batch( self._init_kernel_block_sizes = kernel_block_sizes self._init_max_num_blocks = max_num_blocks self._init_slot_mapping_modes = slot_mapping_modes - self.input_batch = InputBatch( - max_num_reqs=self.max_num_reqs, - max_model_len=max_model_len, - max_num_batched_tokens=self.max_num_tokens, - device=self.device, - vocab_size=self.model_config.get_vocab_size(), - block_sizes=block_sizes, - kernel_block_sizes=kernel_block_sizes, - max_num_blocks_per_req=max_num_blocks, - num_spec_tokens=self.num_spec_tokens, - logitsprocs=self.input_batch.logitsprocs, - logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, - is_pooling_model=self.is_pooling_model, - cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, - reasoning_config=self.vllm_config.reasoning_config, - use_replayssm=self.cache_config.use_replayssm, - slot_mapping_modes=slot_mapping_modes, - ) + # Capture warmup providers registered after final KV-cache geometry is known + with self.jit_warmup_registry.activate(): + self.input_batch = InputBatch( + max_num_reqs=self.max_num_reqs, + max_model_len=max_model_len, + max_num_batched_tokens=self.max_num_tokens, + device=self.device, + vocab_size=self.model_config.get_vocab_size(), + block_sizes=block_sizes, + kernel_block_sizes=kernel_block_sizes, + max_num_blocks_per_req=max_num_blocks, + num_spec_tokens=self.num_spec_tokens, + logitsprocs=self.input_batch.logitsprocs, + logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids, + is_pooling_model=self.is_pooling_model, + cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size, + reasoning_config=self.vllm_config.reasoning_config, + use_replayssm=self.cache_config.use_replayssm, + slot_mapping_modes=slot_mapping_modes, + ) assert self._init_block_sizes == block_sizes, ( f"InputBatch block_sizes {self._init_block_sizes} != " @@ -7914,6 +7938,7 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: with set_current_vllm_config(self.vllm_config): indexes = backend.indexes_kv_by_block_stride() spec = replace(spec, indexes_kv_by_block_stride=indexes) + spec = backend.customize_spec(spec) kv_cache_spec[layer_name] = spec return kv_cache_spec diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 2340b64abaac..a3b00aaad2a2 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -91,6 +91,16 @@ logger = init_logger(__name__) + +def _num_workspace_lanes(vllm_config: VllmConfig, use_v2_model_runner: bool) -> int: + spec_config = vllm_config.speculative_config + return ( + 2 + if use_v2_model_runner and spec_config is not None and spec_config.use_dspark() + else 1 + ) + + if TYPE_CHECKING: from vllm.device_allocator.sleep_mode_backend import SleepModeBackend from vllm.model_executor.model_loader.tensorizer import TensorizerConfig @@ -402,9 +412,13 @@ def init_device(self): else: raise RuntimeError(f"Unsupported device type: {self.device_config.device}") - # Initialize workspace manager + # DSpark target and draft CUDA graphs retain workspace views concurrently. num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 - init_workspace_manager(self.device, num_ubatches) + init_workspace_manager( + self.device, + num_ubatches, + _num_workspace_lanes(self.vllm_config, self.use_v2_model_runner), + ) # Construct the model runner if self.use_v2_model_runner: @@ -869,6 +883,19 @@ def get_model(self) -> nn.Module: def get_draft_model(self) -> nn.Module | None: return self.model_runner.get_draft_model() + def supports_draft_weight_updates(self) -> bool: + engine = self.weight_transfer_engine + speculative_config = self.speculative_config + get_draft_model = getattr(self.model_runner, "get_draft_model", None) + return ( + engine is not None + and engine.supports_draft_weight_update + and callable(get_draft_model) + and get_draft_model() is not None + and speculative_config is not None + and speculative_config.draft_model_config is not None + ) + def _set_draft_weight_update_target(self) -> None: assert self.weight_transfer_engine is not None diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index f87a34a364f9..a0b89303dd86 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -16,6 +16,7 @@ is_conv_state_dim_first, ) from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec @@ -188,19 +189,46 @@ def _copy_mamba_state_block( src_block_id = tl.load(block_table_base + src_col).to(tl.int64) dim_rows = tl.load(state_dim_row_count_ptr + state_idx) row_stride = tl.load(state_dim_row_stride_ptr + state_idx) - per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size - bias_bytes = token_bias.to(tl.int64) * state_elem_size src_block_addr = state_base_addr + src_block_id * state_block_stride offsets = tl.arange(0, COPY_BLOCK_SIZE) - for d in range(0, dim_rows): - row_src = src_block_addr + d * row_stride + bias_bytes - row_dst = dst_addr + d * row_stride - for i in range(0, per_row_bytes, COPY_BLOCK_SIZE): - mask = (i + offsets) < per_row_bytes - curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8)) - curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(curr_src, mask=mask) - tl.store(curr_dst, data, mask=mask) + + # Stable row-to-lane ownership makes left shifts memmove-safe while + # exposing the dimension rows in parallel. All addresses retain + # state_elem_size alignment: tensor strides and token offsets are + # measured in whole elements before conversion to bytes. + num_dst_tokens = conv_width - token_bias + for token_idx in range(0, num_dst_tokens): + for row_base in range(0, dim_rows, COPY_BLOCK_SIZE): + rows = row_base + offsets + mask = rows < dim_rows + src_byte_addr = ( + src_block_addr + + rows * row_stride + + (token_idx + token_bias) * state_elem_size + ) + dst_byte_addr = ( + dst_addr + rows * row_stride + token_idx * state_elem_size + ) + if state_elem_size == 2: + src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16)) + dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16)) + data_u16 = tl.load(src_u16, mask=mask) + tl.store(dst_u16, data_u16, mask=mask) + elif state_elem_size == 4: + src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32)) + dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32)) + data_u32 = tl.load(src_u32, mask=mask) + tl.store(dst_u32, data_u32, mask=mask) + else: + for byte_idx in range(0, state_elem_size): + src_u8 = (src_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + dst_u8 = (dst_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + data_u8 = tl.load(src_u8, mask=mask) + tl.store(dst_u8, data_u8, mask=mask) return if is_conv_state: @@ -209,22 +237,40 @@ def _copy_mamba_state_block( # SD conv: copy # state[bt[src_col], token_bias:] -> # state[bt[dst_col], :conv_width - token_bias] - # Small per-block bytes (~60-80 KiB) make tiling degenerate, so - # conv runs as a single-CTA memcpy (NUM_TILES=1). src_block_id = tl.load(block_table_base + src_col).to(tl.int64) - src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size - src_addr = state_base_addr + src_block_id * state_block_stride + src_offset - copy_size = ( - (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size - ) - _memcpy_u64_tiled( - src_addr, - dst_addr, - copy_size, - tile_idx, - COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, - NUM_TILES=1, - ) + src_block_addr = state_base_addr + src_block_id * state_block_stride + token_bytes = state_inner_size * state_elem_size + num_dst_tokens = conv_width - token_bias + + # Distinct blocks and exact self-copies cannot have a destructive + # overlap, so retain the u64-vectorized single-CTA copy. + if src_block_id != dest_block_id or token_bias == 0: + src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes + copy_size = num_dst_tokens.to(tl.int64) * token_bytes + _memcpy_u64_tiled( + src_addr, + dst_addr, + copy_size, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) + return + + # Copy tokens from low to high. Each token-sized source and destination + # region is disjoint, so same-block left shifts are memmove-safe + # without a barrier. + for token_idx in range(0, num_dst_tokens): + src_token = src_block_addr + (token_idx + token_bias) * token_bytes + dst_token = dst_addr + token_idx * token_bytes + _memcpy_u64_tiled( + src_token, + dst_token, + token_bytes, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) return # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] @@ -521,6 +567,7 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): src_ptr = tl.load(src_ptrs + pid) dst_ptr = tl.load(dst_ptrs + pid) size = tl.load(sizes + pid) + is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr offsets = tl.arange(0, BLOCK_SIZE) for i in range(0, size, BLOCK_SIZE): @@ -530,6 +577,10 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) data = tl.load(curr_src_ptr, mask=mask) + if is_left_overlap: + # Preserve each lane's source before a lower-address lane stores + # over it. The condition is uniform within the program. + tl.debug_barrier() tl.store(curr_dst_ptr, data, mask=mask) @@ -754,7 +805,22 @@ def initialize_from_forward_context( """ if self.is_initialized: return + # This only runs once per worker. + with gpu_sync_allowed(): + self._populate_metadata( + kv_cache_config, + forward_context, + mamba_state_copy_funcs, + block_tables, + ) + def _populate_metadata( + self, + kv_cache_config: KVCacheConfig, + forward_context: dict[str, Any], + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + block_tables: list[torch.Tensor], + ) -> None: idx = 0 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names @@ -762,7 +828,13 @@ def initialize_from_forward_context( attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache - for state_type_idx, state in enumerate(kv_caches): + if len(kv_caches) < self.num_state_types: + raise ValueError( + f"Expected at least {self.num_state_types} Mamba state " + f"tensors, got {len(kv_caches)}" + ) + for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs): + state = kv_caches[state_type_idx] # Base address self.state_base_addrs[idx] = state.data_ptr() @@ -779,7 +851,6 @@ def initialize_from_forward_context( # Element size self.state_elem_sizes[idx] = state.element_size() - copy_func = mamba_state_copy_funcs[state_type_idx] assert ( copy_func is get_conv_copy_spec or copy_func is get_temporal_copy_spec diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 46a73935df8b..d8f36b1d154a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -538,17 +538,7 @@ def bind_kv_cache( # TODO - analyze where runner_kv_caches is used and the right # way to ensure it properly reflects multiple attention layers # in the same decoder block. - if ( - current_platform.is_cuda_alike() - or current_platform.is_xpu() - or current_platform.is_cpu() - ): - # We know that the GPU / CPU runner is not impacted by this - # case. Some test code depends on runner_kv_caches, but - # not in a way that's impacted by ignoring this. - pass - else: - raise NotImplementedError + current_platform.check_runner_kv_caches_multi_layer() for layer_name in layer_names: runner_kv_caches.append(kv_caches[layer_name]) diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 9381d71913dc..998380a8773a 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -125,6 +125,10 @@ def reset_mm_cache(self) -> None: def get_model(self) -> nn.Module: raise NotImplementedError + def supports_draft_weight_updates(self) -> bool: + """Whether this worker can update its configured speculative model.""" + return False + def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R: """Apply a function on the model inside this worker.""" return fn(self.get_model()) diff --git a/vllm/v1/worker/workspace.py b/vllm/v1/worker/workspace.py index 1c502bfd8ff1..39b3b10349c3 100644 --- a/vllm/v1/worker/workspace.py +++ b/vllm/v1/worker/workspace.py @@ -3,6 +3,9 @@ import inspect import os +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar from itertools import accumulate from math import prod @@ -26,22 +29,43 @@ def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int: # Global workspace manager instance _manager: "WorkspaceManager | None" = None +_workspace_lane: ContextVar[int] = ContextVar("vllm_workspace_lane", default=0) + + +@contextmanager +def use_workspace_lane(lane: int) -> Iterator[None]: + """Select an independent workspace owner for this execution context.""" + if lane < 0: + raise ValueError(f"Workspace lane must be non-negative, got {lane}.") + token = _workspace_lane.set(lane) + try: + yield + finally: + _workspace_lane.reset(token) class WorkspaceManager: """Manager for workspace allocation. - Manages one workspace buffer per active ubatch slot. + Manages one workspace buffer per active ``(ubatch, lane)`` slot. Can be locked to prevent further growth during execution. """ - def __init__(self, device: torch.device, num_ubatches: int | None = None): + def __init__( + self, + device: torch.device, + num_ubatches: int | None = None, + num_lanes: int = 1, + ): self._device = device # Cache num ubatches at init based on configuration (default to 1) self._num_ubatches = num_ubatches if num_ubatches is not None else 1 - self._current_workspaces: list[torch.Tensor | None] = [ - None - ] * self._num_ubatches + if num_lanes < 1: + raise ValueError(f"num_lanes must be at least one, got {num_lanes}.") + self._num_lanes = num_lanes + self._current_workspaces: list[torch.Tensor | None] = [None] * ( + self._num_ubatches * self._num_lanes + ) self._locked: bool = False @staticmethod @@ -126,7 +150,14 @@ def _ensure_workspace_size(self, required_bytes: int) -> torch.Tensor: The current workspace tensor. """ ubatch_id = dbo_current_ubatch_id() - current_workspace = self._current_workspaces[ubatch_id] + lane = _workspace_lane.get() + if lane >= self._num_lanes: + raise RuntimeError( + f"Workspace lane {lane} is not configured; manager has " + f"{self._num_lanes} lane(s)." + ) + workspace_id = ubatch_id * self._num_lanes + lane + current_workspace = self._current_workspaces[workspace_id] current_size = self._workspace_size_bytes(current_workspace) if current_size < required_bytes: @@ -161,11 +192,11 @@ def get_caller_info() -> str: "Workspace growth is not allowed after locking." ) - # Only resize the requesting ubatch's workspace. Other - # ubatches resize lazily on their next get_simultaneous call. + # Only resize the requesting ubatch/lane workspace. Other slots + # resize lazily on their next get_simultaneous call. # Resizing all ubatches here would orphan the other ubatch's # old tensor when it still holds views into it (DBO leak). - self._current_workspaces[ubatch_id] = None + self._current_workspaces[workspace_id] = None del current_workspace # Release the freed segment back to CUDA so the caching # allocator can reuse the GPU memory for the larger @@ -173,19 +204,20 @@ def get_caller_info() -> str: # dead segment in reserved memory which can cause higher peak # memory usage. torch.accelerator.empty_cache() - self._current_workspaces[ubatch_id] = torch.empty( + self._current_workspaces[workspace_id] = torch.empty( (required_bytes,), dtype=torch.uint8, device=self._device ) - current_workspace = self._current_workspaces[ubatch_id] + current_workspace = self._current_workspaces[workspace_id] if envs.VLLM_DEBUG_WORKSPACE: logger.info( "[WORKSPACE DEBUG] Resized workspace from '%s': %.2f MB -> " - "%.2f MB (ubatch %d)", + "%.2f MB (ubatch %d, lane %d)", get_caller_info(), current_size / _MB, required_bytes / _MB, ubatch_id, + lane, ) return current_workspace @@ -214,7 +246,9 @@ def current_workspace_manager() -> "WorkspaceManager": def init_workspace_manager( - device: torch.device, num_ubatches: int | None = None + device: torch.device, + num_ubatches: int | None = None, + num_lanes: int = 1, ) -> None: """Initialize the workspace manager with a device. @@ -224,6 +258,7 @@ def init_workspace_manager( Args: device: The device to allocate workspace on. num_ubatches: Number of workspace ubatch slots. Defaults to 1. + num_lanes: Number of independent execution lanes per ubatch. Defaults to 1. """ global _manager if _manager is not None: @@ -233,7 +268,7 @@ def init_workspace_manager( _manager._device, device, ) - _manager = WorkspaceManager(device, num_ubatches) + _manager = WorkspaceManager(device, num_ubatches, num_lanes) def lock_workspace() -> None: diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 81768e12092e..aba0d297272c 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -404,6 +404,22 @@ def flash_attn_varlen_func( from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd + # SM90 FA4 fp8-KV path: fp8 e4m3 paged K/V dequantized and the K/V descale folded + # in-kernel; accepts bf16/fp16 Q and writes O in its native dtype (no Q cast, no + # output copy). Only the (batch, num_kv_heads) f32 K/V descales are forwarded. + fa4_fp8_kv_dequant = ( + k.dtype == torch.float8_e4m3fn + and torch.cuda.get_device_capability()[0] == 9 + ) + if fa4_fp8_kv_dequant: + fa4_q_descale = None + fa4_k_descale = k_descale + fa4_v_descale = v_descale + else: + fa4_q_descale = None + fa4_k_descale = None + fa4_v_descale = None + out, softmax_lse, _, _ = _flash_attn_fwd( q, k, @@ -428,10 +444,11 @@ def flash_attn_varlen_func( block_sparse_tensors=block_sparse_tensors, aux_tensors=aux_tensors, aux_tensor_leading_dims=aux_tensor_leading_dims, - q_descale=q_descale, - k_descale=k_descale, - v_descale=v_descale, + q_descale=fa4_q_descale, + k_descale=fa4_k_descale, + v_descale=fa4_v_descale, output_scale=output_scale, + fp8_kv_dequant=fa4_fp8_kv_dequant, ) else: raise ValueError(f"Unsupported FA version: {fa_version}")