diff --git a/.buildkite/intel_jobs/model_executor_intel.yaml b/.buildkite/intel_jobs/model_executor_intel.yaml index d7c683f8ce..98e4b0ff45 100644 --- a/.buildkite/intel_jobs/model_executor_intel.yaml +++ b/.buildkite/intel_jobs/model_executor_intel.yaml @@ -28,4 +28,13 @@ steps: export VLLM_WORKER_MULTIPROC_METHOD=spawn && export PYTHONFAULTHANDLER=1 && cd tests && - pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload"' + pytest -v -s model_executor -m "not slow_test" --ignore="model_executor/layers/test_rocm_unquantized_gemm.py" --deselect="tests/model_executor/model_loader/test_reload.py::test_kv_scale_reload" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormModule::test_matches_reference_above_cudnn_grid_limit" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype0-out_dtype0]" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype1-out_dtype1]" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype2-out_dtype2]" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype3-out_dtype3]" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype4-out_dtype4]" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormDtypes::test_dtype_combinations[in_dtype5-out_dtype5]" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormShapes" + --deselect="tests/model_executor/layers/test_mm_input_norm.py::TestFusedMMInputNormInputHandling"' diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index 58ecb2d282..eb36e0c02e 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -72,8 +72,7 @@ steps: - >- bash .buildkite/scripts/hardware_ci/run-intel-test.sh 'cd tests && - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma" && - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model' + pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "not qwen2 and not qwen3 and not gemma"' - label: "Multi-Modal Models (Standard) 4: other + whisper" key: multi-modal-models-standard-4-other-whisper diff --git a/.buildkite/performance-benchmarks/tests/latency-tests-arm64-cpu.json b/.buildkite/performance-benchmarks/tests/latency-tests-arm64-cpu.json index 98811049d3..d0693a19dc 100644 --- a/.buildkite/performance-benchmarks/tests/latency-tests-arm64-cpu.json +++ b/.buildkite/performance-benchmarks/tests/latency-tests-arm64-cpu.json @@ -3,7 +3,6 @@ "test_name": "latency_llama8B_tp1", "environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_KVCACHE_SPACE": 40 }, "parameters": { diff --git a/.buildkite/performance-benchmarks/tests/latency-tests-cpu.json b/.buildkite/performance-benchmarks/tests/latency-tests-cpu.json index 5f048df5f6..07e58346fc 100644 --- a/.buildkite/performance-benchmarks/tests/latency-tests-cpu.json +++ b/.buildkite/performance-benchmarks/tests/latency-tests-cpu.json @@ -3,7 +3,6 @@ "test_name": "latency_llama8B_tp2", "environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_KVCACHE_SPACE": 40 }, diff --git a/.buildkite/performance-benchmarks/tests/serving-tests-arm64-cpu.json b/.buildkite/performance-benchmarks/tests/serving-tests-arm64-cpu.json index 75b80b2c21..d68837fa96 100644 --- a/.buildkite/performance-benchmarks/tests/serving-tests-arm64-cpu.json +++ b/.buildkite/performance-benchmarks/tests/serving-tests-arm64-cpu.json @@ -14,7 +14,6 @@ ], "server_environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_KVCACHE_SPACE": 40 }, diff --git a/.buildkite/performance-benchmarks/tests/serving-tests-cpu-asr.json b/.buildkite/performance-benchmarks/tests/serving-tests-cpu-asr.json index 16e8b0600a..756f0f77f2 100644 --- a/.buildkite/performance-benchmarks/tests/serving-tests-cpu-asr.json +++ b/.buildkite/performance-benchmarks/tests/serving-tests-cpu-asr.json @@ -4,9 +4,6 @@ "inf" ], "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200], - "server_environment_variables": { - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120 - }, "server_parameters": { "dtype": "bfloat16", "model": "openai/whisper-large-v3-turbo" diff --git a/.buildkite/performance-benchmarks/tests/serving-tests-cpu-embed.json b/.buildkite/performance-benchmarks/tests/serving-tests-cpu-embed.json index c62f244fc7..1be4c653f1 100644 --- a/.buildkite/performance-benchmarks/tests/serving-tests-cpu-embed.json +++ b/.buildkite/performance-benchmarks/tests/serving-tests-cpu-embed.json @@ -10,7 +10,6 @@ ], "server_environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_KVCACHE_SPACE": 40 }, diff --git a/.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json b/.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json index 9aa76c1108..acc3669dff 100644 --- a/.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json +++ b/.buildkite/performance-benchmarks/tests/serving-tests-cpu-text.json @@ -6,7 +6,6 @@ "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200], "server_environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_KVCACHE_SPACE": 40 }, diff --git a/.buildkite/performance-benchmarks/tests/serving-tests-cpu.json b/.buildkite/performance-benchmarks/tests/serving-tests-cpu.json index 0b7e749996..f1b2f4747a 100644 --- a/.buildkite/performance-benchmarks/tests/serving-tests-cpu.json +++ b/.buildkite/performance-benchmarks/tests/serving-tests-cpu.json @@ -6,7 +6,6 @@ "max_concurrency_list": [12, 16, 24, 32, 64, 128, 200], "server_environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_KVCACHE_SPACE": 40 }, diff --git a/.buildkite/performance-benchmarks/tests/throughput-tests-arm64-cpu.json b/.buildkite/performance-benchmarks/tests/throughput-tests-arm64-cpu.json index 3863cccf43..3bc7dde979 100644 --- a/.buildkite/performance-benchmarks/tests/throughput-tests-arm64-cpu.json +++ b/.buildkite/performance-benchmarks/tests/throughput-tests-arm64-cpu.json @@ -3,7 +3,6 @@ "test_name": "throughput_llama8B_tp1", "environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_KVCACHE_SPACE": 40 }, "parameters": { diff --git a/.buildkite/performance-benchmarks/tests/throughput-tests-cpu.json b/.buildkite/performance-benchmarks/tests/throughput-tests-cpu.json index d3f16eff11..089be00eb4 100644 --- a/.buildkite/performance-benchmarks/tests/throughput-tests-cpu.json +++ b/.buildkite/performance-benchmarks/tests/throughput-tests-cpu.json @@ -3,7 +3,6 @@ "test_name": "throughput_llama8B_tp2", "environment_variables": { "VLLM_ALLOW_LONG_MAX_MODEL_LEN": 1, - "VLLM_ENGINE_ITERATION_TIMEOUT_S": 120, "VLLM_CPU_SGL_KERNEL": 1, "VLLM_CPU_KVCACHE_SPACE": 40 }, diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 2b90fef820..d6e47715fa 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -339,7 +339,8 @@ steps: --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg RUN_WHEEL_CHECK=false \ --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.4 \ - --build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \ + --build-arg FINAL_BASE_IMAGE=nvidia/cuda:13.4.1-base-ubuntu24.04 \ + --build-arg NCCL_VERSION=2.32.3 \ --target vllm-openai \ --progress plain \ -f docker/Dockerfile . @@ -369,7 +370,8 @@ steps: --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg RUN_WHEEL_CHECK=false \ --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.4 \ - --build-arg FINAL_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04 \ + --build-arg FINAL_BASE_IMAGE=nvidia/cuda:13.4.1-base-ubuntu24.04 \ + --build-arg NCCL_VERSION=2.32.3 \ --target vllm-openai \ --progress plain \ -f docker/Dockerfile . diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index d07f3b3830..59240b3d0a 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -1712,8 +1712,7 @@ if is_multi_node "$commands"; then # BASH_REMATCH[2] = comma-separated node0 commands # BASH_REMATCH[3] = comma-separated node1 commands if [[ "$commands" =~ ^(.*)\[(.*)"] && ["(.*)\]$ ]]; then - # shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR - prefix=$(echo "${BASH_REMATCH[1]}" | sed 's/;//g') + prefix=${BASH_REMATCH[1]//;/} echo "PREFIX: ${prefix}" export composite_command="(command rocm-smi || true)" @@ -1728,10 +1727,8 @@ if is_multi_node "$commands"; then fi for i in "${!node0[@]}"; do - # shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR - command_node_0=$(echo "${node0[i]}" | sed 's/\"//g') - # shellcheck disable=SC2001 # verified equivalent behavior; TODO: switch to param expansion in a follow-up cleanup PR - command_node_1=$(echo "${node1[i]}" | sed 's/\"//g') + command_node_0=${node0[i]//\"/} + command_node_1=${node1[i]//\"/} step_cmd="./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 ${image_name} '${command_node_0}' '${command_node_1}'" echo "COMMANDS: ${step_cmd}" diff --git a/.buildkite/scripts/initialized-snapshot-e2e.sh b/.buildkite/scripts/initialized-snapshot-e2e.sh index cd0511ffdd..d989906443 100755 --- a/.buildkite/scripts/initialized-snapshot-e2e.sh +++ b/.buildkite/scripts/initialized-snapshot-e2e.sh @@ -34,6 +34,7 @@ readonly DEADLINE_EPOCH="$(( $(date +%s) + RUN_LIMIT_S ))" RUN_ID="" CONTAINER_NAME="" CONTAINER_ID="" GPU_UUID="" PORT_ONE="" PORT_TWO="" LINK_REMAP_BASELINE="" +CHECK_CLEANUP_STATE=0 SNAPSHOT_PIDS=() die() { echo "initialized snapshot E2E: $*" >&2 @@ -113,7 +114,7 @@ cleanup() { names="$(container_names 2>/dev/null)" || return 1 grep -Fxq "$CONTAINER_NAME" <<< "$names" && failed=1 fi - if [[ -n "$PORT_TWO" ]] && ! wait_clean; then + if (( CHECK_CLEANUP_STATE )) && ! wait_clean; then echo "process, GPU, port, or link-remap residue remained" >&2 failed=1 fi @@ -218,7 +219,8 @@ GPU_PIDS="$(gpu_pids)" || die "failed to inspect the selected GPU" RUN_ID="snapshot-e2e-$(date +%s)-$$-$RANDOM" CONTAINER_NAME="vllm-$RUN_ID" read -r PORT_ONE PORT_TWO < <(python3 -c 'import socket; s=[socket.socket(),socket.socket()]; [x.bind(("127.0.0.1",0)) for x in s]; print(*(x.getsockname()[1] for x in s))') -LINK_REMAP_BASELINE="$(link_remaps)" +LINK_REMAP_BASELINE="$(link_remaps)" \ + || die "cannot inventory /dev/shm; check CI runner permissions" run "pull exact candidate image" 900 docker pull --platform linux/amd64 \ "$IMAGE_TAG" >/dev/null @@ -231,6 +233,8 @@ IFS='|' read -r IMAGE_ID IMAGE_PLATFORM OCI_COMMIT VLLM_COMMIT < <( || die "candidate image identity is invalid" [[ "$OCI_COMMIT" == "$EXPECTED_COMMIT" && "$VLLM_COMMIT" == "$EXPECTED_COMMIT" ]] \ || die "candidate image commit labels do not match" +# A failed preflight has not started any workload whose cleanup we can verify. +CHECK_CLEANUP_STATE=1 CONTAINER_ID="$(run "start exact candidate container" 90 docker run --detach \ --name "$CONTAINER_NAME" --label "ai.vllm.snapshot.e2e.run=$RUN_ID" \ --gpus "device=$GPU_UUID" --user 0 --privileged --pid=host --ipc=host \ diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index acd4d94a53..0a050b6979 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2398,6 +2398,7 @@ steps: - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - tests/quantization - tests/rocm/test_moe_weight_replay.py + - tests/rocm/test_moe_padding.py optional: true commands: # temporary install here since we need nightly, will move to requirements/test.in @@ -2411,6 +2412,7 @@ steps: - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi - label: ":amd: (MI300) Quantized Fusions" # TBD timeout_in_minutes: 45 diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 5fd0af69c1..acc1bafa7f 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -93,6 +93,24 @@ steps: commands: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -s -v tests/compile/passes/distributed + mirror: + amd: + label: ':amd: (MI300) Distributed Compile' + dind: false + device: mi300_2 + timeout_in_minutes: 65 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers + - tests/compile/passes/distributed/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - pytest -s -v tests/compile/passes/distributed - label: ":nvidia: (B200) Fusion and Compile" key: fusion-and-compile-unit-tests-2xb200 @@ -463,3 +481,45 @@ steps: - amd-smi - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)' - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py + +- label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Inductor Partition)' + key: fusion-e2e-tp2-ar-rms-inductor-partition-amd + dind: false + device: mi355_2 + num_devices: 2 + timeout_in_minutes: 65 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers + - tests/compile/fusions_e2e/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - csrc/quantization/ + - vllm/v1/attention/ + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "inductor_partition" + +- label: ':amd: (MI355) Fusion E2E TP2 AR-RMS (Dynamo Partition)' + key: fusion-e2e-tp2-ar-rms-dynamo-partition-amd + dind: false + device: mi355_2 + num_devices: 2 + timeout_in_minutes: 65 + working_dir: /vllm-workspace/ + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/layers + - tests/compile/fusions_e2e/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - csrc/quantization/ + - vllm/v1/attention/ + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY=1 + - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions -k "dynamo_partition" diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 7138e1e065..3710f80bfd 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -17,6 +17,24 @@ steps: - pytest -v -s distributed/test_shm_broadcast.py - pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_storage.py + mirror: + amd: + label: ':amd: (MI300) Distributed Comm Ops' + dind: false + device: mi300_2 + timeout_in_minutes: 30 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed + - tests/distributed + - vllm/platforms/rocm.py + commands: + - pytest -v -s distributed/test_comm_ops.py + - pytest -v -s distributed/test_shm_broadcast.py + - pytest -v -s distributed/test_shm_buffer.py + - pytest -v -s distributed/test_shm_storage.py - label: ":nvidia: (L4) Sharded RDT Weight Transfer" key: sharded-rdt-weight-transfer @@ -176,6 +194,31 @@ steps: - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - pytest -v -s v1/worker/test_worker_memory_snapshot.py + mirror: + amd: + label: ':amd: (MI300) Distributed Torchrun + Shutdown' + dind: false + device: mi300_2 + timeout_in_minutes: 35 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ + - vllm/engine/ + - vllm/v1/executor/ + - vllm/v1/worker/worker_base.py + - vllm/v1/engine/ + - vllm/v1/worker/ + - tests/distributed/ + - tests/v1/shutdown + - tests/v1/worker/test_worker_memory_snapshot.py + - vllm/platforms/rocm.py + commands: + - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' + - HIP_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown + - pytest -v -s v1/worker/test_worker_memory_snapshot.py - label: ":nvidia: (L4) Distributed Torchrun + Examples" key: distributed-torchrun-examples-4-gpus @@ -282,6 +325,31 @@ steps: - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py - pytest -v -s distributed/test_engram_dp_shard.py + mirror: + amd: + label: ':amd: (MI355) Distributed DP Extended' + dind: false + device: mi355_4 + timeout_in_minutes: 60 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ + - tests/v1/distributed + - tests/v1/engine/test_engine_core_client.py + - tests/distributed/test_utils + - vllm/platforms/rocm.py + commands: + - export NCCL_CUMEM_HOST_ENABLE=0 + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py + - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_dense_dp_world_size.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py + - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp + - pytest -v -s distributed/test_utils.py - label: ":nvidia: (L4) Distributed Compile + Comm" key: distributed-compile-comm-4-gpus @@ -304,6 +372,27 @@ steps: - pytest -v -s distributed/test_symm_mem_allreduce.py # test multi-node TP with multiproc executor (simulated on single node) - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node + mirror: + amd: + label: ':amd: (MI355) Distributed Compile + Comm' + dind: false + device: mi355_4 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ + - tests/distributed/test_pynccl + - tests/distributed/test_events + - tests/distributed/test_symm_mem_allreduce.py + - tests/distributed/test_multiproc_executor.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s distributed/test_pynccl.py + - pytest -v -s distributed/test_events.py + - pytest -v -s distributed/test_symm_mem_allreduce.py + - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node - label: ":nvidia: (H100) Distributed DP + EP" key: distributed-tests-8xh100 @@ -359,6 +448,28 @@ steps: - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py + mirror: + amd: + label: ':amd: (MI355) Distributed' + dind: false + device: mi355_4 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/ + - tests/models/test_vision.py + - vllm/platforms/rocm.py + - tests/distributed/ + - tests/basic_correctness/ + - tests/lora/test_mixtral.py + commands: + - pytest -v -s distributed/test_custom_all_reduce.py + - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py + - TARGET_TEST_SUITE=MI355 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - pytest -v -s -x lora/test_mixtral.py + - pytest -v -s models/test_vision.py -m 'distributed(num_gpus=4)' - label: ":nvidia: (H100) Distributed Features" key: distributed-tests-2xh100-2xmi300 @@ -583,3 +694,27 @@ steps: - pytest -v -s distributed/test_ray_v2_executor_e2e.py - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray" + mirror: + amd: + label: ':amd: (MI355) RayExecutorV2' + dind: false + device: mi355_4 + timeout_in_minutes: 75 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/executor/ray_executor_v2.py + - vllm/v1/executor/abstract.py + - vllm/v1/executor/multiproc_executor.py + - tests/distributed/test_ray_v2_executor.py + - tests/distributed/test_ray_v2_executor_e2e.py + - tests/distributed/test_pipeline_parallel.py + - tests/basic_correctness/models/test_basic_correctness.py + - vllm/platforms/rocm.py + commands: + - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 + - pytest -v -s distributed/test_ray_v2_executor.py + - pytest -v -s distributed/test_ray_v2_executor_e2e.py + - pytest -v -s distributed/test_pipeline_parallel.py -k "ray" + - TARGET_TEST_SUITE=MI355 pytest -v -s basic_correctness/models/test_basic_correctness.py -k "ray" diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index 9af5423a4b..b211c47e35 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -150,3 +150,27 @@ steps: - .buildkite/scripts/scheduled_integration_test/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py + +- label: ":amd: (MI355) Qwen3-Next-80B-A3B-Instruct MTP Async EPLB Accuracy" + key: amd-qwen3-next-mtp-async-eplb-accuracy + dind: false + device: mi355_4 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + working_dir: /vllm-workspace + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/spec_decode/ + - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - .buildkite/scripts/scheduled_integration_test/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + num_devices: 4 + commands: + - VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 930a9a448a..76f2835978 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -10,9 +10,11 @@ steps: source_file_dependencies: - vllm/distributed/eplb - tests/distributed/test_eplb_algo.py + - tests/distributed/test_eplb_migration_scheduler.py - tests/distributed/test_eplb_utils.py commands: - pytest -v -s distributed/test_eplb_algo.py + - pytest -v -s distributed/test_eplb_migration_scheduler.py - pytest -v -s distributed/test_eplb_utils.py mirror: amd: diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index a23fcd0a98..04337d03d2 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -753,6 +753,32 @@ steps: parallelism: 4 commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200-shard-$$BUILDKITE_PARALLEL_JOB.txt + mirror: + amd: + label: ":amd: (MI355) MoE Refactor BF16 Integration" + dind: false + device: mi355_2 + num_devices: 2 + parallelism: 1 + timeout_in_minutes: 105 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - tests/evals/gsm8k/ + - tests/utils.py + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/kernels/linear/ + - vllm/v1/attention/ + - vllm/distributed/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt -k "Llama-4-Scout-BF16-triton or Mixtral-8x7B-BF16-triton" - label: ":nvidia: (B200) MoE Refactor DP Integration TEMPORARY" key: moe-refactor-integration-test-b200-dp-temporary @@ -761,6 +787,32 @@ steps: num_devices: 2 commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt + mirror: + amd: + label: ":amd: (MI355) MoE Refactor BF16 DP Integration" + dind: false + device: mi355_2 + num_devices: 2 + timeout_in_minutes: 105 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - tests/evals/gsm8k/ + - tests/utils.py + - vllm/config/parallel.py + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/kernels/linear/ + - vllm/v1/attention/ + - vllm/distributed/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + env: + VLLM_WORKER_MULTIPROC_METHOD: "spawn" + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt -k "Qwen3-30B-A3B-BF16-triton" - label: ":nvidia: (L4) LM Eval Humming TEMPORARY" key: lm-eval-humming-l4 @@ -1068,6 +1120,29 @@ steps: - tests/evals/gpt_oss/ commands: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt + mirror: + amd: + label: ":amd: (MI300) GPQA Eval (GPT-OSS)" + dind: false + device: mi300_2 + num_devices: 2 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - tests/evals/gpt_oss/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-gfx942.txt - label: ":nvidia: (B200) GPQA Eval (GPT-OSS)" key: gpqa-eval-gpt-oss-2xb200 @@ -1176,6 +1251,28 @@ steps: - tests/evals/gpt_oss/ commands: - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt + mirror: + amd: + label: ":amd: (MI355 DPX) GPQA Eval (GPT-OSS TP1)" + # The Spark config is a backend-neutral GPT-OSS TP1 baseline. + dind: false + device: mi355_dpx + num_devices: 1 + timeout_in_minutes: 40 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/fused_moe/ + - tests/evals/gpt_oss/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 35GB) LM Eval KV-Offload" key: kv-offload-small @@ -1292,3 +1389,53 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py - tests/evals/mrcr/ + +- label: ":amd: (MI355) LM Eval Large Models Harness" + key: amd-lm-eval-large-models-harness + dind: false + device: mi355_8 + optional: true + timeout_in_minutes: 85 + depends_on: + - image-build-amd + working_dir: /vllm-workspace/.buildkite/lm-eval-harness + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - .buildkite/lm-eval-harness/ + num_devices: 8 + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # Qwen2-57B has 28 attention heads and cannot run at TP8. + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 + +- label: ":amd: (MI300) LM Eval Large Models ROCm Harness" + key: amd-lm-eval-large-models-rocm-harness + dind: false + device: mi300_8 + optional: true + timeout_in_minutes: 40 + depends_on: + - image-build-amd + working_dir: /vllm-workspace/.buildkite/lm-eval-harness + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/layernorm.py + - csrc/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + - .buildkite/lm-eval-harness/ + num_devices: 8 + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index dea2ca7a6f..602491e0e0 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -560,6 +560,8 @@ steps: - tools/pre_commit/check_test_tethering.py - tools/pre_commit/test_tethering_allowlist.txt - tests/tools/test_check_test_tethering.py + - tools/pre_commit/check_agent_files.py + - tests/tools/test_check_agent_files.py - tests/tool_parsers - tests/tokenizers_ - tests/parser @@ -609,6 +611,7 @@ steps: - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s config - pytest -v -s tools/test_check_test_tethering.py + - pytest -v -s tools/test_check_agent_files.py - label: ":computer: (CPU) Reasoning + Renderers" key: cpu-reasoning-renderers diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 735735c15e..c2b2b054a3 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -124,10 +124,9 @@ steps: - csrc/libtorch_stable/kimi_k3/ - tests/models/kimi_k3/ - tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py - - tests/kernels/test_bf16_skinny_gemm.py commands: # The native NVIDIA Kimi K3 kernels require the SM100 family. - - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py + - pytest -v -s models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py kernels/attention/test_kimi_k3_mla_fused_epilogue.py mirror: amd: label: ":amd: (MI355) Kimi K3" @@ -142,10 +141,30 @@ steps: - vllm/platforms/rocm.py commands: # The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355. - # The two kernels/ files the NVIDIA job adds are CUDA-only (they import - # cuda.bindings.driver), so they are deliberately left out here. + # The kernels/ file the NVIDIA job adds is CUDA-only (it imports + # cuda.bindings.driver), so it is deliberately left out here. - pytest -v -s models/kimi_k3 +- label: ":nvidia: (4xB200) Kimi K3 Prefix Cache Shard %N" + key: kimi-k3-prefix-cache-4xb200 + timeout_in_minutes: 75 + device: b200-k8s + num_devices: 4 + source_file_dependencies: + - vllm/models/kimi_k3/ + - vllm/v1/core/ + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/simple_kv_offload/ + - vllm/v1/spec_decode/ + - vllm/model_executor/layers/mamba/ + - tests/models/kimi_k3/test_prefix_cache.py + - .buildkite/scripts/install-kv-connectors.sh + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + parallelism: 2 + - label: ":nvidia: (H200) GLM5Next" key: glm5next-unit-tests timeout_in_minutes: 20 diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index 5b6260b534..0ba608e725 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -259,6 +259,25 @@ steps: - tests/compile commands: - pytest -s -v compile/fullgraph/test_full_cudagraph.py + mirror: + amd: + label: ':amd: (MI250) PyTorch Fullgraph CUDAGraph Compatibility' + dind: false + device: mi250_1 + timeout_in_minutes: 65 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/compilation/ + - vllm/model_executor/ + - vllm/v1/attention/ + - vllm/config/compilation.py + - csrc/ + - tests/compile + - vllm/platforms/rocm.py + commands: + - pytest -s -v compile/fullgraph/test_full_cudagraph.py - label: ":nvidia: (H200 MIG 18GB) PyTorch Nightly Dependency Override Check" key: pytorch-nightly-dependency-override-check diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 78ea92d232..e73054a4bc 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -38,6 +38,7 @@ steps: - vllm/model_executor/layers/fused_moe/oracle/unquantized.py - vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py - tests/rocm/test_moe_weight_replay.py + - tests/rocm/test_moe_padding.py commands: # Use the ROCm default instead of the NVIDIA runner override. - unset VLLM_USE_V2_MODEL_RUNNER @@ -45,6 +46,7 @@ steps: - uv pip install --system conch-triton-kernels - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi + - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 pytest -v -s rocm/test_moe_padding.py; fi - label: ":nvidia: (H200 MIG 35GB) Quantized Fusions" device: h200_35gb diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index af7e8f071a..6882c24b37 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -97,6 +97,27 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/eagle/ + mirror: + amd: + label: ":amd: (MI355) Spec Decode Eagle Nightly" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 120 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/v1/attention/ + - tests/v1/e2e/spec_decode/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py # MTP correctness is one parametrized test split into per-model folders; # each job below runs whole directories. @@ -234,6 +255,28 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/speculators/ - pytest -v -s v1/e2e/spec_decode/mtp/ + mirror: + amd: + label: ":amd: (MI355) Spec Decode Speculators + MTP Nightly" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 90 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/v1/attention/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H200 MIG 35GB) Spec Decode N-Gram + Suffix" key: spec-decode-ngram-suffix @@ -304,6 +347,27 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ + mirror: + amd: + label: ":amd: (MI355) Spec Decode Draft Model Nightly" + dind: false + device: mi355_1 + num_devices: 1 + timeout_in_minutes: 120 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/v1/attention/ + - tests/v1/e2e/spec_decode/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py - label: ":nvidia: (H100) Speculators Correctness Nightly" key: speculators-correctness diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 24dcf826fb..3e418e005e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -30,7 +30,6 @@ repos: - id: markdownlint-cli2 language_version: lts args: [--config, tools/pre_commit/.markdownlint.yaml, --fix] - exclude: (^|/)CLAUDE\.md$ - repo: https://github.com/rhysd/actionlint rev: v1.7.7 hooks: @@ -177,6 +176,14 @@ repos: files: ^\.github/mergify\.yml$ pass_filenames: false additional_dependencies: [PyYAML, regex] + - id: check-agent-files + name: Check agent instruction files and skills + language: python + entry: python tools/pre_commit/check_agent_files.py + files: ^(\.agents/|\.claude/|(.*/)?(AGENTS|CLAUDE)\.md$|tools/pre_commit/check_agent_files\.py$) + pass_filenames: false + additional_dependencies: [PyYAML, regex] + verbose: true # `note:` lines are advisory, so the hook passes - id: check-buildkite-step-keys name: Check every Buildkite step declares an explicit key language: python diff --git a/benchmarks/kernels/benchmark_top_k_per_row_decode.py b/benchmarks/kernels/benchmark_top_k_per_row_decode.py new file mode 100644 index 0000000000..11337ca110 --- /dev/null +++ b/benchmarks/kernels/benchmark_top_k_per_row_decode.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Benchmark the ROCm sparse-indexer decode top-k in isolation. + +Example usage: +python3 benchmarks/kernels/benchmark_top_k_per_row_decode.py \ + --rows 4 16 32 --lens 2048 8192 16384 65536 --capacity 262144 +""" + +import argparse +import itertools +import json + +import torch + +import vllm._custom_ops # noqa: F401 +from vllm._aiter_ops import rocm_aiter_ops +from vllm.triton_utils import triton + +DISTRIBUTIONS = ("indexer", "fp8", "randn") + + +def make_logits(rows: int, capacity: int, dist: str, spread: float): + normal = torch.randn(rows, capacity, device="cuda", dtype=torch.float32) + if dist == "randn": + return normal + if dist == "fp8": + return normal.to(torch.float8_e4m3fnuz).to(torch.float32) + return normal * spread + 3.0 + + +def make_inputs(rows: int, live: int, capacity: int, k: int, args): + torch.manual_seed(args.seed) + logits = make_logits(rows, capacity, args.dist, args.spread) + seq_lens = torch.full((rows,), live, device="cuda", dtype=torch.int32) + invalid = torch.arange(capacity, device="cuda")[None] >= seq_lens[:, None] + logits.masked_fill_(invalid, -float("inf")) + indices = torch.empty((rows, k), device="cuda", dtype=torch.int32) + return logits, seq_lens, indices + + +def run_native(logits, seq_lens, indices, k): + torch.ops._C.top_k_per_row_decode( + logits, + 1, + seq_lens, + indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + k, + ) + + +def run_aiter(logits, seq_lens, indices, k): + rocm_aiter_ops.indexer_top_k_decode(logits, 1, seq_lens, indices, k) + + +BACKENDS = {"native": run_native, "aiter": run_aiter} + + +def check(logits, indices, k): + selected = logits.gather(1, indices.long()).sort(dim=1, descending=True).values + expected = logits.topk(k, dim=1).values + torch.testing.assert_close(selected, expected, atol=0, rtol=0) + + +def capture(run): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + return graph + + +def benchmark(backend: str, rows: int, live: int, args) -> dict: + logits, seq_lens, indices = make_inputs(rows, live, args.capacity, args.top_k, args) + fn = BACKENDS[backend] + + def run(): + fn(logits, seq_lens, indices, args.top_k) + + run() + torch.accelerator.synchronize() + check(logits, indices, args.top_k) + timed = run + if args.graph: + graph = capture(run) + indices.zero_() + graph.replay() + torch.accelerator.synchronize() + check(logits, indices, args.top_k) + timed = graph.replay + ms = triton.testing.do_bench(timed, warmup=args.warmup, rep=args.rep) + return dict( + backend=backend, + mode="graph" if args.graph else "eager", + dist=args.dist, + rows=rows, + live=live, + capacity=args.capacity, + k=args.top_k, + us=ms * 1000, + live_gbps=rows * live * 4 / (ms * 1e6), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 2, 4, 8]) + parser.add_argument( + "--lens", + type=int, + nargs="+", + default=[32768, 65536, 98304, 131072], + help="Valid columns per row; must be at least --top-k", + ) + parser.add_argument("--capacity", type=int, default=262144) + parser.add_argument("--top-k", type=int, default=512) + parser.add_argument( + "--backends", nargs="+", choices=list(BACKENDS), default=["native", "aiter"] + ) + parser.add_argument( + "--graph", + action="store_true", + help="Time a captured graph replay, matching how decode actually runs", + ) + parser.add_argument( + "--dist", + choices=DISTRIBUTIONS, + default="indexer", + help="Logit value distribution; see the module docstring", + ) + parser.add_argument( + "--spread", + type=float, + default=0.01, + help="Relative width of the --dist indexer cluster", + ) + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--seed", type=int, default=123) + args = parser.parse_args() + + print(json.dumps(dict(gpu=torch.cuda.get_device_name(0), args=vars(args)))) + for backend, rows, live in itertools.product(args.backends, args.rows, args.lens): + if live < args.top_k: + continue + print(json.dumps(benchmark(backend, rows, live, args)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_w8a8_block_fp8.py b/benchmarks/kernels/benchmark_w8a8_block_fp8.py index 3ad7ec8887..f543ee520b 100644 --- a/benchmarks/kernels/benchmark_w8a8_block_fp8.py +++ b/benchmarks/kernels/benchmark_w8a8_block_fp8.py @@ -23,9 +23,11 @@ from vllm.utils.platform_utils import get_device_name_as_file_name mp.set_start_method("spawn", force=True) -assert current_platform.is_cuda() or current_platform.is_rocm(), ( - "Only support tune w8a8 block fp8 kernel on CUDA/ROCm device." -) +assert ( + current_platform.is_cuda() + or current_platform.is_rocm() + or current_platform.is_xpu() +), "Only support tune w8a8 block fp8 kernel on CUDA/ROCm/XPU device." DTYPE_MAP = { "float32": torch.float32, @@ -138,6 +140,29 @@ def get_configs_compute_bound(): return configs +def get_xpu_tuning_space(): + # NOTE: curated subset of get_configs_compute_bound(); the full grid is + # too slow to tune on XPU due to per-config Triton JIT compile time. + configs = [] + for num_stages in [2, 3]: + for block_m in [32, 64, 128]: + for block_k in [64, 128]: + for block_n in [128, 256]: + for num_warps in [4, 8, 16, 32]: + for group_size in [1, 32]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + } + ) + return configs + + def get_weight_shapes(tp_size): # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. # Modify them, if you tune for another different model. @@ -208,13 +233,14 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type): fp8_info = torch.finfo(torch.float8_e4m3fn) fp8_max, fp8_min = fp8_info.max, fp8_info.min + device = current_platform.device_type A_fp32 = ( - (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max ) A = A_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) B_fp32 = ( - (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + (torch.rand(N, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max ) B = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) else: @@ -224,9 +250,9 @@ def tune(M, N, K, block_size, out_dtype, search_space, input_type): 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="cuda") * factor_for_scale + 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="cuda") + torch.rand(n_tiles, k_tiles, dtype=torch.float32, device=device) * factor_for_scale ) @@ -297,7 +323,11 @@ def tune_on_gpu(args_dict): save_path = args.save_path input_type = args.input_type - search_space = get_configs_compute_bound() + search_space = ( + get_xpu_tuning_space() + if current_platform.is_xpu() + else get_configs_compute_bound() + ) search_space = [ config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 ] @@ -342,7 +372,10 @@ def main(args): raise RuntimeError("No GPU available for tuning") print(f"Found {num_gpus} GPUs for parallel tuning") - torch.cuda.init() + if current_platform.is_xpu(): + torch.xpu.init() + else: + torch.cuda.init() if args.batch_size is None: batch_sizes = [ diff --git a/docker/Dockerfile b/docker/Dockerfile index 7dfcd61b9e..d8a1258a74 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -251,33 +251,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi -# CUDA 13.4 Rubin images need a compatibility link for the prerelease -# PyTorch wheel's libcurand.so.11 dependency. CUDA 13.5 must use its native -# libcurand.so.11 instead. -RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \ - && if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \ - CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \ - -print -quit) \ - && test -n "${CURAND_LIBRARY}" \ - && ln -sf "${CURAND_LIBRARY}" \ - /usr/local/cuda/lib64/libcurand.so.11; \ - elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \ - CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \ - -print -quit) \ - && test -n "${CURAND_LIBRARY}" \ - && ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \ - else \ - echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \ - exit 1; \ - fi \ - && printf '%s\n' \ - /opt/venv/lib \ - "/opt/venv/lib/python${PYTHON_VERSION}/site-packages/nvpl/lib" \ - > /etc/ld.so.conf.d/vllm-rubin-python.conf \ - && ldconfig; \ - fi - # Track PyTorch lib versions used during build and match in downstream instances. # We do this for both nightly and release so we can strip dependencies/*.txt as needed. # Otherwise library dependencies can upgrade/downgrade torch incorrectly. @@ -518,20 +491,6 @@ ARG NVSHMEM_VER ARG NCCL_VERSION RUN --mount=type=cache,target=/opt/uv/cache \ mkdir -p /tmp/ep_kernels_workspace/dist && \ - if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - # The Rubin CUDA builder provides NCCL in its architecture-specific \ - # toolkit directory rather than as a Python package. Give DeepEP a \ - # stable prefix that also exists in the Ubuntu runtime image. \ - NCCL_HEADER=$(find /usr/local/cuda/targets -name nccl.h -print -quit) \ - && test -n "${NCCL_HEADER}" \ - && NCCL_LIBRARY=$(find -H /usr/local/cuda/targets \ - -name libnccl.so.2 -print -quit) \ - && test -n "${NCCL_LIBRARY}" \ - && ln -sf "${NCCL_HEADER}" /usr/include/nccl.h \ - && ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2 \ - && ln -sf libnccl.so.2 /usr/lib/libnccl.so \ - && export EP_NCCL_ROOT_DIR=/usr; \ - fi && \ CUDA_MAJOR=$(echo $CUDA_VERSION | cut -d. -f1) && \ if [ -n "$NCCL_VERSION" ]; then \ echo "nvidia-nccl-cu${CUDA_MAJOR}==${NCCL_VERSION}" \ @@ -659,10 +618,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \ if [ "${PYTORCH_NIGHTLY}" = "1" ] || \ [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ python3 tools/use_existing_torch.py --prefix; \ - fi \ - && if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - sed -Ei '/^[[:space:]]*(flashinfer-python|nvidia-cutlass-dsl\[cu13\]|quack-kernels)[[:space:]]*==/Id' \ - requirements/cuda.txt; \ fi # Build the vLLM wheel @@ -764,11 +719,6 @@ ARG DEADSNAKES_MIRROR_URL ARG DEADSNAKES_GPGKEY_URL ARG GET_PIP_URL -# Keep the output image's CUDA metadata aligned with the selected toolkit. -# This also normalizes prerelease base images with four-component metadata for -# OCI runtimes that evaluate NVIDIA's CUDA driver requirement. -ENV CUDA_VERSION=${CUDA_VERSION} - ENV DEBIAN_FRONTEND=noninteractive WORKDIR /vllm-workspace @@ -841,41 +791,6 @@ ENV CUDA_HOME=/usr/local/cuda ENV PATH="${CUDA_HOME}/bin:${PATH}" ENV LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}" -# Apply the CUDA 13.4-only libcurand compatibility link described in the build -# stage. CUDA 13.5 uses its native libcurand.so.11. DeepEP expects NCCL directly -# under its configured /usr prefix, while Ubuntu installs it in a multiarch -# subdirectory - but that symlink is only possible when NCCL is already -# present (e.g. nvcr.io/nvidia/cuda-dl-base bundles it; a bare-CUDA image -# doesn't, and gets it later via the standard nvidia-nccl-cu* pip package, -# UV_OVERRIDE below). Check reality (whether libnccl.so.2 is actually found) -# and tolerate not finding it instead of failing the build. -RUN if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ]; then \ - CUDA_MAJOR_MINOR=$(echo "${CUDA_VERSION}" | cut -d. -f1,2) \ - && if [ "${CUDA_MAJOR_MINOR}" = "13.4" ]; then \ - CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.10 \ - -print -quit) \ - && test -n "${CURAND_LIBRARY}" \ - && ln -sf "${CURAND_LIBRARY}" \ - /usr/local/cuda/lib64/libcurand.so.11; \ - elif [ "${CUDA_MAJOR_MINOR}" = "13.5" ]; then \ - CURAND_LIBRARY=$(find -H /usr/local/cuda -name libcurand.so.11 \ - -print -quit) \ - && test -n "${CURAND_LIBRARY}" \ - && ! readlink -f "${CURAND_LIBRARY}" | grep -q 'libcurand\.so\.10'; \ - else \ - echo "Unsupported Rubin CUDA version: ${CUDA_MAJOR_MINOR}" >&2; \ - exit 1; \ - fi \ - && if NCCL_LIBRARY=$(find /usr/lib -name libnccl.so.2 -print -quit) \ - && [ -n "${NCCL_LIBRARY}" ]; then \ - if [ "${NCCL_LIBRARY}" != "/usr/lib/libnccl.so.2" ]; then \ - ln -sf "${NCCL_LIBRARY}" /usr/lib/libnccl.so.2; \ - fi \ - && ln -sf libnccl.so.2 /usr/lib/libnccl.so; \ - fi \ - && ldconfig; \ - fi - # Install the initialized-snapshot runtime on the architecture supported by # `vllm snapshot`. CRIU and its CUDA plugin are built from the same pinned # source. The `cuda-checkpoint`-compatible helper is built from NVIDIA Dynamo's @@ -990,11 +905,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \ && uv pip install --system -r /tmp/requirements-cuda.txt \ && uv pip install --system \ --prerelease=allow --index-strategy unsafe-first-match \ - -r /tmp/vllm-requirements/rubin-prerelease.txt \ - && printf '%s\n' \ - "/usr/local/lib/python${PYTHON_VERSION}/dist-packages/nvpl/lib" \ - > /etc/ld.so.conf.d/vllm-rubin-python.conf \ - && ldconfig; \ + -r /tmp/vllm-requirements/rubin-prerelease.txt; \ else \ uv pip install --system \ -r /tmp/requirements-cuda.txt \ @@ -1007,10 +918,8 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # From versions.json: .flashinfer.version ARG FLASHINFER_VERSION=0.7.0.post1 RUN --mount=type=cache,target=/opt/uv/cache \ - if [ "${INSTALL_RUBIN_PRERELEASE}" != "true" ]; then \ - uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ - --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ - fi + uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ + --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') # ============================================================ # OPENAI API SERVER DEPENDENCIES @@ -1023,22 +932,13 @@ ARG GDRCOPY_CUDA_VERSION=12.8 ARG GDRCOPY_OS_VERSION=Ubuntu24_04 ARG TARGETPLATFORM COPY tools/install_gdrcopy.sh /tmp/install_gdrcopy.sh -# Some Rubin FINAL_BASE_IMAGE choices (internal Rubin CUDA images, -# nvcr.io/nvidia/cuda-dl-base) already provide GDRCopy; others don't, so it -# falls through to the normal install below. Check reality (ldconfig) rather -# than inferring from which FINAL_BASE_IMAGE flag is set. RUN set -eux; \ - if [ "${INSTALL_RUBIN_PRERELEASE}" = "true" ] && ldconfig -p | grep -q 'libgdrapi\.so'; then \ - true; \ - else \ - case "${TARGETPLATFORM}" in \ - linux/arm64) UUARCH="aarch64" ;; \ - linux/amd64) UUARCH="x64" ;; \ - *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \ - esac; \ - /tmp/install_gdrcopy.sh \ - "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}"; \ - fi; \ + case "${TARGETPLATFORM}" in \ + linux/arm64) UUARCH="aarch64" ;; \ + linux/amd64) UUARCH="x64" ;; \ + *) echo "Unsupported TARGETPLATFORM: ${TARGETPLATFORM}" >&2; exit 1 ;; \ + esac; \ + /tmp/install_gdrcopy.sh "${GDRCOPY_OS_VERSION}" "${GDRCOPY_CUDA_VERSION}" "${UUARCH}" && \ rm /tmp/install_gdrcopy.sh # Install vllm-openai dependencies (saves ~2.6s per build) diff --git a/docs/community/reviewers.md b/docs/community/reviewers.md new file mode 100644 index 0000000000..fe85e1af49 --- /dev/null +++ b/docs/community/reviewers.md @@ -0,0 +1,55 @@ +# Reviewers + +Reviewers are trusted community members who can perform official reviews, helping contributors get their pull requests ready to merge alongside committers, who give final approval. Authors are expected to respond to reviewer feedback, either by making changes or by explaining why a change isn't needed. The role was introduced in [RFC #58253](https://github.com/vllm-project/vllm/issues/58253). + +This page lists the active reviewers and describes the role. See the [committers](../governance/committers.md) page for the committers and area owners who hold final acceptance and merge authority. + +## Active Reviewers + +Community members can reach out to the reviewers below for help with PRs in their areas, for example in the `#pr-reviews` Slack channel. Reviewers will add their areas of expertise to this list. +Sorted alphabetically by GitHub handle: + +- [@andylolu2](https://github.com/andylolu2) +- [@bnellnm](https://github.com/bnellnm) MoE, Quantization layers +- [@BowenBao](https://github.com/BowenBao) +- [@cjackal](https://github.com/cjackal): Multimodality; KV Connector and offload +- [@divakar-amd](https://github.com/divakar-amd) +- [@etelis](https://github.com/etelis): KV cache offloading (CPU, filesystem, and P2P connectors; KV transfer and cache layouts); MoE serving and expert parallelism (developing review focus) +- [@Fangzhou-Ai](https://github.com/Fangzhou-Ai) +- [@fxmarty-amd](https://github.com/fxmarty-amd): Quantization, linear/MOE oracles and modeling definition +- [@GirasoleY](https://github.com/GirasoleY): Kernels and performance +- [@itayalroy](https://github.com/itayalroy): MoE serving (kernels, EP, All2All, EPLB, elasticity/fault tolerance); KV connectors and NIXL integrations +- [@JartX](https://github.com/JartX): ROCm / HIP on Radeon (RDNA3), quantization (W4A16/GPTQ, MXFP4), KV cache quantization +- [@kliuae](https://github.com/kliuae): ROCm feature enablement, model performance, fused kernels +- [@lengrongfu](https://github.com/lengrongfu) +- [@linitra24](https://github.com/linitra24) +- [@LopezCastroRoberto](https://github.com/LopezCastroRoberto): Kernels and quantization +- [@mawong-amd](https://github.com/mawong-amd) +- [@micah-wil](https://github.com/micah-wil): ROCm / AMD GPU integration, CI +- [@netanel-haber](https://github.com/netanel-haber): Nemotron/NVIDIA models, Mamba/Linear attention hybrids, VLMs +- [@Rohan138](https://github.com/Rohan138): ROCm performance, CI, torch.compile fusions +- [@simondanielsson](https://github.com/simondanielsson): ROCm performance (CDNA), MoRI-IO. +- [@taneem-ibrahim](https://github.com/taneem-ibrahim): Pooling models +- [@TheEpicDolphin](https://github.com/TheEpicDolphin) +- [@varun-sundar-rabindranath](https://github.com/varun-sundar-rabindranath): KV cache offloading, LoRA +- [@vllmellm](https://github.com/vllmellm) +- [@wangxiyuan](https://github.com/wangxiyuan): Platform, KV cache, Mooncake + +## The Reviewer Role + +Reviewers assess the problem and proposed solution, guide validation, resolve concerns with authors, and determine when a contribution is ready for a committer. A reviewer's approval means a PR is ready for a committer's final look. + +Reviewers can: + +- Approve PRs or request changes. +- Trigger CI when a contribution is ready for evaluation. +- Triage PRs and issues, apply labels, and flag important PRs for expedited committer attention. +- Close duplicate PRs and issues. +- Escalate design questions to the appropriate committer or SIG. +- Flag PRs as ready for committer review. + +### Becoming a Reviewer + +Candidates should demonstrate a sustained record of useful reviews, sound engineering judgment, constructive communication, and familiarity with vLLM's contribution standards. Deep ownership of a subsystem is not required. + +A reviewer is nominated by two committers from different organizations, who post the nomination in the committers Slack channel. As a rule of thumb, not a hard requirement, candidates have either 15 merged PRs and 30 reviewed PRs, or 45 reviewed PRs. A reviewed PR is one where the candidate gave a substantive review and the PR reached a resolution (merged or closed). diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 89f29afce0..f8e9cfcbb1 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -307,15 +307,23 @@ review process: - After the PR is submitted, the PR will be assigned to a reviewer. Every reviewer will pick up the PRs based on their expertise and availability. + Official reviews can be performed by any of the + [reviewers](../community/reviewers.md) or + [committers](../governance/committers.md); committers perform the final + merge. To find someone to review your PR, see the reviewers' and + committers' areas of expertise or ask in the `#pr-reviews` channel on + [Slack](https://slack.vllm.ai). - After the PR is assigned, the reviewer will provide status updates every 2-3 days. If the PR is not reviewed within 7 days, please feel free to ping the reviewer or the vLLM team. - After the review, the reviewer will put an `action-required` label on the PR if there are changes required. The contributor should address the comments and ping the reviewer to re-review the PR. -- Please respond to all comments within a reasonable time frame. If a comment +- Please respond to all comments within a reasonable time frame. Address each + piece of reviewer feedback with a change or an explanation. If a comment isn't clear or you disagree with a suggestion, feel free to ask for - clarification or discuss the suggestion. + clarification or discuss the suggestion. If you and a reviewer can't reach + agreement, a committer can help decide. - Note that not all CI checks will be executed due to limited computational resources. Reviewers with write access and configured trusted contributors can comment `/ci run` for upstream CI or `/amd-ci run` for AMD CI only when diff --git a/docs/contributing/editing-agent-instructions.md b/docs/contributing/editing-agent-instructions.md index 7be5f594cd..2810b6fdca 100644 --- a/docs/contributing/editing-agent-instructions.md +++ b/docs/contributing/editing-agent-instructions.md @@ -44,6 +44,19 @@ cross-file context, and fixes for repeated mistakes. Each entry should be short, specific, and actionable — e.g., which files to touch, what order to change them in, and which tests to run. +## Skills + +Skills live in `.agents/skills//` and follow the +[Agent Skills specification](https://agentskills.io/specification). + +- Claude Code does not read `.agents/`, so symlink each skill: + `ln -s ../../.agents/skills/ .claude/skills/`. +- Frontmatter: `name` matches the directory. `description` says what the + skill does and when to use it; quote it if it contains a colon. +- Keep `SKILL.md` under 500 lines and move detail into files it links to. +- `tools/pre_commit/check_agent_files.py` runs in pre-commit to check the + skill and `AGENTS.md` rules + ## Keeping Docs Lean - Every addition should trigger review of surrounding content for stale or redundant items. diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index 0b1a9151ec..6c71aa1f46 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -37,6 +37,7 @@ th { | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | | flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| passthrough | standard | nvfp4,mxfp4 | G | N | Y | [`PassThroughPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough.PassThroughPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 182559141c..a88c83ee57 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -135,6 +135,7 @@ Batch invariance has been tested and verified on the following models: - **Granite 3.1 (Dense)**: `ibm-granite/granite-3.1-2b-instruct`, `ibm-granite/granite-3.1-8b-instruct` - **EXAONE 4.0 series**: `LGAI-EXAONE/EXAONE-4.0-1.2B`, `LGAI-EXAONE/EXAONE-4.0.1-32B`, `LGAI-EXAONE/EXAONE-4.0-32B` - **OLMo 2**: `allenai/OLMo-2-0425-1B-Instruct` +- **ERNIE 4.5**: `baidu/ERNIE-4.5-0.3B-PT` - **SmolLM2**: `HuggingFaceTB/SmolLM2-1.7B-Instruct` - **PLaMo3**: `pfnet/plamo-3-nict-2b-base` diff --git a/docs/features/disagg_prefill.md b/docs/features/disagg_prefill.md index 013d57fa87..8fd67ca4c2 100644 --- a/docs/features/disagg_prefill.md +++ b/docs/features/disagg_prefill.md @@ -79,6 +79,12 @@ decode = client.chat.completions.create( If `messages` has non-text content or `echo` is set, the ids are ignored and `messages` is rendered instead, so it must match the prefill request. Otherwise `kv_transfer_params["prompt_token_ids"]` must be a non-empty list of non-negative integers, or the request fails with HTTP 400, as it always does on `/v1/chat/completions/batch`. +## Generate API output modes + +When the prefill and decode stages use the [Generate API](../serving/online_serving/token_in_token_out.md) (`/inference/v1/generate`), only the decode response reaches the client, so set `output_mode` on the decode request only. A proxy that reuses the client's request body for the prefill request must reset `output_mode` to `tokens` there. A prefill instance started with `--tokens-only` has no tokenizer and rejects `output_mode: "text"` with a 400. + +A decode instance that returns text needs a tokenizer, so start it with `--enable-scale-out` and without `--tokens-only`. Prefill instances can keep `--tokens-only`. + ## Development We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance. diff --git a/docs/features/initialized_snapshots.md b/docs/features/initialized_snapshots.md index 15cce7554b..bf7fdbe370 100644 --- a/docs/features/initialized_snapshots.md +++ b/docs/features/initialized_snapshots.md @@ -31,6 +31,11 @@ Snapshots currently require: Current Hugging Face hub clients hold their connections for the process lifetime, so download the model in a separate step and run create with `HF_HUB_OFFLINE=1`, as the quickstart below does. +- Snapshot creation defaults `NCCL_IB_DISABLE=1` for its singleton donor and + inherited workers because CRIU cannot capture live InfiniBand/RDMA state. + An explicit caller value is retained, but creation rejects an open + `/dev/infiniband/` descriptor before CRIU. Close non-NCCL RDMA clients before + capture. - A remote model ID and an immutable 40-character `--revision`. Local model directories and mutable revisions are not supported. - Enough disk for the artifact, with the same installed vLLM package, model @@ -138,9 +143,13 @@ snapshot or external CRIU operation may use a shared `/dev/shm` mount at a time. and directories. Recreate the artifact after an unclean host shutdown. - Restore currently requires the same host, GPU, driver, kernel, Python, PyTorch, installed vLLM version, model revision, engine arguments, - selected environment variables, and CRIU plugin binaries. + selected effective environment variables, and CRIU plugin binaries. An unset + `NCCL_IB_DISABLE` therefore matches the snapshot donor default of `1`, while + an explicit different value does not. - Only dense float16 TP1 has been validated. Other model formats depend on their - existing sleep level 2 reload support; NCCL state is not restored. + existing sleep level 2 reload support; distributed and RDMA snapshots are not + supported. The snapshot-only NCCL default does not change ordinary serve or + multi-GPU defaults. - CRIU support varies by kernel and driver. Preserve package, library, model, and generated-cache paths for the artifact lifetime. - A snapshot can include application secrets or request state present in the diff --git a/docs/features/quantization/modelopt.md b/docs/features/quantization/modelopt.md index c93e7e15df..a624261527 100644 --- a/docs/features/quantization/modelopt.md +++ b/docs/features/quantization/modelopt.md @@ -33,9 +33,7 @@ following `quantization.quant_algo` values: `flashinfer_cutlass`, `flashinfer_cutedsl`, `flashinfer_trtllm`, `flashinfer_cudnn`, and `marlin`; see `KernelConfig` on the [Engine Arguments](../../configuration/engine_args.md) page and shown by - `vllm serve --help=KernelConfig`. For `W4A16_NVFP4`, `auto` currently - selects Marlin. Models with BF16 activations can explicitly select the - FlashInfer CuTe-DSL backend with `--linear-backend flashinfer_cutedsl`. + `vllm serve --help=KernelConfig`. !!! note For models quantized to MXFP8 with BF16 activations on SM100-family GPUs, diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index a0b406e6a4..0d0d0c0c85 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -440,7 +440,8 @@ ref changes. Use `--no-cache-filter extensions-build` to refresh an empty, branch, or tag revision. For `FINAL_BASE_IMAGE`, use the public, multi-arch -`nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04` image. +`nvidia/cuda:13.4.1-base-ubuntu24.04` image. Set `NCCL_VERSION` to 2.32.3 or +newer, the first NCCL release with Rubin (SM107) support. For `BUILD_BASE_IMAGE`, use: - `pytorch/manylinux2_28-builder:cuda13.4` for x86_64 CPUs. @@ -462,7 +463,8 @@ For `BUILD_BASE_IMAGE`, use: --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg CUDA_VERSION=13.4 \ --build-arg BUILD_BASE_IMAGE="pytorch/manylinuxaarch64-builder:cuda13.4" \ - --build-arg FINAL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04" \ + --build-arg FINAL_BASE_IMAGE="nvidia/cuda:13.4.1-base-ubuntu24.04" \ + --build-arg NCCL_VERSION=2.32.3 \ . ``` @@ -482,7 +484,8 @@ For `BUILD_BASE_IMAGE`, use: --build-arg TRITON_INSTALL_FROM_SOURCE_REVISION=3f6e41132b5edf639bfb872ad73d4688765e08b8 \ --build-arg CUDA_VERSION=13.4 \ --build-arg BUILD_BASE_IMAGE="pytorch/manylinux2_28-builder:cuda13.4" \ - --build-arg FINAL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base:26.08-cuda13.4-devel-ubuntu24.04" \ + --build-arg FINAL_BASE_IMAGE="nvidia/cuda:13.4.1-base-ubuntu24.04" \ + --build-arg NCCL_VERSION=2.32.3 \ . ``` @@ -490,9 +493,6 @@ For `BUILD_BASE_IMAGE`, use: Keep the default explicit `torch_cuda_arch_list`. GPU-less BuildKit builds cannot inspect the host GPU. R100 and VR200 report compute capability 10.7, for which the generic `10.0` target provides family-compatible kernels. - The Ubuntu `devel` final image is also required: the corresponding `base` - image lacks the CUDA runtime/JIT package closure used by vLLM and the - prerelease PyTorch wheel. `RUN_WHEEL_CHECK=false` disables only the PyPI publication-size guard for this private staging image. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 545890e152..e461379bcb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -410,7 +410,6 @@ th { | `GlmMoeDsaForCausalLM` | GLM-5, GLM-5.1, GLM-5.2 | `zai-org/GLM-5`, etc. | ✅︎ | ✅︎ | | `GPT2LMHeadModel` | GPT-2 | `openai-community/gpt2`, `openai-community/gpt2-xl`, etc. | | ✅︎ | | `GPTJForCausalLM` | GPT-J | `EleutherAI/gpt-j-6b`, `nomic-ai/gpt4all-j`, etc. | | ✅︎ | -| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ | | `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b`, `openai/gpt-oss-20b` | ✅︎ | ✅︎ | | `GraniteForCausalLM` | Granite 3.0, Granite 3.1, PowerLM | `ibm-granite/granite-3.0-2b-base`, `ibm-granite/granite-3.1-8b-instruct`, `ibm/PowerLM-3b`, etc. | ✅︎ | ✅︎ | | `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ | @@ -426,7 +425,6 @@ th { | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | | | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | | -| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ | | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ | | `K2HorizonForCausalLM` | K2-Horizon | `IFM/K2-Horizon-7B`, `IFM/K2-Horizon-375B` | | ✅︎ | | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ | @@ -456,7 +454,6 @@ th { | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ | | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ | | `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ | -| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ | | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ | | `Plamo3ForCausalLM` | PLaMo3 | `pfnet/plamo-3-nict-2b-base`, `pfnet/plamo-3-nict-8b-base`, etc. | ✅ | ✅︎ | @@ -468,7 +465,6 @@ th { | `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | | | `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ | | `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ | -| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ | | `SolarForCausalLM` | Solar Pro | `upstage/solar-pro-preview-instruct`, etc. | ✅︎ | ✅︎ | | `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | | | `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ | @@ -484,16 +480,20 @@ Some models are supported only via the [Transformers modeling backend](#transfor | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | | `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ | -| `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | | ✅︎ | +| `FlexOlmoForCausalLM` | FlexOlmo | `allenai/FlexOlmo-7x7B-1T`, `allenai/FlexOlmo-7x7B-1T-RT`, etc. | ✅︎ | ✅︎ | | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | +| `GPTNeoXForCausalLM` | GPT-NeoX, Pythia, OpenAssistant, Dolly V2, StableLM | `EleutherAI/gpt-neox-20b`, `EleutherAI/pythia-12b`, `OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5`, `databricks/dolly-v2-12b`, `stabilityai/stablelm-tuned-alpha-7b`, etc. | | ✅︎ | | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ | | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ | -| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ | +| `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | ✅︎ | ✅︎ | +| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | ✅︎ | ✅︎ | | `NanbeigeForCausalLM` | Nanbeige4.2 | `Nanbeige/Nanbeige4.2-3B`, etc. | ✅︎ | ✅︎ | | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ | +| `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ | +| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ | | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ | | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ | | `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ | diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index ea60233645..f3b3296102 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -27,6 +27,7 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios | | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios | | `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads | +| `passthrough` | Any | Internal, not accepted on the CLI: bound automatically for MoE backends that dispatch and combine themselves, so no all-to-all runs in the framework | `--moe-backend flashinfer_moe_ep_*` (FlashInfer MoE-EP megakernels) | | `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes | ## Single Node Deployment diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index a292a4a3a4..8e1a4f3398 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -125,8 +125,10 @@ Scale-out APIs are disabled by default on `vllm serve`. Set `--enable-scale-out` ### Tokens IN <> Tokens OUT APIs -- `/inference/v1/generate` - Generate completions -- `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set) +- [Generate API](token_in_token_out.md) (`/inference/v1/generate`) + - Generate completions from token IDs, optionally with detokenized text (`output_mode`) +- `/inference/v1/abort_requests` - Abort in-flight requests (registered wherever `/inference/v1/generate` is) +- `/abort_requests` - Unauthenticated alias of `/inference/v1/abort_requests` (`--tokens-only` only) ### Renderer APIs diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index cf57f5d8b2..9585031312 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -40,6 +40,8 @@ Seeding only reads the last few IDs (7 today), so a suffix of the prompt is enou Each request wraps the engine's `GenerateResponse`(s) together with the caller metadata needed to reconstruct the final response without a GPU. +Only `output_mode: "tokens"` responses are accepted which is the default. A response from a request with `output_mode: "text"` gets a 400 because its text is already detokenized with stop strings removed. Responses without `output_mode`, from older servers, are treated as `tokens`. + `/v1/chat/completions/derender`: ??? code @@ -127,7 +129,7 @@ Streaming derender is stateless on the server side: all mutable state lives in t - `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings - `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0 -When a streamed `GenerateResponseStreamChoice` carries `logprobs`, the `token_id:N` placeholders are resolved per chunk and the resolved logprobs are attached to the corresponding streamed choice — for chat as `ChatCompletionLogProbs`, for completions converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. +When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the `token_id:N` placeholders are resolved per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. ## Example diff --git a/docs/serving/online_serving/token_in_token_out.md b/docs/serving/online_serving/token_in_token_out.md new file mode 100644 index 0000000000..b4e0d465bd --- /dev/null +++ b/docs/serving/online_serving/token_in_token_out.md @@ -0,0 +1,71 @@ +# Tokens In <> Tokens Out API + +`/inference/v1/generate` takes prompt token IDs and returns generated token IDs. It is the generation step of the `render → generate → derender` pipeline used in disaggregated serving (see the [Renderer APIs](renderer.md) and [Derenderer APIs](derenderer.md)). + +The endpoint is registered on `vllm serve` with `--enable-scale-out`, and always with `--tokens-only`. `--tokens-only` also skips loading the tokenizer, so the server can't turn token IDs back into text. + +## Output modes + +The request field `output_mode` selects how much of the postprocessing the server does for the caller. + +| `output_mode` | Response adds | Server needs | +| --- | --- | --- | +| `tokens` (default) | nothing, token IDs only | nothing | +| `text` | `text` on every choice and decoded tokens with `bytes` in `logprobs` | a tokenizer | + +`text` returns text and token IDs in a single call, without a separate [derender](derenderer.md) hop per chunk. This suits latency sensitive streaming, where the derender hop lands on every output token. + +Every response and every stream chunk echoes `output_mode`. Servers that predate the field ignore it and return token IDs only with a 200, so check that the `output_mode` in the response matches the one you sent. Older servers don't return it. + +```python +import httpx + +response = httpx.post( + "http://localhost:8000/inference/v1/generate", + json={ + "token_ids": [151644, 872, 198], + "sampling_params": {"max_tokens": 32}, + "output_mode": "text", + }, +).json() + +assert response["output_mode"] == "text" +print(response["choices"][0]["text"]) +print(response["choices"][0]["token_ids"]) +``` + +### Text + +- `text` comes from the request's own detokenizer, so `skip_special_tokens`, `spaces_between_special_tokens`, `stop` and `include_stop_str_in_output` behave as they do on `/v1/completions`. +- A matched stop string is cut from `text` unless `include_stop_str_in_output` is set. `token_ids` keeps every generated token, so decoding `token_ids` yourself brings the stop string back. +- `/derender` only accepts `output_mode: "tokens"` responses and returns a 400 for text responses which are already detokenized. +- With `stream: true`, each choice's `text` is the delta since the previous chunk. A chunk is sent whenever the engine output carries new text or a `finish_reason`, even with no new token IDs. That covers text held back for stop string matching and the final output after an abort. + +### Logprobs + +With `output_mode: "tokens"`, logprob entries carry `token_id:N` placeholders. With `output_mode: "text"`, they carry the decoded token strings and their UTF-8 `bytes` for the sampled token and every entry in `top_logprobs`. A server started with `--return-tokens-as-token-ids` always returns the placeholders, as `/v1/completions` does. + +### Errors + +The server returns a 400 for: + +- `output_mode: "text"` on a server without a tokenizer (`--tokens-only` or `--skip-tokenizer-init`). +- `output_mode: "text"` with `sampling_params.detokenize: false`. +- An unsupported `output_mode` value. + +For using `output_mode` with separate prefill and decode pools, see [Disaggregated Prefilling](../../features/disagg_prefill.md#generate-api-output-modes). + +## Aborting requests + +`POST /inference/v1/abort_requests` aborts in-flight requests. It is registered wherever `/inference/v1/generate` is and requires the API key when `--api-key` is set, like `/inference/v1/generate`. + +```bash +curl -X POST http://localhost:8000/inference/v1/abort_requests \ + -H "Authorization: Bearer $VLLM_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"request_ids": ["generate-tokens-42"]}' +``` + +The request ID is the `request_id` the server returned: the `X-Request-Id` header or body `request_id` you sent, prefixed with `generate-tokens-`. A missing `request_ids` returns a 400. The response is empty and the abort finishes in the background. + +With `--tokens-only`, the same handler is also served at `POST /abort_requests` for existing deployments. That path doesn't require the API key, even when `--api-key` is set. See [Security](../../usage/security.md#unprotected-endpoints-no-api-key-required). diff --git a/docs/training/prompt_token_id_logprobs.md b/docs/training/prompt_token_id_logprobs.md index 940facad4f..6414c00bfe 100644 --- a/docs/training/prompt_token_id_logprobs.md +++ b/docs/training/prompt_token_id_logprobs.md @@ -2,14 +2,14 @@ In on-policy distillation, the student records its top-K candidate tokens at every response position, and the teacher provides the log probabilities of -exactly those candidates as the distillation target. The caller passes the -candidates as `prompt_logprob_token_ids`; the teacher runs a single prefill -over the prompt plus response and gathers their log probabilities on the GPU at -every scored row, starting from `prompt_logprob_start` (typically the prompt -length minus one, so only response positions are scored). The result is a -`[rows, K]` matrix whose row `i` scores prompt token -`prompt_logprob_start + i + 1` and whose columns follow the candidate order, so -it lines up position by position with the student's top-K. +exactly those candidates as the distillation target. The caller passes one +candidate list per scored row as `prompt_logprob_token_ids`; the teacher runs a +single prefill over the prompt plus response and gathers their log +probabilities on the GPU, starting from `prompt_logprob_start` (typically the +prompt length minus one, so only response positions are scored). The result is +a `[rows, K]` matrix whose row `i` scores prompt token +`prompt_logprob_start + i + 1` and whose columns follow that row's candidate +order, so it lines up position by position with the student's top-K. ## Quick start @@ -22,13 +22,13 @@ output = llm.generate( TokensPrompt(prompt_token_ids=prompt_ids + response_ids), SamplingParams( max_tokens=1, - prompt_logprob_token_ids=candidate_ids, + prompt_logprob_token_ids=student_topk_ids, # [len(response_ids), K] prompt_logprob_start=len(prompt_ids) - 1, ), ) scores = output[0].prompt_token_id_logprobs -# scores.shape == (len(response_ids), len(candidate_ids)) -# scores[i, j] = log p(candidate_ids[j] | prompt_ids + response_ids[:i]) +# scores.shape == (len(response_ids), K) +# scores[i, j] = log p(student_topk_ids[i][j] | prompt_ids + response_ids[:i]) ``` The same parameters are accepted by the `/inference/v1/generate` HTTP endpoint @@ -42,13 +42,17 @@ import numpy as np scores = np.load(io.BytesIO(base64.b64decode(response["prompt_token_id_logprobs"]))) ``` -The matrix has `max(prompt_len - 1 - prompt_logprob_start, 0)` rows. With a -logits `--logprobs-mode`, it holds logits instead of log probabilities. +`prompt_logprob_token_ids` is an integer array or a list of lists with exactly +`prompt_len - 1 - prompt_logprob_start` rows; shorter rows are padded with `-1`, +and every `-1` entry scores `-inf`. Pass an `int32` NumPy array, such as the +student's top-K IDs, which is sent to the engine without conversion; nested +lists of Python ints are much slower to copy and serialize. With a logits +`--logprobs-mode`, the matrix holds logits instead of log probabilities. ## Requirements - The V2 model runner (`VLLM_USE_V2_MODEL_RUNNER=1`). -- `len(prompt_logprob_token_ids)` at most `--max-logprobs`. +- The longest row at most `--max-logprobs` candidates. - No `--kv-sharing-fast-prefill`. The request skips reading the prefix cache (local and KV connector), since @@ -56,8 +60,6 @@ cached rows have no logits to score. ## Limitations -- One candidate set per request, shared by every scored row; per-position - candidates are planned ([#56860](https://github.com/vllm-project/vllm/issues/56860)). - Non-streaming only; `stream=true` requests are rejected. - Not exposed through the OpenAI-compatible endpoints. - If a prefill starts past the first scored row (e.g. `skip_reading_prefix_cache=False` diff --git a/docs/usage/security.md b/docs/usage/security.md index 2b09ee5266..d350b6bdfd 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -206,7 +206,7 @@ The following endpoints **do not require authentication** even when `--api-key` - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF - `/update_weights` - Update model weights (can alter model behavior) - `/get_world_size` - Get distributed world size -- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`) +- `/abort_requests` - Abort in-flight requests (available with `--tokens-only`. Use the authenticated `/inference/v1/abort_requests` otherwise) **Utility endpoints:** diff --git a/pyproject.toml b/pyproject.toml index b2658168d3..890dd6bdb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -197,7 +197,7 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "examples/pooling/token_embed/*", "tests/models/language/pooling/*", "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", - "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", + "docs/governance/process.md", "docs/community/reviewers.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", "rust/src/parser/**", "rust/src/text/src/output/decoded.rs", diff --git a/requirements/common.txt b/requirements/common.txt index b6b33efa92..6e0b0db783 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -9,7 +9,7 @@ blake3 py-cpuinfo transformers >= 5.16.1, < 5.18.0 huggingface_hub >= 1.31.0 # Required for the `huggingface_hub.utils.httpx` re-export -tokenizers >= 0.22.0 # Required for fast incremental detokenization. +tokenizers >= 0.23.2 # Required for fast incremental detokenization. safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*; platform_machine != "s390x" protobuf == 7.36.1; platform_machine == "s390x" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 5cd1e717be..fbc80b28e1 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -26,7 +26,7 @@ fastsafetensors >= 0.3.3 instanttensor >= 0.1.9 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl[cu13]==4.7.1 +nvidia-cutlass-dsl[cu13]==4.8.0 quack-kernels==0.6.5 # Required by MSA # Tokenspeed_MLA for faster mla with spec decode diff --git a/requirements/rubin-prerelease.txt b/requirements/rubin-prerelease.txt index e5915b6d13..6b416b48f6 100644 --- a/requirements/rubin-prerelease.txt +++ b/requirements/rubin-prerelease.txt @@ -3,24 +3,13 @@ # TODO(rubin-bringup): delete once Rubin APIs ship in regular PyTorch wheels # (vllm-project/vllm#49735). --extra-index-url https://download.pytorch.org/whl/nightly/cu134 ---extra-index-url https://flashinfer.ai/whl/cu134 ---extra-index-url https://flashinfer.ai/whl # Match each architecture's torch pin to the latest torchvision wheel. -torch==2.15.0.dev20260909+cu134; platform_machine == "x86_64" -torch==2.15.0.dev20260908+cu134; platform_machine == "aarch64" -torchvision==0.30.0.dev20260909+cu134 -torchaudio==2.11.0.dev20260911+cu134 - -# docker/Dockerfile strips these packages from requirements/cuda.txt whenever -# INSTALL_RUBIN_PRERELEASE=true, so they must be pinned here instead to -# still get installed. -nvidia-cutlass-dsl[cu13]==4.8.0.dev0 -flashinfer-python==0.7.0.post1 -quack-kernels==0.6.5 +torch==2.15.0.dev20260928+cu134; platform_machine == "x86_64" +torch==2.15.0.dev20260926+cu134; platform_machine == "aarch64" +torchvision==0.30.0.dev20260928+cu134 +torchaudio==2.11.0.dev20260928+cu134 cuda-python==13.4.1 cuda-bindings==13.4.1 -flashinfer-jit-cache==0.7.0.post1+cu134 -flashinfer-cubin==0.7.0.post1 nixl==1.4.1 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index 11fc2991ed..926a5ceb1c 100644 --- a/requirements/test/cpu.txt +++ b/requirements/test/cpu.txt @@ -1165,7 +1165,7 @@ timm==1.0.17 # -r requirements/test/cuda.in # open-clip-torch # segmentation-models-pytorch -tokenizers==0.23.1 +tokenizers==0.23.2 # via # -r requirements/test/../common.txt # cohere diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index cebfd726d5..8f39bb900e 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -1266,7 +1266,7 @@ timm==1.0.17 # -r requirements/test/cuda.in # open-clip-torch # segmentation-models-pytorch -tokenizers==0.23.1 +tokenizers==0.23.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 3714afaf9d..0df6f9c82a 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -1249,7 +1249,7 @@ timm==1.0.17 # -r requirements/test/rocm.in # open-clip-torch # segmentation-models-pytorch -tokenizers==0.23.1 +tokenizers==0.23.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 01a0fda14b..33b89a85fc 100644 --- a/requirements/test/xpu.txt +++ b/requirements/test/xpu.txt @@ -1089,7 +1089,7 @@ timm==1.0.28 # via # -r requirements/test/xpu.in # open-clip-torch -tokenizers==0.23.1 +tokenizers==0.23.2 # via # -c requirements/common.txt # -r requirements/test/../common.txt diff --git a/requirements/tpu.txt b/requirements/tpu.txt index 15e0e9ceab..d8906d8fc1 100644 --- a/requirements/tpu.txt +++ b/requirements/tpu.txt @@ -12,4 +12,4 @@ ray[data] setuptools==78.1.0 setuptools-rust>=1.9.0 nixl==0.3.0 -tpu-inference==0.29.0 +tpu-inference==0.30.0 diff --git a/rust/CLAUDE.md b/rust/CLAUDE.md deleted file mode 100644 index f96fd34544..0000000000 --- a/rust/CLAUDE.md +++ /dev/null @@ -1,4 +0,0 @@ -# CLAUDE.md - -First, check @AGENTS.override.md if exists. -Then, follow instructions in @AGENTS.md. diff --git a/rust/src/bench/CLAUDE.md b/rust/src/bench/CLAUDE.md deleted file mode 100644 index 43c994c2d3..0000000000 --- a/rust/src/bench/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -@AGENTS.md diff --git a/rust/src/engine-core-client/src/protocol/request.rs b/rust/src/engine-core-client/src/protocol/request.rs index 8f38d7da08..02fc1c30ca 100644 --- a/rust/src/engine-core-client/src/protocol/request.rs +++ b/rust/src/engine-core-client/src/protocol/request.rs @@ -152,6 +152,11 @@ impl EngineCoreRequest { feature.extract_aux_frames(&mut aux_frames, threshold); } } + if let Some(params) = &mut self.sampling_params + && let Some(ids) = &mut params.prompt_logprob_token_ids + { + ids.extract_aux_frame(&mut aux_frames, threshold); + } aux_frames } } @@ -168,7 +173,7 @@ mod tests { PlaceholderRange, }; use crate::protocol::sampling::EngineCoreSamplingParams; - use crate::protocol::tensor::{WireArrayData, WireTensor}; + use crate::protocol::tensor::{WireArrayData, WireNdArray, WireTensor}; use crate::protocol::{decode_value, encode_msgpack}; const AUX_FRAME_THRESHOLD: usize = 256; @@ -294,4 +299,27 @@ mod tests { ); assert!(request.extract_aux_frames(AUX_FRAME_THRESHOLD).is_empty()); } + + #[test] + fn engine_core_request_extracts_large_prompt_logprob_token_ids() { + let ids = WireNdArray::from_i32(vec![AUX_FRAME_THRESHOLD, 1], vec![5; AUX_FRAME_THRESHOLD]) + .unwrap(); + let mut request = EngineCoreRequest { + sampling_params: Some(EngineCoreSamplingParams { + prompt_logprob_token_ids: Some(ids), + ..EngineCoreSamplingParams::for_test() + }), + ..EngineCoreRequest::default() + }; + + let aux_frames = request.extract_aux_frames(AUX_FRAME_THRESHOLD); + + assert_eq!(aux_frames.len(), 1); + assert_eq!(aux_frames[0].len(), AUX_FRAME_THRESHOLD * 4); + let params = request.sampling_params.unwrap(); + assert_eq!( + params.prompt_logprob_token_ids.unwrap().data, + WireArrayData::AuxIndex(1) + ); + } } diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index 0cd11b5397..577dcb651f 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; use crate::protocol::structured_outputs::StructuredOutputsParams; +use crate::protocol::tensor::WireNdArray; fn default_top_p() -> f32 { 1.0 @@ -101,9 +102,9 @@ pub struct EngineCoreSamplingParams { /// /// `None` disables prompt logprobs. `-1` requests the full vocabulary. pub prompt_logprobs: Option, - /// Candidate token IDs scored at every scored causal prompt row, where row - /// `i` scores them as predictions of prompt token `i + 1`. - pub prompt_logprob_token_ids: Option>, + /// `[num_rows, num_ids]` candidate token IDs per scored causal prompt row, + /// `-1` padding shorter rows. + pub prompt_logprob_token_ids: Option, /// First causal prompt row to score; `None` scores from the first row. pub prompt_logprob_start: Option, /// Minimum probability threshold for token sampling. diff --git a/rust/src/engine-core-client/src/protocol/tensor.rs b/rust/src/engine-core-client/src/protocol/tensor.rs index 2bd193f76b..a79c7df5fd 100644 --- a/rust/src/engine-core-client/src/protocol/tensor.rs +++ b/rust/src/engine-core-client/src/protocol/tensor.rs @@ -121,6 +121,18 @@ impl + Serialize + DeserializeOwned> WireArray { )) } + /// Build an int32 tensor/ndarray backed by native-endian raw-view bytes. + /// + /// Takes ownership of the backing buffer without copying its data. + pub fn from_i32(shape: Vec, data: Vec) -> Result { + validate_element_count(&shape, data.len())?; + Ok(Self::from_raw_bytes( + TensorDtype::I32, + shape, + bytes_from_pod_vec(data), + )) + } + /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes. /// /// Takes ownership of the backing buffer without copying its data. diff --git a/rust/src/server/src/routes/inference/generate/validate.rs b/rust/src/server/src/routes/inference/generate/validate.rs index 9b1779d6e9..f9e821f945 100644 --- a/rust/src/server/src/routes/inference/generate/validate.rs +++ b/rust/src/server/src/routes/inference/generate/validate.rs @@ -181,7 +181,7 @@ mod tests { let request: GenerateRequest = serde_json::from_value(json!({ "token_ids": [11, 22], "stream": true, - "sampling_params": {"prompt_logprob_token_ids": [1, 2]} + "sampling_params": {"prompt_logprob_token_ids": [[1, 2]]} })) .expect("parse request"); assert!(validate_request_compat(&request, &served(&["test-model"])).is_err()); diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 01a4fb25b8..cf53704374 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -7,7 +7,7 @@ pub(crate) mod logprobs; pub(crate) mod sampling; pub(crate) mod token_ids; -use logprobs::validate_logprobs; +use logprobs::{lower_prompt_logprob_token_ids, validate_logprobs}; use sampling::validate_resolved_sampling_params; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use vllm_engine_core_client::protocol::sampling::{ @@ -50,7 +50,11 @@ pub fn lower_text_request( // multimodal tensor payloads. mm_features: request.mm_features.take(), sampling_params: lower_sampling_params( - request.sampling_params.clone(), + SamplingParams { + // Move the candidate table rather than cloning it; only lowering reads it. + prompt_logprob_token_ids: request.sampling_params.prompt_logprob_token_ids.take(), + ..request.sampling_params.clone() + }, sampling_hints, sampling_limits, prompt_len, @@ -180,7 +184,12 @@ pub fn lower_sampling_params( thinking_token_budget, logprobs, prompt_logprobs, - prompt_logprob_token_ids, + prompt_logprob_token_ids: lower_prompt_logprob_token_ids( + prompt_logprob_token_ids, + prompt_logprob_start, + prompt_len, + sampling_limits.model_vocab_size, + )?, prompt_logprob_start, min_p, frequency_penalty, @@ -332,6 +341,7 @@ mod tests { use vllm_engine_core_client::protocol::multimodal::{ MmFeatureSpec, MmModality, PlaceholderRange, }; + use vllm_engine_core_client::protocol::tensor::WireNdArray; use vllm_tokenizer::test_utils::TestTokenizer; use super::*; @@ -1140,7 +1150,7 @@ mod tests { #[test] fn lower_sampling_params_validates_prompt_logprob_token_ids() { - let lower = |ids: Option>, start: Option| { + let lower = |ids: Option>>, start: Option| { lower_sampling_params_with_limits( SamplingParams { prompt_logprob_token_ids: ids, @@ -1155,23 +1165,30 @@ mod tests { Ok(_) => panic!("expected rejection"), }; - let params = lower(Some(vec![1, 2]), Some(1)).unwrap(); - assert_eq!(params.prompt_logprob_token_ids, Some(vec![1, 2])); - assert_eq!(params.prompt_logprob_start, Some(1)); + let params = lower(Some(vec![vec![1, 2], vec![3]]), None).unwrap(); + assert_eq!( + params.prompt_logprob_token_ids, + Some(WireNdArray::from_i32(vec![2, 2], vec![1, 2, 3, -1]).unwrap()) + ); + assert_eq!( + lower(Some(vec![vec![5]]), Some(1)).unwrap().prompt_logprob_start, + Some(1) + ); expect_test::expect![[r#" [ - "prompt_logprob_token_ids must not be empty.", + "prompt_logprob_token_ids must be a non-empty integer array of shape [num_rows, num_ids].", "requested prompt_logprob_token_ids of 21, which is greater than max allowed: 20", - "token_id(s) [1000] in prompt_logprob_token_ids are out of vocabulary. Vocabulary size: 1000", - "prompt_logprob_token_ids must not contain duplicates.", + "prompt_logprob_token_ids contain out-of-vocab token ids (-1 pads a row). Vocabulary size: 1000", + "prompt_logprob_token_ids contain out-of-vocab token ids (-1 pads a row). Vocabulary size: 1000", + "prompt_logprob_token_ids has 1 rows, but the prompt has 2 scored rows (prompt_len - 1 - prompt_logprob_start).", "prompt_logprob_start requires prompt_logprob_token_ids.", ] - "#]] - .assert_debug_eq(&[ - rejects(Some(vec![]), None), - rejects(Some((0..21).collect()), None), - rejects(Some(vec![1000]), None), - rejects(Some(vec![1, 1]), None), + "#]].assert_debug_eq(&[ + rejects(Some(vec![vec![]]), None), + rejects(Some(vec![(0..21).collect(), vec![0]]), None), + rejects(Some(vec![vec![1000], vec![0]]), None), + rejects(Some(vec![vec![-2], vec![0]]), None), + rejects(Some(vec![vec![1]]), None), rejects(None, Some(0)), ]); } diff --git a/rust/src/text/src/lower/logprobs.rs b/rust/src/text/src/lower/logprobs.rs index 58ad5d694b..407a47c9e1 100644 --- a/rust/src/text/src/lower/logprobs.rs +++ b/rust/src/text/src/lower/logprobs.rs @@ -3,11 +3,11 @@ //! Python-compatible validation for logprobs sampling params. //! -//! `-1` is expanded only for bounds checks. The original request values are -//! passed through to engine-core. +//! `-1` counts are expanded only for bounds checks and passed to engine-core +//! unchanged; `prompt_logprob_token_ids` rows are padded into an array. -use itertools::Itertools as _; use thiserror::Error; +use vllm_engine_core_client::protocol::tensor::WireNdArray; use crate::backend::SamplingLimits; @@ -36,10 +36,21 @@ pub enum LogprobsError { len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}." )] TokenIdsMismatch { logprobs: i32, num_token_ids: usize }, - #[error("prompt_logprob_token_ids must not be empty.")] + #[error( + "prompt_logprob_token_ids must be a non-empty integer array of shape \ + [num_rows, num_ids]." + )] EmptyPromptLogprobTokenIds, - #[error("prompt_logprob_token_ids must not contain duplicates.")] - DuplicatePromptLogprobTokenIds, + #[error( + "prompt_logprob_token_ids contain out-of-vocab token ids (-1 pads a row). \ + Vocabulary size: {vocab_size}" + )] + PromptLogprobTokenIdsOutOfVocab { vocab_size: usize }, + #[error( + "prompt_logprob_token_ids has {rows} rows, but the prompt has {expected} scored \ + rows (prompt_len - 1 - prompt_logprob_start)." + )] + PromptLogprobRowsMismatch { rows: usize, expected: usize }, #[error("prompt_logprob_start requires prompt_logprob_token_ids.")] PromptLogprobStartWithoutTokenIds, } @@ -49,7 +60,7 @@ pub(super) fn validate_logprobs( logprobs: Option, prompt_logprobs: Option, logprob_token_ids: Option<&[u32]>, - prompt_logprob_token_ids: Option<&[u32]>, + prompt_logprob_token_ids: Option<&[Vec]>, prompt_logprob_start: Option, sampling_limits: SamplingLimits, ) -> Result<(), LogprobsError> { @@ -60,7 +71,7 @@ pub(super) fn validate_logprobs( validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?; validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?; validate_logprobs_count( - prompt_logprob_token_ids.map(|ids| ids.len() as i32), + prompt_logprob_token_ids.map(|rows| rows.iter().map(Vec::len).max().unwrap_or(0) as i32), max_logprobs, vocab_size, "prompt_logprob_token_ids", @@ -70,17 +81,49 @@ pub(super) fn validate_logprobs( } fn validate_prompt_logprob_token_ids( - token_ids: Option<&[u32]>, + rows: Option<&[Vec]>, start: Option, ) -> Result<(), LogprobsError> { - match token_ids { - Some([]) => Err(LogprobsError::EmptyPromptLogprobTokenIds), - Some(ids) if !ids.iter().all_unique() => Err(LogprobsError::DuplicatePromptLogprobTokenIds), + match rows { + Some(rows) if rows.iter().all(Vec::is_empty) => { + Err(LogprobsError::EmptyPromptLogprobTokenIds) + } None if start.is_some() => Err(LogprobsError::PromptLogprobStartWithoutTokenIds), _ => Ok(()), } } +/// Check per-row candidates against the vocabulary and the prompt as Python +/// does, and pad them with `-1` into the engine's `[num_rows, num_ids]` array. +pub(super) fn lower_prompt_logprob_token_ids( + rows: Option>>, + start: Option, + prompt_len: u32, + vocab_size: usize, +) -> Result, LogprobsError> { + let Some(rows) = rows else { + return Ok(None); + }; + if rows.iter().flatten().any(|&id| id < -1 || i64::from(id) >= vocab_size as i64) { + return Err(LogprobsError::PromptLogprobTokenIdsOutOfVocab { vocab_size }); + } + let expected = (prompt_len as usize).saturating_sub(1 + start.unwrap_or(0) as usize); + if rows.len() != expected { + return Err(LogprobsError::PromptLogprobRowsMismatch { + rows: rows.len(), + expected, + }); + } + let width = rows.iter().map(Vec::len).max().unwrap_or(0); + let data = rows + .iter() + .flat_map(|row| row.iter().copied().chain(std::iter::repeat_n(-1, width - row.len()))) + .collect(); + Ok(Some( + WireNdArray::from_i32(vec![rows.len(), width], data).expect("padded rows are rectangular"), + )) +} + fn validate_logprobs_count( requested: Option, max_logprobs: usize, diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs index 240b0528e1..8512a830e6 100644 --- a/rust/src/text/src/lower/token_ids.rs +++ b/rust/src/text/src/lower/token_ids.rs @@ -95,14 +95,6 @@ pub(crate) fn validate_vocab_range( )?; } - if let Some(token_ids) = params.prompt_logprob_token_ids.as_deref() { - validate_param( - "prompt_logprob_token_ids", - token_ids.iter().copied(), - limits.model_vocab_size, - )?; - } - if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() { validate_param( "bad_words", diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 9c9f22a38e..c37e21060e 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -97,9 +97,9 @@ pub struct SamplingParams { /// /// `None` disables prompt logprobs. `-1` requests the full vocabulary. pub prompt_logprobs: Option, - /// Candidate token IDs scored at every scored causal prompt row, where row - /// `i` scores them as predictions of prompt token `i + 1`. - pub prompt_logprob_token_ids: Option>, + /// Candidate token IDs per scored causal prompt row, where row `i` scores + /// its IDs as predictions of prompt token `prompt_logprob_start + i + 1`. + pub prompt_logprob_token_ids: Option>>, /// First causal prompt row to score; `None` scores from the first row. pub prompt_logprob_start: Option, /// Minimum probability threshold for token sampling. `None` means no diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index 36a54ae219..fa06facf29 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -327,6 +327,35 @@ def test_reentered_tag(free_x_early): allocator.release_pools() +@pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"]) +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") +@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator") +def test_runtime_state_survives_sleep(level): + """Runtime state is offloaded at both levels and restored by the first + selective wake; a nested kv_cache pool keeps its own tag.""" + from vllm.device_allocator.sleep_mode_backend import CuMemBackend + + allocator = get_mem_allocator_instance() + with allocator.use_memory_pool("runtime"): + state = torch.arange(1 << 20, device=DEVICE_TYPE) + with allocator.use_memory_pool("kv_cache"): + kv = torch.zeros(1 << 20, device=DEVICE_TYPE) + assert {d.tag for d in allocator.pointer_to_data.values()} == { + "runtime", + "kv_cache", + } + + backend = CuMemBackend() + backend.suspend(level=level) + assert mapped_usage(allocator) == 0 + + backend.resume(tags=["weights"]) + assert torch.equal(state, torch.arange(1 << 20, device=DEVICE_TYPE)) + backend.resume(tags=["kv_cache"]) + kv.fill_(1) + assert int(kv.sum()) == kv.numel() + + @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator") def test_level2_discards_ordinary_tensor_with_weights_tag(): diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index a2cbe0135b..db7ba05a57 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -6,14 +6,17 @@ from collections import defaultdict import pytest import regex as re -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode from .common import FUSION_LOG_PATTERNS, AttentionBackendCase, Matches def run_model( - compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs + vllm_runner, + compile_config: CompilationMode | CompilationConfig, + model: str, + **model_kwargs, ): """Run a model with the given compilation config for E2E fusion tests.""" compilation_config = ( @@ -36,38 +39,48 @@ def run_model( # No cudagraphs by default if compilation_config.cudagraph_mode is None: compilation_config.cudagraph_mode = CUDAGraphMode.NONE - llm = LLM( - model=model, + + # trust_remote_code/enable_chunked_prefill/block_size/kernel_config are + # pinned to LLM()'s own defaults since VllmRunner overrides them for + # other test needs (e.g. it forces block_size=16, which some attention + # backends like FLASHINFER_MLA reject, and disables JIT warmup). + model_kwargs.setdefault("kernel_config", {"enable_jit_warmup": True}) + with vllm_runner( + model, compilation_config=compilation_config, + trust_remote_code=False, + enable_chunked_prefill=None, + block_size=None, **model_kwargs, - ) - outputs = llm.generate(prompts, sampling_params) + ) as vllm_model: + llm = vllm_model.llm + outputs = llm.generate(prompts, sampling_params) - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") + # Print the outputs. + for output in outputs: + prompt = output.prompt + generated_text = output.outputs[0].text + print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - # Get the compile ranges endpoints after vllm config post init - # in order to compute compile ranges correctly - compilation_config.compile_ranges_endpoints = ( - llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints - ) + # Get the compile ranges endpoints after vllm config post init + # in order to compute compile ranges correctly + compilation_config.compile_ranges_endpoints = ( + llm.llm_engine.vllm_config.compilation_config.compile_ranges_endpoints + ) - # Fetch match table from each worker via RPC and sum across workers. - worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc( - "get_compilation_match_table" - ) - combined: defaultdict[str, int] = defaultdict(int) - for table in worker_tables: - for k, v in table.items(): - combined[k] += v - return dict(combined) + # Fetch match table from each worker via RPC and sum across workers. + worker_tables: list[dict[str, int]] = llm.llm_engine.engine_core.collective_rpc( + "get_compilation_match_table" + ) + combined: defaultdict[str, int] = defaultdict(int) + for table in worker_tables: + for k, v in table.items(): + combined[k] += v + return dict(combined) @pytest.fixture -def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): +def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn, vllm_runner): def run( model_name: str, matches: Matches, @@ -137,7 +150,9 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): ) with caplog_mp_spawn(logging.DEBUG) as log_holder: - match_table = run_model(full_compilation_config, model_name, **model_kwargs) + match_table = run_model( + vllm_runner, full_compilation_config, model_name, **model_kwargs + ) num_compile_ranges = len(full_compilation_config.get_compile_ranges()) assert num_compile_ranges in [1, 2, 3] diff --git a/tests/compile/fusions_e2e/models.py b/tests/compile/fusions_e2e/models.py index eadf1739fd..082ab0bf52 100644 --- a/tests/compile/fusions_e2e/models.py +++ b/tests/compile/fusions_e2e/models.py @@ -101,6 +101,20 @@ llama3_8b_fp4 = ModelFusionInfo( ), ) +# ModelOpt MIXED_PRECISION: FP8 attention projections, NVFP4 MLPs. It resolves to +# modelopt_mixed, not modelopt_fp4, and its MLP has no manual act+quant fusion, so +# every MLP boundary is left to ActivationQuantFusionPass. +qwen3_8_27b_mixed_fp4 = ModelFusionInfo( + model_name="nvidia/Qwen3.8-27B-NVFP4", + hf_overrides=lambda n_layers: { + "text_config": { + "num_hidden_layers": n_layers, + "layer_types": ["linear_attention"] * (n_layers - 1) + ["full_attention"], + } + }, + matches=lambda n_layers: Matches(act_quant_fusion=n_layers), +) + # MoEs cannot do act+quant fusion because those ops are hidden from torch.compile. # MoEs also only expose 1 rms+quant fusion because the quant for up_proj is hidden. # TODO(luka): https://github.com/vllm-project/vllm/issues/31985 diff --git a/tests/compile/fusions_e2e/test_tp1_quant.py b/tests/compile/fusions_e2e/test_tp1_quant.py index fa1af8e1c6..cec0cbecae 100644 --- a/tests/compile/fusions_e2e/test_tp1_quant.py +++ b/tests/compile/fusions_e2e/test_tp1_quant.py @@ -30,6 +30,7 @@ from .models import ( llama3_8b_fp8, llama4_scout_fp4, llama4_scout_fp8, + qwen3_8_27b_mixed_fp4, qwen3_a3b_fp8, ) @@ -202,3 +203,46 @@ def test_tp1_fp4_fusions( compilation_config, matches_check, ) + + +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", [qwen3_8_27b_mixed_fp4] +) +@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +@pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4") +def test_tp1_fp4_mixed_precision_act_fusion( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + inductor_graph_partition: bool, + run_e2e_fusion_test, +): + """Act+quant fusion matches on a MIXED_PRECISION NVFP4 checkpoint.""" + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False} + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + pass_config=PassConfig( + fuse_norm_quant=False, + fuse_act_quant=True, + fuse_attn_quant=False, + enable_qk_norm_rope_fusion=False, + ), + ) + + run_e2e_fusion_test( + model_name, + matches_fn(n_layers), + model_kwargs, + attn_backend, + compilation_config, + ["act_quant_fusion"], + ) diff --git a/tests/config/test_modelopt_mixed_nvfp4.py b/tests/config/test_modelopt_mixed_nvfp4.py new file mode 100644 index 0000000000..b9a890144a --- /dev/null +++ b/tests/config/test_modelopt_mixed_nvfp4.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 detection for ModelOpt MIXED_PRECISION checkpoints. + +A checkpoint that quantizes some layers to NVFP4 and others to FP8 declares +``quant_algo: MIXED_PRECISION`` and resolves to ``modelopt_mixed``, so NVFP4 +has to be detected from the per-layer algorithms. +""" + +import pytest + +from vllm.config.model import _modelopt_mixed_has_nvfp4 + + +def _quantized_layers(*algos: str) -> dict: + return { + "quantized_layers": { + f"model.layers.{i}.proj": {"quant_algo": algo} + for i, algo in enumerate(algos) + } + } + + +def test_detects_nvfp4_alongside_another_algorithm(): + """The shipped shape: FP8 attention projections, NVFP4 MLP.""" + assert _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "NVFP4", "FP8")) + + +def test_ignores_a_checkpoint_with_no_nvfp4_layer(): + assert not _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "MXFP8")) + + +def test_weight_only_nvfp4_does_not_count(): + """W4A16_NVFP4 leaves activations in high precision, so nothing to fuse.""" + assert not _modelopt_mixed_has_nvfp4(_quantized_layers("FP8", "W4A16_NVFP4")) + + +@pytest.mark.parametrize( + "quant_config", + [None, {}, {"quantized_layers": None}, {"quantized_layers": ["NVFP4"]}], +) +def test_absent_or_malformed_config_is_not_nvfp4(quant_config): + assert not _modelopt_mixed_has_nvfp4(quant_config) diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 0158b342f3..ac73188534 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -511,7 +511,7 @@ def _check_dummy_hash_model_forward( model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model) torch.nn.Module.__init__(model) - model.use_mega_moe = False + model.use_native_mega_moe = False model.use_sequence_parallel = False model.fuse_mhc_all_reduce = False model.engram_hash = state diff --git a/tests/distributed/test_eplb_migration_scheduler.py b/tests/distributed/test_eplb_migration_scheduler.py new file mode 100644 index 0000000000..1d63afb268 --- /dev/null +++ b/tests/distributed/test_eplb_migration_scheduler.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +from vllm.config.parallel import EPLBConfig +from vllm.distributed.eplb.eplb_communicator import EplbCommunicator +from vllm.distributed.eplb.migration_scheduler import ( + MigrationFlow, + schedule_migration_batches, + schedule_migration_batches_for_layers, +) +from vllm.distributed.eplb.rebalance_execute import move_from_buffer, move_to_buffer + + +class _MockEplbCommunicator(EplbCommunicator): + def __init__(self) -> None: + self.send_calls: list[tuple[list[torch.Tensor], int, int]] = [] + self.recv_calls: list[tuple[list[torch.Tensor], int, int]] = [] + self.context_calls = 0 + self.execute_count = 0 + + def add_send( + self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int + ) -> None: + self.send_calls.append((tensors, dst_rank, expert_id)) + + def add_recv( + self, tensors: list[torch.Tensor], src_rank: int, expert_id: int + ) -> None: + self.recv_calls.append((tensors, src_rank, expert_id)) + + def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None: + self.context_calls += 1 + + def execute(self) -> None: + self.execute_count += 1 + + +def test_migration_batching_is_disabled_by_default_and_can_be_enabled() -> None: + assert not EPLBConfig().enable_migration_batching + assert EPLBConfig(enable_migration_batching=True).enable_migration_batching + + +def test_migration_batching_only_applies_to_async_eplb() -> None: + assert not EPLBConfig(use_async=True).migration_batching_enabled + assert EPLBConfig( + use_async=True, enable_migration_batching=True + ).migration_batching_enabled + assert not EPLBConfig( + use_async=True, enable_migration_batching=False + ).migration_batching_enabled + assert not EPLBConfig( + use_async=False, enable_migration_batching=True + ).migration_batching_enabled + + +def test_schedule_migration_batches_is_deterministic() -> None: + transfers = [ + (1, 3, 0), + (2, 4, 1), + (0, 1, 2), + (0, 2, 3), + (0, 3, 4), + (0, 4, 5), + ] + num_local_experts, old_indices, new_indices = _placements_for_transfers( + 5, transfers + ) + + batches = schedule_migration_batches(num_local_experts, old_indices, new_indices) + + assert [_endpoints(batch) for batch in batches] == [ + [(1, 3), (2, 4)], + [(0, 1)], + [(0, 2)], + [(0, 3)], + [(0, 4)], + ] + for batch in batches: + _assert_no_endpoint_conflict(batch) + + +def test_schedule_migration_batches_coalesces_rank_pair() -> None: + transfers = [ + (0, 1, 0), + (0, 1, 1), + (2, 3, 2), + ] + num_local_experts, old_indices, new_indices = _placements_for_transfers( + 4, transfers + ) + + assert schedule_migration_batches(num_local_experts, old_indices, new_indices) == [ + [ + MigrationFlow(0, 1, expert_ids=(0, 1)), + MigrationFlow(2, 3, expert_ids=(2,)), + ] + ] + + +def test_schedule_migration_batches_covers_random_instructions() -> None: + rng = np.random.default_rng(42) + transfers = [] + for expert_id in range(80): + src_rank, dst_rank = rng.integers(0, 8, size=2).tolist() + if src_rank != dst_rank: + transfers.append((src_rank, dst_rank, expert_id)) + num_local_experts, old_indices, new_indices = _placements_for_transfers( + 8, transfers + ) + + batches = schedule_migration_batches(num_local_experts, old_indices, new_indices) + + scheduled = _flatten_transfers(batches) + assert len(scheduled) == len(transfers) + assert set(scheduled) == set(transfers) + for batch in batches: + _assert_no_endpoint_conflict(batch) + + +def test_schedule_migration_batches_excludes_local_copies() -> None: + # Old: rank 0 [0, 1], rank 1 [1, 2]. New: rank 1 needs expert 0. + old_indices = np.array([0, 1, 1, 2], dtype=np.int64) + new_indices = np.array([0, 1, 0, 2], dtype=np.int64) + + assert schedule_migration_batches(2, old_indices, new_indices) == [ + [MigrationFlow(0, 1, expert_ids=(0,))] + ] + + +def test_schedule_migration_batches_requires_matching_shapes() -> None: + old_indices = np.array([0, 1], dtype=np.int64) + new_indices = np.array([0], dtype=np.int64) + + with pytest.raises(AssertionError): + schedule_migration_batches(1, old_indices, new_indices) + + +def test_schedule_migration_batches_balances_replicas() -> None: + # Rank 0 and rank 1 hold expert 0; ranks 2, 3, and 4 need it. + old_indices = np.array([0, 0, 1, 1, 1], dtype=np.int64) + new_indices = np.array([0, 0, 0, 0, 0], dtype=np.int64) + + batches = schedule_migration_batches(1, old_indices, new_indices) + + assert set(_flatten_transfers(batches)) == { + (0, 2, 0), + (0, 4, 0), + (1, 3, 0), + } + + +def test_schedule_migration_batches_for_layers_matches_individual_calls() -> None: + old_indices = np.array([[0, 1, 2, 3], [3, 2, 1, 0]], dtype=np.int64) + new_indices = np.array([[1, 2, 3, 0], [0, 3, 2, 1]], dtype=np.int64) + + assert schedule_migration_batches_for_layers(1, old_indices, new_indices) == [ + schedule_migration_batches(1, old_indices[0], new_indices[0]), + schedule_migration_batches(1, old_indices[1], new_indices[1]), + ] + + +@pytest.mark.parametrize("use_precomputed", [False, True]) +def test_move_to_buffer_uses_multiple_batches(use_precomputed: bool) -> None: + old_indices = np.array([0, 1, 2, 3], dtype=np.int64) + new_indices = np.array([1, 2, 3, 0], dtype=np.int64) + communicator = _MockEplbCommunicator() + migration_batches = ( + schedule_migration_batches(1, old_indices, new_indices) + if use_precomputed + else None + ) + + move_to_buffer( + num_local_experts=1, + old_indices=old_indices, + new_indices=new_indices, + expert_weights=[torch.zeros(1, 1)], + expert_weights_buffers=[torch.zeros(1, 1)], + stream=None, + ep_rank=0, + communicator=communicator, + layer_idx=7, + enable_migration_batching=True, + migration_batches=migration_batches, + ) + + assert communicator.context_calls == 2 + assert communicator.execute_count == 2 + assert [(dst, expert) for _, dst, expert in communicator.send_calls] == [(3, 0)] + assert [(src, expert) for _, src, expert in communicator.recv_calls] == [(1, 1)] + + +def test_move_to_buffer_uses_primary_duplicate_destination() -> None: + old_indices = np.array([0, 1, 2, 3], dtype=np.int64) + new_indices = np.array([2, 2, 0, 3], dtype=np.int64) + expert_weights = [torch.zeros(2, 1)] + recv_buffer = torch.zeros(2, 1) + communicator = _MockEplbCommunicator() + + transfer_metadata = move_to_buffer( + num_local_experts=2, + old_indices=old_indices, + new_indices=new_indices, + expert_weights=expert_weights, + expert_weights_buffers=[recv_buffer], + stream=None, + ep_rank=0, + communicator=communicator, + enable_migration_batching=True, + ) + + assert len(communicator.recv_calls) == 1 + recv_tensors, src_rank, expert_id = communicator.recv_calls[0] + assert (src_rank, expert_id) == (1, 2) + assert recv_tensors[0].data_ptr() == recv_buffer[0].data_ptr() + + recv_buffer[0].fill_(42) + move_from_buffer( + expert_weights=expert_weights, + expert_weights_buffers=[recv_buffer], + transfer_metadata=transfer_metadata, + new_indices=new_indices, + ep_rank=0, + ) + torch.testing.assert_close(expert_weights[0], torch.full((2, 1), 42.0)) + + +def _endpoints( + flows: list[MigrationFlow], +) -> list[tuple[int, int]]: + return [(item.src_rank, item.dst_rank) for item in flows] + + +def _assert_no_endpoint_conflict(batch: list[MigrationFlow]) -> None: + endpoints: set[int] = set() + for flow in batch: + assert flow.src_rank not in endpoints + assert flow.dst_rank not in endpoints + endpoints.update((flow.src_rank, flow.dst_rank)) + + +def _flatten_transfers( + batches: list[list[MigrationFlow]], +) -> list[tuple[int, int, int]]: + return [ + (flow.src_rank, flow.dst_rank, expert_id) + for batch in batches + for flow in batch + for expert_id in flow.expert_ids + ] + + +def _placements_for_transfers( + num_ranks: int, + transfers: list[tuple[int, int, int]], +) -> tuple[int, np.ndarray, np.ndarray]: + old_counts = [0] * num_ranks + new_counts = [0] * num_ranks + for src_rank, dst_rank, _ in transfers: + old_counts[src_rank] += 1 + new_counts[dst_rank] += 1 + num_local_experts = max(1, *old_counts, *new_counts) + old_indices = np.full(num_ranks * num_local_experts, -1, dtype=np.int64) + new_indices = np.full_like(old_indices, -1) + old_counts = [0] * num_ranks + new_counts = [0] * num_ranks + for src_rank, dst_rank, expert_id in transfers: + old_offset = src_rank * num_local_experts + old_counts[src_rank] + new_offset = dst_rank * num_local_experts + new_counts[dst_rank] + old_indices[old_offset] = expert_id + new_indices[new_offset] = expert_id + old_counts[src_rank] += 1 + new_counts[dst_rank] += 1 + return num_local_experts, old_indices, new_indices diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 89fe6eb6ad..f7dac37940 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -231,6 +231,8 @@ def test_one_sided_combine_into_compatibility(supports_output): payload, runtime_max_tokens_per_rank, output=None, + *, + use_low_precision=False, ): result = payload + runtime_max_tokens_per_rank if output is None: @@ -249,6 +251,33 @@ def test_one_sided_combine_into_compatibility(supports_output): torch.testing.assert_close(output, payload + 2) +@pytest.mark.parametrize("low_precision_combine", [False, True]) +def test_one_sided_combine_into_low_precision(low_precision_combine): + """The fp8 combine setting is passed through to the kernel.""" + from vllm.distributed.device_communicators.all2all import ( + FlashInferNVLinkOneSidedManager, + ) + + seen_kwargs = {} + + class FakeMoeAlltoAll: + def combine(self, payload, runtime_max_tokens_per_rank, output, **kwargs): + seen_kwargs.update(kwargs) + output.copy_(payload + runtime_max_tokens_per_rank) + + manager = FlashInferNVLinkOneSidedManager.__new__(FlashInferNVLinkOneSidedManager) + manager.moe_alltoall = FakeMoeAlltoAll() + manager._combine_supports_output = True + manager.low_precision_combine = low_precision_combine + payload = torch.arange(4, dtype=torch.float32) + output = torch.empty_like(payload) + + manager.combine_into(payload, runtime_max_tokens_per_rank=2, output=output) + + assert seen_kwargs == {"use_low_precision": low_precision_combine} + torch.testing.assert_close(output, payload + 2) + + # --------------------------------------------------------------------------- # Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init) # --------------------------------------------------------------------------- @@ -915,6 +944,106 @@ def test_one_sided_dispatch_combine(world_size): _spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size) +def _one_sided_low_precision_worker(rank, world_size): + from vllm.distributed.device_communicators.all2all import ( + FlashInferNVLinkOneSidedManager, + ) + from vllm.distributed.parallel_state import get_dp_group + + device = torch.device(f"cuda:{rank}") + hidden_size = 256 + tokens_per_rank = 32 + experts_per_token = 2 + num_experts = world_size * 8 + + manager = FlashInferNVLinkOneSidedManager(get_dp_group().cpu_group) + manager.initialize( + max_num_tokens=tokens_per_rank, + top_k=experts_per_token, + num_experts=num_experts, + hidden_size=hidden_size, + x_bytes_per_token=hidden_size // 2, + x_sf_bytes_per_token=hidden_size // 16, + ) + assert manager.low_precision_combine + + torch.manual_seed(rank + 42) + x = torch.randint( + 0, 256, (tokens_per_rank, hidden_size // 2), device=device, dtype=torch.uint8 + ) + x_sf = torch.randint( + 0, 256, (tokens_per_rank, hidden_size // 16), device=device, dtype=torch.uint8 + ) + topk_ids = torch.randint( + 0, + num_experts, + (tokens_per_rank, experts_per_token), + device=device, + dtype=torch.int32, + ) + topk_weights = torch.rand( + tokens_per_rank, experts_per_token, device=device, dtype=torch.float32 + ) + + manager.moe_alltoall.dispatch( + token_selected_experts=topk_ids, + input_payloads=[x, x_sf, topk_ids, topk_weights], + runtime_max_tokens_per_rank=tokens_per_rank, + ) + + # Go through combine_into rather than moe_alltoall.combine: the workspace + # was sized for fp8, so the kernel must be told to send fp8. Calling the + # raw kernel here would send bf16 into fp8-sized slots and corrupt it. + expert_output = torch.ones( + world_size, tokens_per_rank, hidden_size, device=device, dtype=torch.bfloat16 + ) + output = torch.empty( + tokens_per_rank, hidden_size, device=device, dtype=torch.bfloat16 + ) + manager.combine_into( + payload=expert_output, + runtime_max_tokens_per_rank=tokens_per_rank, + output=output, + ) + + # Integers 1..top_k are exact in fp8_e4m3, so the fp8 transport must match + # bf16 exactly: one contribution per distinct expert-owning rank. + experts_per_rank = num_experts // world_size + expert_ranks = topk_ids // experts_per_rank + num_distinct = torch.tensor( + [len(set(row.tolist())) for row in expert_ranks], + device=device, + dtype=torch.bfloat16, + ).unsqueeze(1) + torch.testing.assert_close(output, num_distinct.expand_as(output)) + + torch.distributed.barrier() + manager.cleanup() + + +@requires_multi_gpu +@requires_one_sided +@requires_ptrace +@pytest.mark.parametrize("world_size", [2]) +def test_one_sided_combine_low_precision(world_size, monkeypatch): + """Exercise the real fp8 combine transport, not a fake kernel. + + A workspace/dtype mismatch corrupts the combine buffer silently rather + than raising, so this asserts the combined values are exactly right. + """ + from flashinfer.comm.trtllm_moe_alltoall import MoeAlltoAll + + from vllm.utils.func_utils import supports_kw + + if not supports_kw( + MoeAlltoAll.combine, "use_low_precision", allow_var_kwargs=False + ): + pytest.skip("installed FlashInfer has no use_low_precision support") + + monkeypatch.setenv("VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE", "1") + _spawn_workers(_one_sided_low_precision_worker, world_size, dp_size=world_size) + + # --------------------------------------------------------------------------- # Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle # --------------------------------------------------------------------------- diff --git a/tests/distributed/test_shm_storage.py b/tests/distributed/test_shm_storage.py index 762ffad8e6..66c2751708 100644 --- a/tests/distributed/test_shm_storage.py +++ b/tests/distributed/test_shm_storage.py @@ -291,6 +291,93 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase): reader.touch("key", address, monotonic_id, signature) +class TestShmObjectStorageReferenceCounting(unittest.TestCase): + """One writer and four readers, as with tensor-parallel workers.""" + + def setUp(self): + ring_buffer = SingleWriterShmRingBuffer( + data_buffer_size=1024 * 100, create=True + ) + self.writer = SingleWriterShmObjectStorage( + max_object_size=1024 * 10, + n_readers=4, + ring_buffer=ring_buffer, + serde_class=MsgpackSerde, + ) + self.addCleanup(self.writer.close) + handle = self.writer.handle() + handle.reader_lock = Lock() + self.readers = [ + SingleWriterShmObjectStorage.create_from_handle(handle) for _ in range(4) + ] + for reader in self.readers: + self.addCleanup(reader.close) + + def _reference(self, key: str) -> tuple[int, int, list[int], str]: + """Take one writer reference to the key, as one request item does.""" + if self.writer.is_cached(key): + address, monotonic_id = self.writer.get_cached(key) + else: + address, monotonic_id = self.writer.put(key, "value") + return address, monotonic_id, self.writer.get_signature(key), key + + @staticmethod + def _touch(reader: SingleWriterShmObjectStorage, reference) -> None: + address, monotonic_id, signature, key = reference + reader.touch(key, address, monotonic_id, signature) + + def _is_evictable(self, key: str) -> bool: + self.writer.free_unused() + return not self.writer.is_cached(key) + + def test_item_referenced_again_before_every_reader_fetched_it(self): + """A second request must not leave the item unevictable forever.""" + first = self._reference("image") + self.writer.touch("image") + second = self._reference("image") + self.writer.release_touches() + + # Every reader starts on the first request; the ones that finish it + # go on to the second while the others are still fetching. + for reader in self.readers: + self._touch(reader, first) + for reader in self.readers: + reader.get(*first) + self._touch(reader, second) + reader.get(*second) + + self.assertTrue(self._is_evictable("image")) + + def test_new_item_twice_in_one_request_is_not_evicted_early(self): + """The item must stay until the last reader has fetched it.""" + self.writer.touch("image") + references = [self._reference("image"), self._reference("image")] + self.writer.release_touches() + + for reader in self.readers[:-1]: + for reference in references: + self._touch(reader, reference) + for reference in references: + reader.get(*reference) + for reference in references: + self._touch(self.readers[-1], reference) + self.assertFalse(self._is_evictable("image")) + + for reference in references: + self.readers[-1].get(*reference) + self.assertTrue(self._is_evictable("image")) + + def test_touch_protects_item_until_released(self): + reference = self._reference("image") + for reader in self.readers: + reader.get(*reference) + + self.writer.touch("image") + self.assertFalse(self._is_evictable("image")) + self.writer.release_touches() + self.assertTrue(self._is_evictable("image")) + + # Reader process function def reader_process(process_id, storage_handle, items_to_read): """Reader process that connects to existing shared memory and reads data.""" diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 85ae7f58c7..d044ef3bb7 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -7,12 +7,12 @@ Integration tests for NCCL and IPC weight transfer between processes using Ray. """ import builtins +import importlib.util import pickle import runpy import threading import time from contextlib import nullcontext -from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -1497,10 +1497,12 @@ class TestDeferredProcessingContract: that caller ends up reaching through a getattr.""" def _engines(self): - return { - name: loader() - for name, loader in WeightTransferEngineFactory._registry.items() - } + registry = dict(WeightTransferEngineFactory._registry) + # ModelExpress is an optional, separately installed package, but its + # backend is always registered. Skip it when the package is missing. + if importlib.util.find_spec("modelexpress") is None: + registry.pop("modelexpress") + return {name: loader() for name, loader in registry.items()} def test_every_engine_declares_whether_it_defers(self): for name, cls in self._engines().items(): @@ -2054,11 +2056,11 @@ def test_sparse_nccl_trainer_non_sender_skips_client(): def _import_modelexpress_shim(): - path = ( - Path(__file__).parents[2] - / "vllm/distributed/weight_transfer/modelexpress_engine.py" + spec = importlib.util.find_spec( + "vllm.distributed.weight_transfer.modelexpress_engine" ) - runpy.run_path(str(path), run_name="_test_modelexpress_shim") + assert spec is not None and spec.origin is not None + runpy.run_path(spec.origin, run_name="_test_modelexpress_shim") @pytest.mark.parametrize( diff --git a/tests/entrypoints/launchers/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py index fc83010698..55093c5685 100644 --- a/tests/entrypoints/launchers/test_launch_cli.py +++ b/tests/entrypoints/launchers/test_launch_cli.py @@ -237,6 +237,46 @@ def test_snapshot_environment_contract( ) +@pytest.mark.parametrize( + ("configured", "expected"), + [(None, "1"), ("0", "0"), ("1", "1")], +) +def test_snapshot_child_defaults_nccl_ib_without_mutating_parent( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + configured: str | None, + expected: str, +): + if configured is None: + monkeypatch.delenv("NCCL_IB_DISABLE", raising=False) + else: + monkeypatch.setenv("NCCL_IB_DISABLE", configured) + parent_environment = os.environ.copy() + process = create_autospec(subprocess.Popen, instance=True) + process.pid = 100 + popen = create_autospec(subprocess.Popen, return_value=process) + monkeypatch.setattr(snapshot_runtime.subprocess, "Popen", popen) + + LocalSnapshotTools().launch_child(tmp_path, ("model",)) + + assert popen.call_args.kwargs["env"]["NCCL_IB_DISABLE"] == expected + assert os.environ == parent_environment + + +def test_snapshot_environment_identity_applies_effective_nccl_ib_default( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv("NCCL_IB_DISABLE", raising=False) + implicit = LocalSnapshotTools()._environment_identity() + monkeypatch.setenv("NCCL_IB_DISABLE", "1") + explicit_disabled = LocalSnapshotTools()._environment_identity() + monkeypatch.setenv("NCCL_IB_DISABLE", "0") + explicit_enabled = LocalSnapshotTools()._environment_identity() + + assert implicit == explicit_disabled + assert explicit_enabled != implicit + + @pytest.mark.parametrize( "case", [ @@ -369,6 +409,33 @@ ss() { :; } assert result.returncode == (0 if probe is None else 1), result.stderr +@pytest.mark.parametrize("container_attempted", [False, True]) +def test_snapshot_e2e_cleanup_checks_state_only_after_container_attempt( + container_attempted: bool, +): + script = ( + Path(__file__).resolve().parents[3] + / ".buildkite/scripts/initialized-snapshot-e2e.sh" + ) + if not script.exists(): + pytest.skip("snapshot E2E script is not packaged in this test image") + source = script.read_text() + cleanup = source[source.index("cleanup() {") : source.index("on_exit() {")] + shell = ( + "set -Eeuo pipefail\n" + "CONTAINER_NAME='' PORT_TWO=18002\n" + f"CHECK_CLEANUP_STATE={int(container_attempted)}\n" + "wait_clean() { echo checked-state; return 1; }\n" + + cleanup + + "\nif cleanup; then exit 0; else exit 1; fi\n" + ) + + result = subprocess.run(["bash", "-c", shell], capture_output=True, timeout=10) + + assert result.returncode == int(container_attempted), result.stderr + assert (b"checked-state" in result.stdout) == container_attempted + + def test_snapshot_create_cli_accepts_only_pinned_compact_mode( monkeypatch: pytest.MonkeyPatch, ): @@ -652,6 +719,111 @@ def test_snapshot_create_rolls_back_failed_dump(tmp_path: Path): assert not target.exists() +@pytest.mark.parametrize( + "failure_kind", ["called-process", "timeout", "missing", "unreadable"] +) +def test_snapshot_create_preserves_bounded_criu_diagnostics_after_cleanup( + tmp_path: Path, failure_kind: str, caplog: pytest.LogCaptureFixture +): + target = tmp_path / "snapshot" + tools = _fake_snapshot_tools() + criu_tools = LocalSnapshotTools() + criu_tools.plugin_dir = tmp_path + original_run = criu_tools._run + criu_tools._privileged = lambda: [] # type: ignore[method-assign] + captured = "captured stderr " + "x" * 9000 + if failure_kind == "timeout": + primary: subprocess.SubprocessError = subprocess.TimeoutExpired( + ["criu", "dump"], 1, output=b"captured stdout\xff", stderr=captured.encode() + ) + else: + primary = subprocess.CalledProcessError( + 1, ["criu", "dump"], output="captured stdout", stderr=captured + ) + + def run(command: list[str], **_kwargs: object): + if "tail" not in command: + raise primary + if failure_kind == "unreadable": + raise PermissionError("log is root-owned") + if failure_kind == "timeout": + return subprocess.CompletedProcess(command, 0, "timeout log tail", "") + return original_run(command, **_kwargs) + + criu_tools._run = run # type: ignore[method-assign] + + def fail_dump(workdir: Path, _inventory: ProcessInventory) -> None: + images = workdir / "images" + images.mkdir() + if failure_kind != "missing": + (images / "dump.log").write_text( + "omitted dump prefix" + "d" * 9000 + "distinctive dump diagnostic" + ) + criu_tools._criu("dump", workdir, []) + + tools.dump.side_effect = fail_dump + with pytest.raises(type(primary)) as excinfo: + create_snapshot( + argparse.Namespace(snapshot_dir=str(target), model_tag="Qwen/Qwen3-0.6B"), + tools=tools, + ) + + assert excinfo.value is primary + assert not target.exists() + diagnostics = "\n".join(caplog.messages) + assert "captured stdout" in diagnostics + assert "captured stderr" not in diagnostics + assert "x" * 100 in diagnostics + assert len(diagnostics) < 17000 + if failure_kind in {"missing", "unreadable"}: + assert "dump.log unavailable" in diagnostics + expected_error = ( + "CalledProcessError" if failure_kind == "missing" else "PermissionError" + ) + assert expected_error in diagnostics + elif failure_kind == "called-process": + assert "distinctive dump diagnostic" in diagnostics + assert "omitted dump prefix" not in diagnostics + else: + assert "timeout log tail" in diagnostics + if failure_kind == "timeout": + assert "captured stdout�" in diagnostics + + +def test_snapshot_criu_failure_selects_restore_log( + tmp_path: Path, caplog: pytest.LogCaptureFixture +): + tools = LocalSnapshotTools() + tools.plugin_dir = tmp_path + tools._privileged = lambda: [] # type: ignore[method-assign] + artifact = tmp_path / "snapshot" + images = artifact / "images" + images.mkdir(parents=True) + (images / "restore.log").write_text( + "omitted restore prefix" + "r" * 9000 + "distinctive restore diagnostic" + ) + primary = subprocess.CalledProcessError(1, ["criu", "restore"]) + original_run = tools._run + tailed: list[str] = [] + + def run(command: list[str], **_kwargs: object): + if "tail" not in command: + raise primary + tailed.append(command[-1]) + return original_run(command, **_kwargs) + + tools._run = run # type: ignore[method-assign] + + with pytest.raises(subprocess.CalledProcessError) as excinfo: + tools._criu("restore", artifact, []) + + assert excinfo.value is primary + assert tailed == [str(images / "restore.log")] + diagnostics = "\n".join(caplog.messages) + assert "distinctive restore diagnostic" in diagnostics + assert "omitted restore prefix" not in diagnostics + + def _restore_fixture(tmp_path: Path) -> argparse.Namespace: artifact = tmp_path / "snapshot" artifact.mkdir(mode=0o700) @@ -868,6 +1040,96 @@ def test_snapshot_private_path_and_link_remap_security( tools.restore(artifact, _manifest()) +@pytest.mark.parametrize("inventory_error", [None, PermissionError("denied")]) +def test_snapshot_dump_rejects_unavailable_link_remap_inventory_before_criu( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + inventory_error: PermissionError | None, +): + tools = LocalSnapshotTools() + workdir = tmp_path / "snapshot" + workdir.mkdir() + tools.shm_dir = tmp_path / "missing-shm" + if inventory_error is not None: + tools.shm_dir.mkdir() + + def fail_inventory(_path: Path): + raise inventory_error + + monkeypatch.setattr(Path, "iterdir", fail_inventory) + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + + with pytest.raises(SnapshotCreateError, match="inventory CRIU link remaps"): + tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU")) + + +def test_snapshot_restore_rejects_saved_remap_without_prefix( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + _mark_snapshot_pids_free(monkeypatch) + artifact = _local_restore_artifact(tmp_path) + saved = artifact / "link-remaps" + saved.mkdir() + (saved / "270").write_bytes(b"invalid remap name") + tools = LocalSnapshotTools() + tools.shm_dir = tmp_path / "shm" + tools.shm_dir.mkdir() + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + + with pytest.raises(SnapshotRestoreError, match="invalid saved link remap"): + tools.restore(artifact, _manifest()) + + assert not tuple(tools.shm_dir.iterdir()) + + +def test_snapshot_link_remap_inventory_matches_only_numeric_suffixes(tmp_path: Path): + for name in ( + "link_remap.0", + "link_remap.270", + "link_remap.", + "link_remap.12a", + "link-remap.12", + "prefix.link_remap.12", + "270", + ): + (tmp_path / name).touch() + tools = LocalSnapshotTools() + tools.shm_dir = tmp_path + + assert tools._link_remap_names() == {"link_remap.0", "link_remap.270"} + + +def test_snapshot_dump_preserves_primary_failure_when_cleanup_inventory_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + tools = LocalSnapshotTools() + workdir = tmp_path / "snapshot" + workdir.mkdir() + inventory_calls = 0 + + def inventory_names() -> set[str]: + nonlocal inventory_calls + inventory_calls += 1 + if inventory_calls > 1: + raise SnapshotCreateError( + "could not inventory CRIU link remaps: PermissionError: denied" + ) + return set() + + monkeypatch.setattr(tools, "_link_remap_names", inventory_names) + + def fail_dump(*_args): + raise RuntimeError("dump failed") + + monkeypatch.setattr(tools, "_criu", fail_dump) + + with pytest.raises(SnapshotCreateError) as failure: + tools.dump(workdir, ProcessInventory(100, (100,), (100,), "GPU")) + + assert "RuntimeError: dump failed" in str(failure.value) + assert "cleanup inventory failed" in str(failure.value) + + @pytest.mark.parametrize( ("blocked_state", "message"), [("tcp", "external established TCP"), ("io_uring", "kernel.io_uring_disabled=1")], @@ -900,6 +1162,28 @@ def test_snapshot_rejects_unsafe_process_state_before_criu( assert "1.1.1.1:443" in str(excinfo.value) +def test_snapshot_rejects_open_infiniband_device_before_criu( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + tools = LocalSnapshotTools() + monkeypatch.setattr(tools, "_tree_pids", lambda _root_pid: (100, 101)) + monkeypatch.setattr(tools, "_cuda_process_rows", lambda: ("101, GPU-abc",)) + monkeypatch.setattr( + tools, + "_descriptor_targets", + lambda pid: ("/dev/infiniband/uverbs7",) if pid == 101 else (), + ) + monkeypatch.setattr(tools, "_criu", lambda *_args: pytest.fail("CRIU called")) + + with pytest.raises(SnapshotCreateError) as failure: + tools.inventory(100, tmp_path) + + message = str(failure.value) + assert "pid 101" in message + assert "/dev/infiniband/uverbs7" in message + assert "NCCL_IB_DISABLE=1" in message + + def test_decode_endpoint_families(): assert _decode_endpoint("AF_INET", "3C00000A:D6F6") == "10.0.0.60:55030" assert ( diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_video.py b/tests/entrypoints/multimodal/openai/chat_completion/test_video.py index a6ca410e07..2431dd5339 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_video.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_video.py @@ -40,13 +40,11 @@ def server(): "--trust-request-mm-kwargs", ] - # ROCm: Increase timeouts to handle potential network delays and slower - # video processing when downloading multiple videos from external sources + # ROCm: Allow more time to download videos from external sources. env_overrides = {} if current_platform.is_rocm(): env_overrides = { "VLLM_VIDEO_FETCH_TIMEOUT": "120", - "VLLM_ENGINE_ITERATION_TIMEOUT_S": "300", } with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_overrides) as remote_server: diff --git a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py index 56a7cf6a13..6438d50a78 100644 --- a/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py +++ b/tests/entrypoints/multimodal/openai/chat_completion/test_vision.py @@ -88,13 +88,11 @@ def server(): *ROCM_EXTRA_ARGS, ] - # ROCm: Increase timeouts to handle potential network delays and slower - # video processing when downloading multiple videos from external sources + # ROCm: Allow more time to download videos from external sources. env_overrides = { **( { "VLLM_VIDEO_FETCH_TIMEOUT": "120", - "VLLM_ENGINE_ITERATION_TIMEOUT_S": "300", } if current_platform.is_rocm() else {} diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index 5a0ff9363e..b8511ea2c9 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -1,20 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ScoringIOProcessor post-tokenization helpers.""" +"""Unit tests for ScoringIOProcessor helpers.""" from dataclasses import dataclass from types import SimpleNamespace import pytest +import torch from tokenizers import Tokenizer, models, pre_tokenizers, processors from transformers import TokenizersBackend from vllm import TokensPrompt from vllm.entrypoints.pooling.scoring.io_processor import ( + BiEncoderIOProcessor, CrossEncoderIOProcessor, + JinaRankingIOProcessor, + LateInteractionIOProcessor, _apply_post_tokenization_to_token_type_ids, ) from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids +from vllm.outputs import PoolingOutput, PoolingRequestOutput, RequestError from vllm.renderers import TokenizeParams pytestmark = pytest.mark.skip_global_cleanup @@ -113,3 +118,59 @@ def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt(): assert first_doc == 10 assert prompt_token_ids[:first_doc] == list(range(10, num_query)) assert prompt_token_ids[first_doc:40] == list(range(num_query, 50)) + + +def _pooling_output( + request_id: str, + data: torch.Tensor, + error: RequestError | None = None, +) -> PoolingRequestOutput: + return PoolingRequestOutput( + request_id=request_id, + outputs=PoolingOutput(data), + prompt_token_ids=[1], + num_cached_tokens=0, + finished=True, + error=error, + ) + + +@pytest.mark.parametrize( + "processor_type", [BiEncoderIOProcessor, LateInteractionIOProcessor] +) +@pytest.mark.parametrize("failed_index", [0, 1]) +def test_paired_scoring_preserves_request_error(processor_type, failed_index): + processor = processor_type.__new__(processor_type) + processor.pad_token_id = None + error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + outputs = [ + _pooling_output("query", torch.tensor([1.0, 2.0])), + _pooling_output("document", torch.tensor([3.0, 4.0])), + ] + outputs[failed_index] = _pooling_output( + outputs[failed_index].request_id, + torch.empty(0), + error, + ) + + (result,) = processor._post_process(outputs, n_queries=1) + + assert result.request_id == "query_document" + assert result.error is error + assert result.outputs is outputs[failed_index].outputs + + +def test_jina_scoring_preserves_request_error_before_indexing_output(): + processor = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) + error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + failed = _pooling_output("request", torch.empty(0), error) + + assert processor._post_process([failed], n_queries=1) == [failed] diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py b/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py index def1dc264f..f295c5fc04 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py @@ -49,10 +49,60 @@ def _query_key(context: PoolingServeContext) -> str: @pytest.mark.asyncio -async def test_colliding_request_ids_use_distinct_query_cache_keys( +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("failed_stage", ["query", "doc"]) +async def test_flash_late_interaction_cleans_up_queries_on_failure( + failed_stage: str, monkeypatch: pytest.MonkeyPatch, ): serving = object.__new__(ServingScores) + serving.io_processor = Mock() + serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock()) + + context = _make_context() + monkeypatch.setattr(serving, "_init_ctx", AsyncMock(return_value=context)) + monkeypatch.setattr(serving, "_preprocessing", AsyncMock()) + + async def encode_queries(ctx): + ctx.late_interaction_query_keys = ["query-key"] + if failed_stage == "query": + raise RuntimeError("query failed") + + async def encode_docs(ctx): + ctx.late_interaction_doc_keys = ["doc-key"] + if failed_stage == "doc": + raise RuntimeError("doc failed") + + monkeypatch.setattr( + serving, + "_flash_late_interaction_encode_queries", + AsyncMock(side_effect=encode_queries), + ) + monkeypatch.setattr( + serving, + "_flash_late_interaction_encode_docs", + AsyncMock(side_effect=encode_docs), + ) + + with pytest.raises(RuntimeError, match=f"{failed_stage} failed"): + await serving.flash_late_interaction() + + expected_doc_keys = [] if failed_stage == "query" else ["doc-key"] + serving.engine_client.abort.assert_awaited_once_with( + ["query-key", *expected_doc_keys] + ) + serving.engine_client.collective_rpc.assert_awaited_once_with( + "release_late_interaction_query_cache", + args=(["query-key"],), + ) + + +@pytest.mark.asyncio +async def test_colliding_request_ids_use_distinct_late_interaction_keys( + monkeypatch: pytest.MonkeyPatch, +): + serving = object.__new__(ServingScores) + serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock()) prepare_generators = AsyncMock() monkeypatch.setattr(serving, "_prepare_generators", prepare_generators) monkeypatch.setattr(serving, "_collect_batch", AsyncMock()) @@ -68,9 +118,24 @@ async def test_colliding_request_ids_use_distinct_query_cache_keys( first_doc_key = _query_key(prepared_contexts[1]) second_query_key = _query_key(prepared_contexts[2]) second_doc_key = _query_key(prepared_contexts[3]) + first_doc_request_ids = prepared_contexts[1].prompt_request_ids + second_doc_request_ids = prepared_contexts[3].prompt_request_ids assert first_query_key == first_doc_key assert second_query_key == second_doc_key assert first_query_key != second_query_key assert "caller-controlled-id" not in first_query_key assert "caller-controlled-id" not in second_query_key + assert first_doc_request_ids is not None + assert second_doc_request_ids is not None + assert first.late_interaction_doc_keys == first_doc_request_ids + assert second.late_interaction_doc_keys == second_doc_request_ids + assert first_doc_request_ids != second_doc_request_ids + assert all("caller-controlled-id" not in key for key in first_doc_request_ids) + assert all("caller-controlled-id" not in key for key in second_doc_request_ids) + + assert first.late_interaction_query_keys is not None + await serving._cleanup_flash_late_interaction(first) + serving.engine_client.abort.assert_awaited_once_with( + [*first.late_interaction_query_keys, *first_doc_request_ids] + ) diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index a89f8d68db..66d3e83969 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -22,6 +22,10 @@ The streaming cases do the same for `stream=true`. Each coupled chunk's `token_ids` becomes one `/inference/v1/generate` stream chunk, fed through the chunked derender path with `stream_state` threaded across calls the way a client would, so both parsers see the same chunk boundaries. + +The text level cases pin `/inference/v1/generate` with `output_mode="text"` +against `/v1/completions/derender` on the same token IDs, batch and streaming, +plus logprob resolution and stop string handling. """ import json @@ -440,3 +444,241 @@ async def test_stream_parity_reasoning_and_tool_call(client, chunking): assert any(len(ch["token_ids"]) > 1 for ch in coupled_choices) if not (coupled["reasoning"] and coupled["tool_calls"]): pytest.skip("Model did not emit both a block and a tool call") + + +# --------------------------------------------------------------------------- +# Inline text level: generate(output_mode="text") == derender(generate(...)) +# --------------------------------------------------------------------------- + +TEXT_MESSAGES = [ + {"role": "user", "content": "What is 2+2? Answer in one short sentence."} +] + + +async def _render_token_ids( + client: httpx.AsyncClient, messages: list[dict] +) -> list[int]: + resp = await client.post( + "/v1/chat/completions/render", json={"model": MODEL, "messages": messages} + ) + assert resp.status_code == 200, resp.text + return resp.json()["token_ids"] + + +def _generate_payload( + token_ids: list[int], output_mode: str, stream: bool = False, **sampling +) -> dict: + return { + "model": MODEL, + "token_ids": token_ids, + "sampling_params": {"temperature": 0, "max_tokens": 128, **sampling}, + "output_mode": output_mode, + "stream": stream, + } + + +async def _generate( + client: httpx.AsyncClient, token_ids: list[int], output_mode: str, **sampling +) -> dict: + resp = await client.post( + "/inference/v1/generate", + json=_generate_payload(token_ids, output_mode, **sampling), + ) + assert resp.status_code == 200, resp.text + data = resp.json() + assert data["output_mode"] == output_mode + return data + + +async def _generate_text_stream( + client: httpx.AsyncClient, token_ids: list[int], **sampling +) -> list[dict]: + """Stream generate at the text level and return the choice of every chunk + that has one.""" + choices: list[dict] = [] + async with client.stream( + "POST", + "/inference/v1/generate", + json=_generate_payload(token_ids, "text", stream=True, **sampling), + ) as resp: + assert resp.status_code == 200, await resp.aread() + async for line in resp.aiter_lines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + chunk = json.loads(line[len("data: ") :]) + assert "error" not in chunk, chunk + assert chunk["output_mode"] == "text" + choices.extend(chunk["choices"]) + return choices + + +async def _derender_completion_text( + client: httpx.AsyncClient, output_ids: list[int], finish_reason: str +) -> str: + resp = await client.post( + "/v1/completions/derender", + json={ + "model": MODEL, + "generate_responses": [ + { + "request_id": "text-parity", + "choices": [ + { + "index": 0, + "token_ids": output_ids, + "finish_reason": finish_reason, + } + ], + } + ], + }, + ) + assert resp.status_code == 200, resp.text + return resp.json()["choices"][0]["text"] + + +async def _derender_completion_stream_text( + client: httpx.AsyncClient, choices: list[dict] +) -> str: + """Replay each generate chunk's token IDs through the streaming + completions derender endpoint with the same chunk boundaries.""" + state = None + text = "" + for choice in choices: + resp = await client.post( + "/v1/completions/derender", + json={ + "stream": True, + "model": MODEL, + "generate_chunk": { + "request_id": "text-parity-stream", + "choices": [ + { + "index": 0, + "token_ids": choice.get("token_ids") or [], + "finish_reason": choice.get("finish_reason"), + } + ], + }, + "stream_state": state, + }, + ) + assert resp.status_code == 200, resp.text + data = resp.json() + state = data["stream_state"] + text += "".join(ch["text"] for ch in data["chunk"]["choices"]) + return text + + +@pytest.mark.asyncio +async def test_text_parity_batch(client): + """Inline text matches /v1/completions/derender on the same token IDs.""" + token_ids = await _render_token_ids(client, TEXT_MESSAGES) + inline = (await _generate(client, token_ids, "text"))["choices"][0] + + derendered = await _derender_completion_text( + client, inline["token_ids"], inline["finish_reason"] + ) + + assert inline["text"] + assert inline["text"] == derendered + + +@pytest.mark.asyncio +async def test_text_parity_stream(client): + """Streamed inline text matches both the streaming and the non-streaming + derender of the same token IDs.""" + token_ids = await _render_token_ids(client, TEXT_MESSAGES) + choices = await _generate_text_stream(client, token_ids) + inline = "".join(ch["text"] for ch in choices) + output_ids = [t for ch in choices for t in ch.get("token_ids") or []] + finish_reason = choices[-1]["finish_reason"] + + assert finish_reason is not None + assert inline == await _derender_completion_stream_text(client, choices) + assert inline == await _derender_completion_text(client, output_ids, finish_reason) + + +@pytest.mark.asyncio +async def test_text_parity_logprobs(client): + """Inline logprob tokens and bytes match /derender resolving the + token_id:N placeholders of the tokens level.""" + token_ids = await _render_token_ids(client, TEXT_MESSAGES) + sampling = {"logprobs": 3} + tokens_choice = (await _generate(client, token_ids, "tokens", **sampling))[ + "choices" + ][0] + text_choice = (await _generate(client, token_ids, "text", **sampling))["choices"][0] + assert text_choice["token_ids"] == tokens_choice["token_ids"], ( + "greedy (temperature=0) generation was expected to be deterministic " + "across the two generate calls" + ) + + disagg = await _disagg( + client, + tokens_choice["token_ids"], + len(token_ids), + tokens_choice["finish_reason"], + {"model": MODEL, "messages": TEXT_MESSAGES}, + logprobs=tokens_choice["logprobs"], + ) + + inline_content = text_choice["logprobs"]["content"] + derender_content = disagg["choices"][0]["logprobs"]["content"] + assert len(inline_content) == len(derender_content) + for inline_entry, derender_entry in zip(inline_content, derender_content): + assert inline_entry["token"] == derender_entry["token"] + assert inline_entry["bytes"] == derender_entry["bytes"] + assert [ + (top["token"], top["bytes"]) for top in inline_entry["top_logprobs"] + ] == [(top["token"], top["bytes"]) for top in derender_entry["top_logprobs"]] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_text_parity_stop_string(client, stream): + """Inline text follows the engine's stop handling, like /v1/completions. + + token_ids keep every generated token, so /derender brings the matched + stop string back. Its text must start with the inline text and only + differ after it. + """ + stop = "\n" + token_ids = await _render_token_ids( + client, + [{"role": "user", "content": "Write two short lines about the sea."}], + ) + + resp = await client.post( + "/v1/completions", + json={ + "model": MODEL, + "prompt": token_ids, + "temperature": 0, + "max_tokens": 128, + "stop": [stop], + }, + ) + assert resp.status_code == 200, resp.text + coupled = resp.json()["choices"][0] + + if stream: + choices = await _generate_text_stream(client, token_ids, stop=[stop]) + inline_text = "".join(ch["text"] for ch in choices) + output_ids = [t for ch in choices for t in ch.get("token_ids") or []] + finish_reason = choices[-1]["finish_reason"] + else: + inline = (await _generate(client, token_ids, "text", stop=[stop]))["choices"][0] + inline_text = inline["text"] + output_ids = inline["token_ids"] + finish_reason = inline["finish_reason"] + + assert inline_text == coupled["text"] + assert finish_reason == coupled["finish_reason"] + assert stop not in inline_text + + derendered = await _derender_completion_text(client, output_ids, finish_reason) + assert derendered.startswith(inline_text) + if coupled["stop_reason"] != stop: + pytest.skip("Model did not emit the stop string") + assert derendered[len(inline_text) :].startswith(stop) diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 5862256835..0e4b168677 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -50,10 +50,11 @@ from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderCompletionRequest, DerenderCompletionStreamRequest, DerenderStreamState, - GenerateResponse, - GenerateResponseChoice, - GenerateResponseStreamChoice, GenerateStreamResponse, + GenerateTokensChoice, + GenerateTokensResponse, + GenerateTokensStreamChoice, + GenerateTokensStreamResponse, ) from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo from vllm.parser import Parser @@ -195,10 +196,10 @@ def _make_stream_chunk( prompt_token_ids: list[int] | None = None, ) -> GenerateStreamResponse: """Build a GenerateStreamResponse SSE chunk.""" - return GenerateStreamResponse( + return GenerateTokensStreamResponse( request_id=request_id, choices=[ - GenerateResponseStreamChoice( + GenerateTokensStreamChoice( index=index, token_ids=token_ids, finish_reason=finish_reason, @@ -233,7 +234,7 @@ def _make_usage_chunk( request_id: str = "test-req", ) -> GenerateStreamResponse: """Build a usage only final SSE chunk (empty choices).""" - return GenerateStreamResponse( + return GenerateTokensStreamResponse( request_id=request_id, choices=[], usage=UsageInfo( @@ -326,11 +327,11 @@ def test_non_streaming_derender_initializes_parser_prefix( ) monkeypatch.setattr(derenderer, "parser", MagicMock(return_value=parser)) generated_ids = tokenizer.encode("answer", add_special_tokens=False) - response = GenerateResponse( + response = GenerateTokensResponse( request_id="test", prompt_token_ids=[11, 12], choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=generated_ids, finish_reason="length", @@ -583,10 +584,10 @@ class TestPromptSeededLeadingSpace: ): prompt_ids, output_ids, expected = leading_space_ids choices = await derenderer.derender_chat( - GenerateResponse( + GenerateTokensResponse( request_id="t", choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=output_ids, finish_reason="stop" ) ], @@ -601,10 +602,10 @@ class TestPromptSeededLeadingSpace: ): prompt_ids, output_ids, expected = leading_space_ids choices = await derenderer.derender_chat( - GenerateResponse( + GenerateTokensResponse( request_id="t", choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=output_ids, finish_reason="stop" ) ], @@ -620,10 +621,10 @@ class TestPromptSeededLeadingSpace: prompt_ids, output_ids, expected = leading_space_ids wrong_ids = tokenizer.encode("a different prompt", add_special_tokens=False) choices = await derenderer.derender_chat( - GenerateResponse( + GenerateTokensResponse( request_id="t", choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=output_ids, finish_reason="stop" ) ], @@ -639,10 +640,10 @@ class TestPromptSeededLeadingSpace: ): _, output_ids, _ = leading_space_ids choices = await derenderer.derender_chat( - GenerateResponse( + GenerateTokensResponse( request_id="t", choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=output_ids, finish_reason="stop" ) ], @@ -655,10 +656,10 @@ class TestPromptSeededLeadingSpace: prompt_ids, output_ids, expected = leading_space_ids choices, _, _ = await derenderer.derender_completion( [ - GenerateResponse( + GenerateTokensResponse( request_id="t", choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=output_ids, finish_reason="stop" ) ], @@ -1494,10 +1495,10 @@ class TestDerenderChatStreamHarmony: content = "".join(d.content or "" for d in deltas) batch_choices = await harmony_derenderer.derender_chat( - GenerateResponse( + GenerateTokensResponse( request_id="test-harmony-batch", choices=[ - GenerateResponseChoice( + GenerateTokensChoice( index=0, token_ids=output_ids, finish_reason="stop" ) ], @@ -1743,10 +1744,10 @@ class TestServingDerenderStreamValidation: could never fire since derender_chat_stream itself rejects anything above 1 first.""" serving = self._make_serving(parser_configured=False) - two_choices = GenerateStreamResponse( + two_choices = GenerateTokensStreamResponse( request_id="t", choices=[ - GenerateResponseStreamChoice(index=i, token_ids=[1]) for i in range(2) + GenerateTokensStreamChoice(index=i, token_ids=[1]) for i in range(2) ], ) request = DerenderChatStreamRequest( diff --git a/tests/entrypoints/scale_out/test_factories.py b/tests/entrypoints/scale_out/test_factories.py index 6d8335395e..ea96fff876 100644 --- a/tests/entrypoints/scale_out/test_factories.py +++ b/tests/entrypoints/scale_out/test_factories.py @@ -6,6 +6,8 @@ from argparse import Namespace from fastapi import FastAPI from vllm.entrypoints.launchers.api_server.routers import register_api_routers +from vllm.entrypoints.scale_out.token_in_token_out.api_router import attach_router +from vllm.entrypoints.serve.middleware.authenticate import GUARDED_PREFIX from vllm.tasks import SupportedTask RENDER_PATHS = { @@ -17,7 +19,9 @@ RENDER_PATHS = { "/v1/responses/render", } GENERATE_PATH = "/inference/v1/generate" -SCALE_OUT_PATHS = RENDER_PATHS | {GENERATE_PATH} +ABORT_PATH = "/inference/v1/abort_requests" +UNAUTHENTICATED_ABORT_PATH = "/abort_requests" +SCALE_OUT_PATHS = RENDER_PATHS | {GENERATE_PATH, ABORT_PATH} def registered_paths( @@ -66,4 +70,31 @@ def test_render_routes_remain_enabled_for_render_server(): def test_tokens_only_mode_enables_generate_routes_when_flag_is_unset(): paths = registered_paths(("generate",), tokens_only=True) - assert paths >= {GENERATE_PATH, "/abort_requests"} + assert paths >= {GENERATE_PATH, ABORT_PATH, UNAUTHENTICATED_ABORT_PATH} + + +def test_abort_route_without_tokens_only_requires_api_key(): + """A decode pool that loads a tokenizer serves generate without + --tokens-only and still needs to abort, but only behind --api-key.""" + paths = registered_paths(("generate",), enable_scale_out=True, tokens_only=False) + + assert ABORT_PATH in paths + assert ABORT_PATH.startswith(GUARDED_PREFIX) + assert UNAUTHENTICATED_ABORT_PATH not in paths + + +def test_unauthenticated_abort_route_is_not_registered_twice(): + """The RLHF dev router registers /abort_requests before the scale-out + routers, so attaching the scale-out one too would add a shadowed duplicate.""" + app = FastAPI() + app.state.args = Namespace(tokens_only=True) + + @app.post(UNAUTHENTICATED_ABORT_PATH) + async def existing_abort_requests(): + return None + + attach_router(app) + + paths = [route.path for route in app.routes] + assert paths.count(UNAUTHENTICATED_ABORT_PATH) == 1 + assert ABORT_PATH in paths diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py index 53ff1dc598..6b6d732f07 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py @@ -20,7 +20,8 @@ from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import ( ) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, - GenerateResponse, + GenerateTextResponse, + GenerateTokensResponse, MultiModalFeatures, PlaceholderRangeInfo, ) @@ -144,6 +145,7 @@ def _make_request_output( num_cached_tokens: int | None = None, index: int = 0, spec_decode_metrics: RequestSpecDecodeMetrics | None = None, + text: str = "", ) -> RequestOutput: return RequestOutput( request_id=request_id, @@ -153,7 +155,7 @@ def _make_request_output( outputs=[ CompletionOutput( index=index, - text="", + text=text, token_ids=token_ids, cumulative_logprob=None, logprobs=logprobs, @@ -239,7 +241,7 @@ async def test_serve_tokens_skips_mm_cache_for_remote_engine_execution(): response = await serving.serve_tokens(request) - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) assert ( serving.online_renderer.preprocess_completion.call_args.kwargs["skip_mm_cache"] is True @@ -411,7 +413,7 @@ async def test_serve_tokens_returns_spec_decode_metrics(stream: bool): if isinstance(chunk, dict) and chunk.get("metrics") )["metrics"]["speculative_decoding"] else: - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) assert response.metrics is not None assert response.metrics.speculative_decoding is not None payload = response.metrics.speculative_decoding.model_dump() @@ -511,7 +513,7 @@ async def test_serve_tokens_omits_spec_decode_metrics_for_parallel_sampling( if isinstance(chunk, dict) ) else: - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) assert response.metrics is None @@ -540,7 +542,7 @@ async def test_non_stream_returns_final_prompt_token_ids_when_requested(): response = await serving.serve_tokens(request) - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) assert response.prompt_token_ids == final_prompt_token_ids @@ -622,7 +624,7 @@ async def test_non_stream_prompt_metadata_requires_return_token_ids( response = await serving.serve_tokens(request) - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) if return_token_ids: assert response.prompt_token_ids == [1, 2, 3] assert response.model_dump()["mm_placeholders"] == MM_PLACEHOLDERS @@ -652,7 +654,7 @@ async def test_non_stream_text_only_prompt_has_no_mm_placeholders(): response = await serving.serve_tokens(request) - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) assert response.prompt_token_ids == [1, 2, 3] assert response.mm_placeholders is None @@ -678,7 +680,7 @@ async def test_non_stream_encoder_decoder_omits_mm_placeholders(): response = await serving.serve_tokens(request) - assert isinstance(response, GenerateResponse) + assert isinstance(response, GenerateTokensResponse) assert response.prompt_token_ids == [1, 2, 3] assert response.mm_placeholders is None @@ -1195,3 +1197,251 @@ async def test_stream_sampling_mask_matches_each_token_chunk(with_mask): [[token, token + 1] for token in choice["token_ids"]] if with_mask else None ) assert choice["sampling_mask"] == expected + + +async def _collect_sse(response) -> list[Any]: + return _parse_sse_chunks([chunk async for chunk in response]) + + +def _text_request(**kwargs) -> GenerateRequest: + kwargs.setdefault("sampling_params", SamplingParams(max_tokens=10)) + return GenerateRequest( + token_ids=[1, 2, 3], model=MODEL_NAME, output_mode="text", **kwargs + ) + + +@pytest.mark.asyncio +async def test_non_stream_tokens_mode_omits_text_and_echoes_output_mode(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], text="ignored", finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens( + GenerateRequest( + token_ids=[1, 2, 3], + sampling_params=SamplingParams(max_tokens=10), + model=MODEL_NAME, + ) + ) + + assert isinstance(response, GenerateTokensResponse) + dumped = response.model_dump() + assert dumped["output_mode"] == "tokens" + assert "text" not in dumped["choices"][0] + + +@pytest.mark.asyncio +async def test_non_stream_text_mode_returns_text_and_token_ids(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10, 20], + text="Hello", + finish_reason="stop", + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens(_text_request()) + + assert isinstance(response, GenerateTextResponse) + assert response.output_mode == "text" + assert response.choices[0].text == "Hello" + assert response.choices[0].token_ids == [10, 20] + assert response.choices[0].finish_reason == "stop" + assert response.usage is not None + assert response.usage.completion_tokens == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("no_tokenizer", ["tokens_only", "skip_tokenizer_init"]) +async def test_text_mode_rejected_without_tokenizer(no_tokenizer: str): + """The no-op detokenizer returns "", so text mode must fail instead + of answering 200 with empty text.""" + engine = _mock_engine() + kwargs: dict[str, Any] = {} + if no_tokenizer == "tokens_only": + kwargs["force_no_detokenize"] = True + else: + engine.model_config.skip_tokenizer_init = True + engine.renderer = _build_renderer(engine.model_config) + engine.generate = MagicMock() + serving = _build_serving_tokens(engine, **kwargs) + + response = await serving.serve_tokens(_text_request()) + + assert isinstance(response, ErrorResponse) + assert response.error.code == 400 + assert "requires a tokenizer" in response.error.message + engine.generate.assert_not_called() + + +@pytest.mark.asyncio +async def test_text_mode_logprobs_have_decoded_tokens_and_bytes(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + text="é", + logprobs=[ + { + 10: Logprob(logprob=-0.5, decoded_token="é"), + 11: Logprob(logprob=-1.5, decoded_token="e"), + } + ], + finish_reason="stop", + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens( + _text_request(sampling_params=SamplingParams(max_tokens=10, logprobs=2)) + ) + + assert isinstance(response, GenerateTextResponse) + logprobs = response.choices[0].logprobs + assert logprobs is not None and logprobs.content is not None + entry = logprobs.content[0] + assert (entry.token, entry.bytes) == ("é", [195, 169]) + assert [(t.token, t.bytes) for t in entry.top_logprobs] == [ + ("é", [195, 169]), + ("e", [101]), + ] + + +@pytest.mark.asyncio +async def test_text_mode_logprobs_keep_placeholders_with_return_tokens_as_ids(): + """--return-tokens-as-token-ids keeps token_id:N placeholders, as it does + on /v1/completions, while text is still decoded.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", + token_ids=[10], + text="é", + logprobs=[{10: Logprob(logprob=-0.5, decoded_token="é")}], + finish_reason="stop", + finished=True, + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine, return_tokens_as_token_ids=True) + + response = await serving.serve_tokens( + _text_request(sampling_params=SamplingParams(max_tokens=10, logprobs=1)) + ) + + assert isinstance(response, GenerateTextResponse) + assert response.choices[0].text == "é" + logprobs = response.choices[0].logprobs + assert logprobs is not None and logprobs.content is not None + entry = logprobs.content[0] + assert (entry.token, entry.bytes) == ("token_id:10", None) + assert [t.token for t in entry.top_logprobs] == ["token_id:10"] + + +@pytest.mark.asyncio +async def test_stream_text_mode_emits_text_deltas(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10], text="Hel") + yield _make_request_output("req-1", token_ids=[20], text="lo") + yield _make_request_output( + "req-1", token_ids=[30], text=" world", finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens(_text_request(stream=True)) + parsed = await _collect_sse(response) + + assert parsed[-1] == "[DONE]" + data_chunks = parsed[:-1] + assert {c["output_mode"] for c in data_chunks} == {"text"} + assert [c["choices"][0]["text"] for c in data_chunks] == ["Hel", "lo", " world"] + assert [c["choices"][0]["token_ids"] for c in data_chunks] == [[10], [20], [30]] + assert data_chunks[-1]["choices"][0]["finish_reason"] == "stop" + + +@pytest.mark.asyncio +async def test_stream_text_mode_emits_chunks_without_new_token_ids(): + """Text flushed from the stop string buffer and the final abort output + carry no token IDs but must still reach the client.""" + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10], text="a") + yield _make_request_output("req-1", token_ids=[], text="b") + yield _make_request_output( + "req-1", token_ids=[], text="", finish_reason="abort", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens(_text_request(stream=True)) + data_chunks = (await _collect_sse(response))[:-1] + + assert [c["choices"][0]["text"] for c in data_chunks] == ["a", "b", ""] + assert data_chunks[1]["choices"][0]["token_ids"] == [] + assert data_chunks[2]["choices"][0]["finish_reason"] == "abort" + + +@pytest.mark.asyncio +async def test_stream_text_mode_skips_outputs_without_text_or_finish(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output("req-1", token_ids=[10], text="a") + yield _make_request_output("req-1", token_ids=[], text="") + yield _make_request_output( + "req-1", token_ids=[20], text="b", finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens(_text_request(stream=True)) + data_chunks = (await _collect_sse(response))[:-1] + + assert [c["choices"][0]["text"] for c in data_chunks] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_stream_text_mode_usage_chunk_echoes_output_mode(): + engine = _mock_engine() + + async def mock_generate(*args, **kwargs): + yield _make_request_output( + "req-1", token_ids=[10], text="a", finish_reason="stop", finished=True + ) + + engine.generate = MagicMock(side_effect=mock_generate) + serving = _build_serving_tokens(engine) + + response = await serving.serve_tokens( + _text_request(stream=True, stream_options=StreamOptions(include_usage=True)) + ) + parsed = await _collect_sse(response) + + usage_chunk = parsed[-2] + assert usage_chunk["output_mode"] == "text" + assert usage_chunk["choices"] == [] + assert usage_chunk["usage"]["completion_tokens"] == 1 diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py index e51530ce45..9894d09180 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_protocol.py @@ -10,9 +10,21 @@ fail loudly if the validator semantics ever drift. import json import pytest +from pydantic import TypeAdapter, ValidationError from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( + DerenderChatRequest, + DerenderChatStreamRequest, + DerenderCompletionRequest, + DerenderCompletionStreamRequest, GenerateRequest, + GenerateResponse, + GenerateStreamResponse, + GenerateTextResponse, + GenerateTextStreamResponse, + GenerateTokensChoice, + GenerateTokensResponse, + GenerateTokensStreamResponse, MultiModalFeatures, PlaceholderRangeInfo, ) @@ -110,3 +122,120 @@ def test_generate_request_rejects_placeholder_outside_prompt(): kwargs_data={"image": ["encoded"]}, ), ) + + +def test_output_mode_defaults_to_tokens(): + assert GenerateRequest.model_validate(_base_payload()).output_mode == "tokens" + + +def test_output_mode_text_is_accepted(): + payload = {**_base_payload(), "output_mode": "text"} + assert GenerateRequest.model_validate(payload).output_mode == "text" + + +@pytest.mark.parametrize("output_mode", ["derender", "bogus", None]) +def test_output_mode_rejects_unsupported_levels(output_mode): + """Levels the server does not implement fail instead of falling back to + tokens which would answer 200 with the wrong response shape.""" + payload = {**_base_payload(), "output_mode": output_mode} + with pytest.raises(ValidationError, match="output_mode"): + GenerateRequest.model_validate(payload) + + +def test_text_mode_rejects_detokenize_false(): + payload = { + "token_ids": [1, 2, 3], + "sampling_params": {"detokenize": False}, + "output_mode": "text", + } + with pytest.raises(ValidationError, match="detokenize"): + GenerateRequest.model_validate(payload) + + +def test_tokens_mode_allows_detokenize_false(): + payload = { + "token_ids": [1, 2, 3], + "sampling_params": {"detokenize": False}, + "output_mode": "tokens", + } + assert GenerateRequest.model_validate(payload).output_mode == "tokens" + + +def test_response_without_output_mode_parses_as_tokens(): + """Existing clients and older servers never send output_mode.""" + parsed = TypeAdapter(GenerateResponse).validate_python( + {"choices": [{"index": 0, "token_ids": [1]}]} + ) + assert isinstance(parsed, GenerateTokensResponse) + assert parsed.output_mode == "tokens" + + +def test_stream_chunk_without_output_mode_parses_as_tokens(): + parsed = TypeAdapter(GenerateStreamResponse).validate_python( + {"choices": [{"index": 0, "token_ids": [1]}]} + ) + assert isinstance(parsed, GenerateTokensStreamResponse) + + +def test_text_response_and_chunk_dispatch_on_output_mode(): + response = TypeAdapter(GenerateResponse).validate_python( + {"output_mode": "text", "choices": [{"index": 0, "text": "hi"}]} + ) + chunk = TypeAdapter(GenerateStreamResponse).validate_python( + {"output_mode": "text", "choices": [{"index": 0, "text": "hi"}]} + ) + assert isinstance(response, GenerateTextResponse) + assert isinstance(chunk, GenerateTextStreamResponse) + + +def test_text_response_requires_text_on_every_choice(): + with pytest.raises(ValidationError, match="text"): + TypeAdapter(GenerateResponse).validate_python( + {"output_mode": "text", "choices": [{"index": 0}]} + ) + + +@pytest.mark.parametrize("output_mode", ["derender", None, 5]) +def test_response_rejects_unknown_output_mode(output_mode): + with pytest.raises(ValidationError): + TypeAdapter(GenerateResponse).validate_python( + {"output_mode": output_mode, "choices": []} + ) + + +def test_derender_requests_accept_responses_without_output_mode(): + DerenderChatRequest.model_validate( + {"generate_response": {"choices": [{"index": 0, "token_ids": [1]}]}} + ) + DerenderChatStreamRequest.model_validate( + { + "stream": True, + "generate_chunk": {"choices": [{"index": 0, "token_ids": [1]}]}, + } + ) + + +def test_derender_requests_reject_text_responses(): + """Text responses are already detokenized with stop strings applied, so + re-decoding their token_ids would silently undo the truncation.""" + text = {"output_mode": "text", "choices": [{"index": 0, "text": "hi"}]} + with pytest.raises(ValidationError, match="output_mode"): + DerenderChatRequest.model_validate({"generate_response": text}) + with pytest.raises(ValidationError, match="output_mode"): + DerenderCompletionRequest.model_validate({"generate_responses": [text]}) + with pytest.raises(ValidationError, match="output_mode"): + DerenderChatStreamRequest.model_validate( + {"stream": True, "generate_chunk": text} + ) + with pytest.raises(ValidationError, match="output_mode"): + DerenderCompletionStreamRequest.model_validate( + {"stream": True, "generate_chunk": text} + ) + + +def test_tokens_response_echoes_output_mode_without_text(): + dumped = GenerateTokensResponse( + choices=[GenerateTokensChoice(index=0, token_ids=[1])] + ).model_dump() + assert dumped["output_mode"] == "tokens" + assert "text" not in dumped["choices"][0] diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index 18513f83dc..2921422f01 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -19,6 +19,7 @@ from vllm.v1.engine.detokenizer import check_stop_strings MODEL_NAME = "Qwen/Qwen3-0.6B" GEN_ENDPOINT = "/inference/v1/generate" +ABORT_ENDPOINT = "/inference/v1/abort_requests" def get_vocab_size(model_name): @@ -317,11 +318,11 @@ async def test_generate_logprobs(client, logprobs_value): @pytest.mark.asyncio async def test_generate_prompt_token_id_logprobs(client): token_ids = [11, 22, 33, 44, 55] - candidates = [22, 33, 44, 55] + rows = [[33], [44, 100], [55]] sampling_params = { "max_tokens": 1, "prompt_logprobs": 0, - "prompt_logprob_token_ids": candidates, + "prompt_logprob_token_ids": rows, "prompt_logprob_start": 1, } payload = { @@ -334,14 +335,13 @@ async def test_generate_prompt_token_id_logprobs(client): data = resp.json() scores = np.load(io.BytesIO(base64.b64decode(data["prompt_token_id_logprobs"]))) - assert scores.shape == (len(token_ids) - 2, len(candidates)) + assert scores.shape == (3, 2) for row, target in enumerate(token_ids[2:]): expected = data["prompt_logprobs"][row + 2][str(target)]["logprob"] - assert scores[row, candidates.index(target)] == pytest.approx( - expected, abs=1e-3 - ) + assert scores[row, 0] == pytest.approx(expected, abs=1e-3) + assert np.isneginf(scores[[0, 2], 1]).all() - stream_only_ids = {"prompt_logprob_token_ids": candidates} + stream_only_ids = {"prompt_logprob_token_ids": rows} resp = await client.post( GEN_ENDPOINT, json={**payload, "sampling_params": stream_only_ids, "stream": True}, @@ -349,15 +349,12 @@ async def test_generate_prompt_token_id_logprobs(client): assert resp.status_code == 400 assert "prompt_logprob_token_ids" in resp.text - past_end = {**sampling_params, "prompt_logprob_start": len(token_ids)} + too_few_rows = {**sampling_params, "prompt_logprob_token_ids": rows[:2]} resp = await client.post( - GEN_ENDPOINT, json={**payload, "sampling_params": past_end} + GEN_ENDPOINT, json={**payload, "sampling_params": too_few_rows} ) - resp.raise_for_status() - empty = np.load( - io.BytesIO(base64.b64decode(resp.json()["prompt_token_id_logprobs"])) - ) - assert empty.shape == (0, len(candidates)) + assert resp.status_code == 400 + assert "scored rows" in resp.text @pytest.mark.asyncio @@ -590,3 +587,187 @@ async def test_generate_with_lora_adapter(client, tokenizer, messages): completions_res = completions_data["choices"][0]["message"]["content"] assert generate_res == completions_res + + +def _chat_prompt_token_ids(tokenizer, messages) -> list[int]: + return tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + enable_thinking=False, # default with Qwen3 + return_dict=True, # default with Transformers v5 + ).input_ids + + +def _text_mode_payload(token_ids, **sampling_params) -> dict: + return { + "model": MODEL_NAME, + "token_ids": token_ids, + "sampling_params": {"max_tokens": 24, "temperature": 0.0, **sampling_params}, + "output_mode": "text", + "stream": False, + } + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="output_mode is not supported by the Rust frontend", +) +async def test_text_mode_matches_chat_completions(client, tokenizer, messages): + token_ids = _chat_prompt_token_ids(tokenizer, messages) + + resp = await client.post(GEN_ENDPOINT, json=_text_mode_payload(token_ids)) + resp.raise_for_status() + data = resp.json() + assert data["output_mode"] == "text" + + payload = { + "model": MODEL_NAME, + "messages": messages, + "max_tokens": 24, + "temperature": 0.0, + "stream": False, + "chat_template_kwargs": {"enable_thinking": False}, + } + chat_resp = await client.post("/v1/chat/completions", json=payload) + chat_content = chat_resp.json()["choices"][0]["message"]["content"] + + assert data["choices"][0]["text"] == chat_content + assert data["choices"][0]["token_ids"] + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="output_mode is not supported by the Rust frontend", +) +async def test_text_mode_stream_matches_non_stream(client, tokenizer, messages): + token_ids = _chat_prompt_token_ids(tokenizer, messages) + payload = _text_mode_payload(token_ids) + + non_stream = await client.post(GEN_ENDPOINT, json=payload) + non_stream.raise_for_status() + + chunks = [] + async with client.stream( + "POST", GEN_ENDPOINT, json={**payload, "stream": True} + ) as resp: + resp.raise_for_status() + async for line in resp.aiter_lines(): + if line.startswith("data: ") and line != "data: [DONE]": + chunks.append(json.loads(line[len("data: ") :])) + + assert {chunk["output_mode"] for chunk in chunks} == {"text"} + assert chunks[-1]["choices"][0]["finish_reason"] is not None + streamed_text = "".join(chunk["choices"][0]["text"] for chunk in chunks) + assert streamed_text == non_stream.json()["choices"][0]["text"] + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="output_mode is not supported by the Rust frontend", +) +@pytest.mark.parametrize( + "overrides, expected", + [ + pytest.param( + {"sampling_params": {"detokenize": False}}, "detokenize", id="no-detok" + ), + pytest.param({"output_mode": "derender"}, "output_mode", id="unsupported"), + ], +) +async def test_text_mode_invalid_requests_return_400(client, overrides, expected): + payload = { + **_text_mode_payload([1, 2, 3], max_tokens=5), + **overrides, + } + resp = await client.post(GEN_ENDPOINT, json=payload) + + assert resp.status_code == 400 + assert expected in resp.json()["error"]["message"] + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="--tokens-only is not supported by the Rust frontend", +) +@pytest.mark.parametrize("server", [["--tokens-only"]], indirect=True) +async def test_text_mode_rejected_when_tokens_only(client): + """Without a tokenizer the detokenizer returns "", so the request must + fail instead of answering 200 with empty text.""" + resp = await client.post(GEN_ENDPOINT, json=_text_mode_payload([1, 2, 3])) + assert resp.status_code == 400 + assert "requires a tokenizer" in resp.json()["error"]["message"] + + payload = {**_text_mode_payload([1, 2, 3]), "output_mode": "tokens"} + payload["sampling_params"]["detokenize"] = False + resp = await client.post(GEN_ENDPOINT, json=payload) + resp.raise_for_status() + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="output_mode is not supported by the Rust frontend", +) +async def test_text_mode_stream_delivers_abort_finish_chunk( + client, tokenizer, messages +): + """The final output after an abort has no new token IDs, so it + reaches a text stream only because text mode emits finish-only chunks.""" + payload = { + **_text_mode_payload( + _chat_prompt_token_ids(tokenizer, messages), + max_tokens=900, + ignore_eos=True, + ), + "stream": True, + "request_id": "text-abort-e2e", + } + + chunks = [] + aborted = False + async with client.stream("POST", GEN_ENDPOINT, json=payload) as resp: + resp.raise_for_status() + async for line in resp.aiter_lines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + chunk = json.loads(line[len("data: ") :]) + chunks.append(chunk) + if not aborted and chunk["choices"] and chunk["choices"][0]["text"]: + abort = await client.post( + ABORT_ENDPOINT, + json={"request_ids": [chunk["request_id"]]}, + ) + assert abort.status_code == 200 + aborted = True + + assert aborted + assert {chunk["output_mode"] for chunk in chunks} == {"text"} + final_choice = chunks[-1]["choices"][0] + assert final_choice["finish_reason"] == "abort" + generated = sum(len(chunk["choices"][0]["token_ids"] or []) for chunk in chunks) + assert generated < 900 + + +@pytest.mark.asyncio +@pytest.mark.skipif( + envs.VLLM_USE_RUST_FRONTEND, + reason="the Rust frontend does not serve /inference/v1/abort_requests", +) +async def test_abort_requests_is_served_without_tokens_only(client): + resp = await client.post( + ABORT_ENDPOINT, json={"request_ids": ["generate-tokens-unknown"]} + ) + assert resp.status_code == 200 + + resp = await client.post(ABORT_ENDPOINT, json={}) + assert resp.status_code == 400 + + # The unauthenticated path is only served with --tokens-only. + resp = await client.post( + "/abort_requests", json={"request_ids": ["generate-tokens-unknown"]} + ) + assert resp.status_code == 404 diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py index f42a406f41..c417a5fc2b 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens from vllm.logprobs import Logprob @@ -36,3 +38,50 @@ def test_logprobs_minus_one_emits_all_tokens(): num_output_top_logprobs=-1, ) assert len(result.content[0].top_logprobs) == 2 + + +def test_tokens_level_logprobs_have_no_bytes(): + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[7], + top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}], + num_output_top_logprobs=2, + ) + entry = result.content[0] + assert entry.bytes is None + assert all(t.bytes is None for t in entry.top_logprobs) + + +def test_resolved_logprobs_use_engine_decoded_tokens_and_bytes(): + tokenizer = MagicMock() + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[7], + top_logprobs=[ + {7: Logprob(-0.9, decoded_token="é"), 8: Logprob(-1.1, decoded_token="e")} + ], + num_output_top_logprobs=2, + tokenizer=tokenizer, + ) + entry = result.content[0] + assert (entry.token, entry.bytes) == ("é", [195, 169]) + assert [(t.token, t.bytes) for t in entry.top_logprobs] == [ + ("é", [195, 169]), + ("e", [101]), + ] + tokenizer.decode.assert_not_called() + + +def test_resolved_logprobs_decode_tokens_the_engine_did_not(): + tokenizer = MagicMock() + tokenizer.decode.side_effect = lambda ids: f"<{ids[0]}>" + result = ServingTokens._create_tokens_logprobs( + None, + token_ids=[7, 9], + top_logprobs=[{7: Logprob(-0.9)}, None], + num_output_top_logprobs=1, + tokenizer=tokenizer, + ) + assert [e.token for e in result.content] == ["<7>", "<9>"] + assert result.content[0].top_logprobs[0].token == "<7>" + assert result.content[1].bytes == list(b"<9>") diff --git a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py index 7fb7a5522d..c74246b2f8 100644 --- a/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py +++ b/tests/entrypoints/speech_to_text/realtime/test_realtime_validation.py @@ -14,12 +14,6 @@ from tests.utils import ROCM_EXTRA_ARGS, RemoteOpenAIServer from vllm.assets.audio import AudioAsset from vllm.multimodal.media.audio import load_audio -# Increase engine iteration timeout for ROCm where first-use JIT compilation -# can exceed the default 60s, causing a silent deadlock in feed_tokens. -REALTIME_ENV_OVERRIDES = { - "VLLM_ENGINE_ITERATION_TIMEOUT_S": "600", -} - MISTRAL_FORMAT_ARGS = [ "--tokenizer_mode", "mistral", @@ -77,9 +71,7 @@ async def test_multi_chunk_streaming(model_name, mary_had_lamb_audio_chunks): if model_name.startswith("mistralai"): server_args += MISTRAL_FORMAT_ARGS - with RemoteOpenAIServer( - model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES - ) as remote_server: + with RemoteOpenAIServer(model_name, server_args) as remote_server: ws_url = _get_websocket_url(remote_server) async with websockets.connect(ws_url) as ws: # Receive session.created @@ -183,9 +175,7 @@ async def test_empty_commit_does_not_crash_engine( if model_name.startswith("mistralai"): server_args += MISTRAL_FORMAT_ARGS - with RemoteOpenAIServer( - model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES - ) as remote_server: + with RemoteOpenAIServer(model_name, server_args) as remote_server: ws_url = _get_websocket_url(remote_server) # --- First connection: empty commit (no audio appended) ---------- @@ -272,9 +262,7 @@ async def test_session_update_invalid_model_returns_error(model_name): if model_name.startswith("mistralai"): server_args += MISTRAL_FORMAT_ARGS - with RemoteOpenAIServer( - model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES - ) as remote_server: + with RemoteOpenAIServer(model_name, server_args) as remote_server: ws_url = _get_websocket_url(remote_server) async with websockets.connect(ws_url) as ws: event = await receive_event(ws, timeout=30.0) @@ -301,9 +289,7 @@ async def test_commit_without_session_update_returns_error(model_name): if model_name.startswith("mistralai"): server_args += MISTRAL_FORMAT_ARGS - with RemoteOpenAIServer( - model_name, server_args, env_dict=REALTIME_ENV_OVERRIDES - ) as remote_server: + with RemoteOpenAIServer(model_name, server_args) as remote_server: ws_url = _get_websocket_url(remote_server) async with websockets.connect(ws_url) as ws: event = await receive_event(ws, timeout=30.0) diff --git a/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py b/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py index 247faba17f..dcf7a6bc03 100644 --- a/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py +++ b/tests/kernels/attention/test_minimax_m3_fp8_triton_indexer.py @@ -107,3 +107,37 @@ def test_fp8_index_query_and_cache_prefill_topk_matches_bf16() -> None: return {int(b) for b in idx[0, 0].tolist() if b >= 0} _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn))) + + +@pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required") +@torch.inference_mode() +def test_bf16_query_against_fp8_cache_decode_topk_matches_bf16() -> None: + """A mixed pair must upcast rather than emit a mixed-dtype fp8 dot. + + Production allocates index_q with the index-K cache dtype, so the scorers + normally see a matched pair and keep the native FP8 MMA. This pins the + other branch: a bf16 query against an e4m3 cache still compiles and still + ranks the planted blocks. + """ + dev = "cuda" + q, keys = _make_inputs(dev) + block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0) + seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev) + + def select(cache: torch.Tensor) -> set[int]: + idx = minimax_m3_index_decode( + q, # bf16 query, deliberately not cast to the cache dtype + cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(), + block_table, + seq_lens, + SEQ_LEN, + TOPK, + 0, + 0, + HEADS, + 1, + 1, + ) + return {int(b) for b in idx[0, 0].tolist() if b >= 0} + + _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn))) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 387e690e1b..3b200be514 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -596,6 +596,114 @@ def test_sparse_attn_prefill_ragged_kernel() -> None: torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) +def _sparse_prefill_ragged_inputs(nope_dim: int, rope_dim: int) -> dict: + device = torch.device("cuda") + set_random_seed(7) + head_dim = nope_dim + rope_dim + return dict( + q=torch.randn(3, 3, head_dim, dtype=torch.bfloat16, device=device) * 0.125, + kv=torch.randn(5, head_dim, dtype=torch.bfloat16, device=device) * 0.125, + indices=torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device), + indptr=torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device), + scale=head_dim**-0.5, + attn_sink=torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device), + nope_head_dim=nope_dim, + rope_head_dim=rope_dim, + ) + + +# 448+64 is DeepSeek V4/V4.1, 512+64 is V3.2, and 256+0 stands in for the +# NoPE-only layouts where the destination spans the whole head dim. +SPARSE_PREFILL_DIMS = [(NOPE_HEAD_DIM, ROPE_HEAD_DIM), (512, 64), (256, 0)] + + +@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS) +@torch.inference_mode() +def test_sparse_attn_prefill_ragged_out_matches_allocated( + nope_dim: int, rope_dim: int +) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_prefill_ragged_triton, + ) + + kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim) + allocated = _rocm_sparse_attn_prefill_ragged_triton(**kwargs) + + q = kwargs["q"] + dest = torch.empty(q.shape[0], q.shape[1], nope_dim, dtype=q.dtype, device=q.device) + returned = _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest) + + assert returned is dest + torch.testing.assert_close(dest, allocated[..., :nope_dim], atol=0.0, rtol=0.0) + + +@pytest.mark.parametrize(("nope_dim", "rope_dim"), SPARSE_PREFILL_DIMS) +@torch.inference_mode() +def test_sparse_attn_prefill_ragged_out_view_leaves_neighbors_untouched( + nope_dim: int, rope_dim: int +) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_prefill_ragged_triton, + ) + + kwargs = _sparse_prefill_ragged_inputs(nope_dim, rope_dim) + q = kwargs["q"] + num_queries, num_heads = q.shape[0], q.shape[1] + + buffer = torch.full( + (num_queries + 4, num_heads + 2, nope_dim), + -7.0, + dtype=q.dtype, + device=q.device, + ) + dest = buffer[2 : 2 + num_queries, :num_heads] + assert not dest.is_contiguous() + _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest) + + expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :nope_dim] + torch.testing.assert_close(dest, expected, atol=0.0, rtol=0.0) + + untouched = torch.ones_like(buffer, dtype=torch.bool) + untouched[2 : 2 + num_queries, :num_heads] = False + assert torch.all(buffer[untouched] == -7.0) + + +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.float16]) +@torch.inference_mode() +def test_sparse_attn_prefill_ragged_out_casts_to_dest_dtype( + out_dtype: torch.dtype, +) -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_prefill_ragged_triton, + ) + + kwargs = _sparse_prefill_ragged_inputs(512, 64) + q = kwargs["q"] + dest = torch.empty(q.shape[0], q.shape[1], 512, dtype=out_dtype, device=q.device) + _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest) + + expected = _rocm_sparse_attn_prefill_ragged_triton(**kwargs)[..., :512] + assert dest.dtype == out_dtype + torch.testing.assert_close(dest.float(), expected.float(), atol=2e-2, rtol=2e-2) + + +@torch.inference_mode() +def test_sparse_attn_prefill_head_dim_wide_out_keeps_rope_rows() -> None: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + _rocm_sparse_attn_prefill_ragged_triton, + ) + + kwargs = _sparse_prefill_ragged_inputs(NOPE_HEAD_DIM, ROPE_HEAD_DIM) + q = kwargs["q"] + dest = torch.empty_like(q) + _rocm_sparse_attn_prefill_ragged_triton(**kwargs, out=dest) + + expected = _ref_sparse_prefill_ragged( + q, kwargs["kv"], [[0, 2], [1, 3, 4], []], kwargs["scale"], kwargs["attn_sink"] + ) + torch.testing.assert_close(dest, expected, atol=2e-2, rtol=2e-2) + + @pytest.mark.parametrize( ("num_queries", "on_gfx950", "expected"), [(1023, True, False), (1024, True, True), (1024, False, False)], @@ -718,10 +826,11 @@ def test_sparse_attn_prefill_preserves_dense_triton_fallback(monkeypatch) -> Non output = torch.empty_like(q) dense_fallback_calls = 0 - def fake_dense_fallback(*args, **kwargs): + def fake_dense_fallback(*args, out, **kwargs): nonlocal dense_fallback_calls dense_fallback_calls += 1 - return torch.zeros_like(q) + out.zero_() + return out monkeypatch.setattr(mod, "_can_use_aiter_sparse_prefill_opus", lambda *args: True) monkeypatch.setattr(mod, "_get_aiter_sparse_prefill_opus", lambda: None) diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py index ea062b50c2..b6a101be07 100644 --- a/tests/kernels/mamba/test_gdn_fused_mtp.py +++ b/tests/kernels/mamba/test_gdn_fused_mtp.py @@ -194,7 +194,7 @@ def test_fused_forward_uses_packed_entrypoint() -> None: ba = torch.randn(num_tokens, 2 * HV, dtype=torch.bfloat16, device=device) layer = types.SimpleNamespace( prefix=PREFIX, - enable_fused_gdn_decode=True, + enable_fused_gdn_spec_decode=True, norm=types.SimpleNamespace( weight=torch.empty(V, dtype=torch.bfloat16, device=device) ), diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 0554e9ca66..00e9834d62 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -167,7 +167,7 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch): monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams) draft = SimpleNamespace( - use_mega_moe=False, + use_native_mega_moe=False, use_sequence_parallel=False, hc_mult=hc_mult, layers=[make_layer(mix) for mix in mixes], diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 1944aaac57..9ae6092c94 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -88,6 +88,15 @@ class Config: torch_trace_dir_path: str | None = None + # Force AiterExperts's hidden_pad/intermediate_pad computation + # (`experts/rocm_aiter_moe.py`) to diverge from the padded K/N sizes above. + # None (default) preserves today's behavior: FusedMoEConfig defaults both + # to the (unpadded) K/intermediate_size_per_partition, so hidden_pad and + # intermediate_pad come out to 0. + # See https://github.com/vllm-project/vllm/issues/54966 ("Test padding"). + hidden_dim_unpadded: int | None = None + intermediate_size_per_partition_unpadded: int | None = None + def __post_init__(self): if self.quant_config is None: self.quant_config = TestMoEQuantConfig(None, False, False, None) @@ -664,6 +673,10 @@ def make_modular_kernel( activation=config.activation, device=vllm_config.device_config.device, routing_method=RoutingMethodType.DeepSeekV3, + hidden_dim_unpadded=config.hidden_dim_unpadded, + intermediate_size_per_partition_unpadded=( + config.intermediate_size_per_partition_unpadded + ), ) prepare_finalize = maybe_make_prepare_finalize( diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 2a46a0a56a..629759de40 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -295,6 +295,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( expert_map=None, apply_router_weight_on_input=True, routed_scaling_factor=1.0, + routing_sink=None, ) check_accuracy( @@ -968,6 +969,7 @@ def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_in router_logits=logits, e_score_correction_bias=bias, routed_scaling_factor=scale, + routing_sink=None, **common, ) topk_weights, topk_ids = fused_topk_bias( diff --git a/tests/kernels/moe/test_flashinfer_moe_ep.py b/tests/kernels/moe/test_flashinfer_moe_ep.py new file mode 100644 index 0000000000..c3223b7d98 --- /dev/null +++ b/tests/kernels/moe/test_flashinfer_moe_ep.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the FlashInfer MoE-EP adapter on real Blackwell GPUs. + +Each test spawns one process per GPU, bootstraps vLLM's expert-parallel +groups, builds a ``FlashInferMoeEp`` adapter from bf16 expert weights at the +MoE geometry of a real model, warms it up and runs a forward whose inputs are +chosen so that every quantization step in the megakernel is exact, for the +CuTeDSL kernel (NVFP4 activations and weights, E4M3 scales per 16) as well as +the DeepGEMM kernel (FP8 activations and FP4 weights, power-of-two scales per +32). The expected output is then known in closed form and compared bit for +bit. +""" + +from dataclasses import dataclass + +import pytest +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.distributed.eplb_utils import distributed_run, set_env_vars_and_device +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM +from vllm.distributed.parallel_state import ensure_model_parallel_initialized +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( + FlashInferMoeEpExperts, +) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + FlashInferMoeEp, + FlashInferMoeEpWeights, + apply_topk_in_fc1, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Static +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_deep_gemm + + +@dataclass(frozen=True) +class Geometry: + name: str + hidden_size: int + intermediate_size: int + num_experts: int + top_k: int + + +# Routed-expert geometry of the models the branch was evaluated on. DeepGEMM +# keeps one UE8M0 scale byte per 32 elements and requires the per-token scale +# rows to be TMA aligned, so hidden and intermediate sizes must be multiples +# of 512 (FlashInfer's own validation only checks 128); all four qualify. +GEOMETRIES = ( + Geometry("dsv4-pro", 7168, 3072, 384, 6), + Geometry("dsv4-flash", 4096, 2048, 256, 6), + Geometry("mistral-small-4", 4096, 2048, 128, 4), + Geometry("mistral-large-3", 7168, 4096, 128, 4), +) +BACKENDS = (FLASHINFER_MOE_EP_CUTEDSL, FLASHINFER_MOE_EP_DEEP_GEMM) + +NUM_TOKENS_PER_RANK = 4 +MAX_TOKENS_PER_RANK = 128 + +# Every block scale below is either 1 or a power of two, so the FP4 codes +# and the FP8 values survive quantization exactly: NVFP4 scales a 16-block by +# absmax / 6 in E4M3, DeepGEMM scales a 32-block by absmax / 6 (FP4) or +# absmax / 448 (FP8) rounded up to a power of two. A 6.0 in a block pins the +# FP4 scale to 1; 2.0, 6.0 and 144 are exact FP8 values at any such scale. +# +# Token t carries X_SIGNAL at hidden column 16 * t and a 6.0 sentinel next to +# it; every weight column the sentinel meets is zero. +X_SIGNAL = 2.0 +X_SENTINEL = 6.0 +# Gate and up rows are identical, so h = silu(12) * 12 = 144 = 6 * 24 +# whichever half the kernel treats as the gate, and 144 requantizes exactly. +W13_DIAG = 6.0 +H_SIGNAL = 144.0 +# Expert e writes E2M1_CODES[e % 7] * H_SIGNAL into output channel e. The 6.0 +# sentinel sits next to a zero intermediate channel and pins the block scale. +W2_SENTINEL = 6.0 +E2M1_CODES = (0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _expert_out(expert: int) -> float: + return E2M1_CODES[expert % len(E2M1_CODES)] + + +def _topk_weights(top_k: int) -> tuple[float, ...]: + # The DeepSeek-V4 routing policy applies the router weight before fc2, so + # weight * 144 must requantize exactly too: powers of two do. + return tuple(2.0 ** -(k + 1) for k in range(top_k)) + + +def _routing_method(backend: str) -> RoutingMethodType: + if backend == FLASHINFER_MOE_EP_DEEP_GEMM: + return RoutingMethodType.DeepseekV4 + return RoutingMethodType.Default + + +def _routing( + geometry: Geometry, rank: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """top_k distinct experts per token, strided across ranks so dispatch goes + remote.""" + base = ( + torch.arange(NUM_TOKENS_PER_RANK, device=device) + rank * NUM_TOKENS_PER_RANK + ) % geometry.num_experts + stride = geometry.num_experts // geometry.top_k + offsets = torch.arange(geometry.top_k, device=device) * stride + topk_ids = ((base[:, None] + offsets[None, :]) % geometry.num_experts).to( + torch.int32 + ) + topk_weights = torch.tensor( + [_topk_weights(geometry.top_k)] * NUM_TOKENS_PER_RANK, + dtype=torch.float32, + device=device, + ) + return topk_ids, topk_weights + + +def _hidden_states(geometry: Geometry, device: torch.device) -> torch.Tensor: + x = torch.zeros( + NUM_TOKENS_PER_RANK, geometry.hidden_size, dtype=torch.bfloat16, device=device + ) + for t in range(NUM_TOKENS_PER_RANK): + x[t, 16 * t] = X_SIGNAL + x[t, 16 * t + 1] = X_SENTINEL + return x + + +def _local_expert_weights( + geometry: Geometry, rank: int, num_local_experts: int, device: torch.device +) -> FlashInferMoeEpWeights: + w13 = torch.zeros( + num_local_experts, + 2 * geometry.intermediate_size, + geometry.hidden_size, + dtype=torch.bfloat16, + device=device, + ) + w2 = torch.zeros( + num_local_experts, + geometry.hidden_size, + geometry.intermediate_size, + dtype=torch.bfloat16, + device=device, + ) + token_columns = [16 * t for t in range(NUM_TOKENS_PER_RANK)] + for local in range(num_local_experts): + expert = rank * num_local_experts + local + w13[local, 0, token_columns] = W13_DIAG + w13[local, geometry.intermediate_size, token_columns] = W13_DIAG + w2[local, expert, 0] = _expert_out(expert) + w2[local, expert, 2] = W2_SENTINEL + return FlashInferMoeEpWeights(w13=w13, w2=w2) + + +def _expected_output( + geometry: Geometry, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + expected = torch.zeros( + NUM_TOKENS_PER_RANK, + geometry.hidden_size, + dtype=torch.float32, + device=topk_ids.device, + ) + for t in range(NUM_TOKENS_PER_RANK): + for k in range(geometry.top_k): + expert = int(topk_ids[t, k]) + weight = float(topk_weights[t, k]) + expected[t, expert] += weight * _expert_out(expert) * H_SIGNAL + return expected.to(torch.bfloat16) + + +def _moe_config( + vllm_config: VllmConfig, + geometry: Geometry, + world_size: int, + device: torch.device, + backend: str, +) -> FusedMoEConfig: + parallel = FusedMoEParallelConfig.make( + tp_size_=1, + pcp_size_=1, + dp_size_=world_size, + sp_size_=1, + vllm_parallel_config=vllm_config.parallel_config, + ) + return FusedMoEConfig( + num_experts=geometry.num_experts, + experts_per_token=geometry.top_k, + hidden_dim=geometry.hidden_size, + intermediate_size=geometry.intermediate_size, + num_local_experts=geometry.num_experts // world_size, + num_logical_experts=geometry.num_experts, + activation=MoEActivation.SILU, + device=device, + routing_method=_routing_method(backend), + moe_parallel_config=parallel, + in_dtype=torch.bfloat16, + moe_backend=backend, + max_num_tokens=MAX_TOKENS_PER_RANK, + ) + + +def _check_expert_views_alias_kernel_memory( + adapter: FlashInferMoeEp, + geometry: Geometry, + rank: int, + num_local_experts: int, + device: torch.device, +) -> None: + """EPLB moves experts through ``expert_weight_views``; the kernel must see it. + + Swap local experts 0 and 1 in place on every rank, then route through the + swapped physical slots: the output must be the one the unswapped logical + experts would have produced. + """ + for view in adapter.expert_weight_views(): + first = view[0].clone() + view[0].copy_(view[1]) + view[1].copy_(first) + torch.distributed.barrier() + + topk_ids, topk_weights = _routing(geometry, rank, device) + slot = topk_ids % num_local_experts + physical_ids = torch.where(slot == 0, topk_ids + 1, topk_ids) + physical_ids = torch.where(slot == 1, topk_ids - 1, physical_ids) + output = adapter(_hidden_states(geometry, device), physical_ids, topk_weights) + torch.accelerator.synchronize() + assert torch.equal(output, _expected_output(geometry, topk_ids, topk_weights)) + + +def _init_warmup_forward( + env: dict[str, str], world_size: int, geometry: Geometry, backend: str +) -> None: + set_env_vars_and_device(env) + vllm_config = VllmConfig() + vllm_config.parallel_config.data_parallel_size = world_size + vllm_config.parallel_config.enable_expert_parallel = True + vllm_config.kernel_config.moe_backend = backend + + with set_current_vllm_config(vllm_config): + ensure_model_parallel_initialized( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank}") + moe = _moe_config(vllm_config, geometry, world_size, device, backend) + supported, reason = FlashInferMoeEpExperts.is_supported_config( + FlashInferMoeEpExperts, + moe, + kMxfp4Static, + None, + mk.FusedMoEActivationFormat.Standard, + ) + assert supported, reason + + adapter = FlashInferMoeEp( + moe, + _local_expert_weights(geometry, rank, moe.num_local_experts, device), + apply_topk_in_fc1=apply_topk_in_fc1(moe), + ) + try: + adapter.warmup() + + topk_ids, topk_weights = _routing(geometry, rank, device) + output = adapter(_hidden_states(geometry, device), topk_ids, topk_weights) + torch.accelerator.synchronize() + + expected = _expected_output(geometry, topk_ids, topk_weights) + assert output.dtype == expected.dtype and output.shape == expected.shape + assert torch.equal(output, expected), ( + f"rank {rank}: max abs diff " + f"{(output.float() - expected.float()).abs().max().item()}, " + f"nonzero output channels per token " + f"{[torch.nonzero(row).flatten().tolist() for row in output]}" + ) + + if backend == FLASHINFER_MOE_EP_CUTEDSL: + _check_expert_views_alias_kernel_memory( + adapter, geometry, rank, moe.num_local_experts, device + ) + finally: + adapter.destroy() + + +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="The FlashInfer MoE-EP megakernels require SM100", +) +@pytest.mark.parametrize("geometry", GEOMETRIES, ids=[g.name for g in GEOMETRIES]) +@pytest.mark.parametrize("backend", BACKENDS, ids=("cutedsl", "deep_gemm")) +@pytest.mark.parametrize("world_size", [2, 4]) +def test_init_warmup_forward_is_exact_on_grid_inputs( + world_size: int, backend: str, geometry: Geometry +): + pytest.importorskip("flashinfer.moe_ep") + if backend == FLASHINFER_MOE_EP_DEEP_GEMM and not has_deep_gemm(): + pytest.skip("DeepGEMM is not available") + if torch.accelerator.device_count() < world_size: + pytest.skip(f"Need at least {world_size} GPUs") + distributed_run(_init_warmup_forward, world_size, geometry, backend) diff --git a/tests/kernels/moe/test_modular_kernel_combinations.py b/tests/kernels/moe/test_modular_kernel_combinations.py index 253abc5c38..4b769931ea 100644 --- a/tests/kernels/moe/test_modular_kernel_combinations.py +++ b/tests/kernels/moe/test_modular_kernel_combinations.py @@ -42,7 +42,7 @@ from .modular_kernel_tools.parallel_utils import ( ProcessGroupInfo, parallel_launch_with_config, ) -from .utils import check_accuracy +from .utils import check_accuracy, make_test_weights has_any_multi_gpu_package = ( has_deep_ep() or has_deep_gemm() or has_flashinfer_cutlass_fused_moe() @@ -691,6 +691,303 @@ def test_aiter_moe_dispatch_policy_forwarded_through_apply(dispatch_policy: int) ) +# --- 4b: hidden_dim_unpadded/intermediate_size_per_partition_unpadded matrix - + +# K/N intentionally unaligned to AITER's 64/128 granularity (unlike the +# shared mk_objects.py defaults) so hidden_pad/intermediate_pad are forced. +_PADDING_E = 8 +_PADDING_M = 16 +_PADDING_TOPK = 2 +_PADDING_K_UNPADDED = 896 # 7 * 128 +_PADDING_K_PADDED = 1024 # +128 +_PADDING_N_UNPADDED = 384 # 3 * 128 +_PADDING_N_PADDED = 512 # +128 +# hidden_pad = (K_padded - hidden_dim_unpadded) // 128 * 128 (rocm_aiter_moe.py) +_PADDING_HIDDEN_PAD_EXPECTED = 128 +# intermediate_pad = (N_padded - intermediate_size_per_partition_unpadded) +# // 64 * 64 * (2 if tp_size == 1 else 1) (rocm_aiter_moe.py); tp_size==1 here. +_PADDING_INTERMEDIATE_PAD_EXPECTED = 256 + +# {no padding, hidden-only, intermediate-only, both} -> (pad_hidden, pad_intermediate) +_PADDING_MODES: dict[str, tuple[bool, bool]] = { + "none": (False, False), + "hidden": (True, False), + "intermediate": (False, True), + "both": (True, True), +} + +# unquantized + the 4 fp8 configs AiterExperts supports (excludes MXFP4, +# stubbed below pending MI350/gfx950, and MK_QUANT_CONFIGS[1], unsupported). +_PADDING_QUANT_CONFIGS = [ + MK_QUANT_CONFIGS[0], # unquantized + MK_QUANT_CONFIGS[2], # fp8 channel weights / per-token activations + MK_QUANT_CONFIGS[3], # fp8 per-tensor weights / per-tensor activations + MK_QUANT_CONFIGS[4], # fp8 per-tensor weights / per-token activations + MK_QUANT_CONFIGS[5], # fp8 128x128-block weights / 128-block activations +] +_PADDING_QUANT_IDS = [ + "unquantized", + "fp8_channel_token", + "fp8_tensor_tensor", + "fp8_tensor_token", + "fp8_block_token", +] + + +def _slice_gate_up_rows(t: torch.Tensor, real_per_half: int) -> torch.Tensor: + """Slice a (E, 2*padded_per_half, ...) tensor's dim=1 down to the real + gate/up halves of `real_per_half` size each.""" + padded_per_half = t.shape[1] // 2 + return torch.cat( + [ + t[:, :real_per_half], + t[:, padded_per_half : padded_per_half + real_per_half], + ], + dim=1, + ) + + +def _slice_unpadded_weights( + weights: WeightTensors, + quant_config: TestMoEQuantConfig | None, + k_unpadded: int, + n_unpadded: int, +) -> WeightTensors: + """Derive the unpadded sub-block of a padded WeightTensors: same values, + restricted to the region hidden_pad/intermediate_pad keeps. + + Scale-shape choice (per_out_ch) must match _make_padding_matrix_weights().""" + block_shape = quant_config.block_shape if quant_config is not None else None + per_out_ch = quant_config is not None and quant_config.per_out_ch_quant + + w1 = _slice_gate_up_rows(weights.w1[:, :, :k_unpadded], n_unpadded) + w2 = weights.w2[:, :k_unpadded, :n_unpadded] + + if weights.w1_scale is None: + w1_scale = w2_scale = None + elif weights.w2_scale is None: + raise AssertionError("w1_scale and w2_scale must both be set or both None") + elif block_shape is not None: + block_n, block_k = block_shape + assert n_unpadded % block_n == 0 and k_unpadded % block_k == 0 + n_scale, k_scale = n_unpadded // block_n, k_unpadded // block_k + w1_scale = _slice_gate_up_rows(weights.w1_scale, n_scale)[..., :k_scale] + w2_scale = weights.w2_scale[:, :k_scale, :n_scale] + elif per_out_ch: + w1_scale = _slice_gate_up_rows(weights.w1_scale, n_unpadded) + w2_scale = weights.w2_scale[:, :k_unpadded, :] + else: + # Per-tensor scale: shape (E, 1, 1), independent of K/N -- reuse as-is. + w1_scale = weights.w1_scale + w2_scale = weights.w2_scale + + return WeightTensors(w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale) + + +def _make_padding_matrix_weights(config: Config) -> WeightTensors: + """Like WeightTensors.make(), but scales weights by + config.is_per_out_ch_quant instead of is_per_act_token_quant. + + WeightTensors.make() ties weight-scale shape to activation quant, so + fp8_tensor_token would get per-channel scales despite being a per-tensor + weight scheme. Scoped here rather than fixing WeightTensors.make(), + which other tests depend on.""" + (_, w1, w1_scale, w1_gs), (_, w2, w2_scale, w2_gs) = make_test_weights( + e=config.E, + n=config.N, + k=config.K, + in_dtype=config.dtype, + quant_dtype=config.quant_dtype, + block_shape=config.quant_block_shape, + per_out_ch_quant=config.is_per_out_ch_quant, + ) + return WeightTensors( + w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale, w1_gs=w1_gs, w2_gs=w2_gs + ) + + +def _slice_unpadded_rank_tensors( + rank_tensors: RankTensors, k_unpadded: int +) -> RankTensors: + return RankTensors( + hidden_states=rank_tensors.hidden_states[:, :k_unpadded].contiguous(), + hidden_states_scale=rank_tensors.hidden_states_scale, + topk_weights=rank_tensors.topk_weights, + topk_ids=rank_tensors.topk_ids, + expert_map=rank_tensors.expert_map, + ) + + +def _aiter_padding_matrix_worker( + pgi: ProcessGroupInfo, + vllm_config: VllmConfig, + cpu_group, + padded_config: Config, + padded_weights: WeightTensors, + verbose: bool, + hidden_pad_expected: int, + intermediate_pad_expected: int, +): + device = torch.device(f"cuda:{pgi.local_rank}") + init_workspace_manager(device) + set_random_seed(pgi.rank) + + weights = copy.deepcopy(padded_weights) + weights.to_current_device() + + rank_tensors = RankTensors.make(padded_config, pgi) + + with mock.patch.object( + rocm_aiter_ops, "fused_moe", wraps=rocm_aiter_ops.fused_moe + ) as fused_moe_mock: + mk_out = run_modular_kernel( + pgi, vllm_config, padded_config, weights, rank_tensors + ) + + assert fused_moe_mock.call_count > 0, ( + "Expected AiterExperts.apply() to call rocm_aiter_ops.fused_moe." + ) + for call in fused_moe_mock.call_args_list: + assert call.kwargs["hidden_pad"] == hidden_pad_expected, ( + f"AiterExperts.apply() forwarded hidden_pad=" + f"{call.kwargs['hidden_pad']}, expected {hidden_pad_expected}." + ) + assert call.kwargs["intermediate_pad"] == intermediate_pad_expected, ( + f"AiterExperts.apply() forwarded intermediate_pad=" + f"{call.kwargs['intermediate_pad']}, expected " + f"{intermediate_pad_expected}." + ) + + # AiterExperts.apply() never slices its own output -- the buffer stays + # raw (padded) width; only the caller-sliced unpadded prefix is checked. + assert mk_out.shape[-1] == padded_config.K, ( + f"AiterExperts output width {mk_out.shape[-1]} != raw hidden_dim " + f"{padded_config.K}." + ) + mk_out = mk_out[..., :_PADDING_K_UNPADDED] + # Only the real slice is checked: beyond it is unspecified/reused memory. + assert torch.isfinite(mk_out).all(), ( + "AiterExperts output must not contain NaN/Inf in the real " + "hidden_dim_unpadded-wide output slice." + ) + + unpadded_weights = _slice_unpadded_weights( + weights, + padded_config.quant_config, + _PADDING_K_UNPADDED, + _PADDING_N_UNPADDED, + ) + unpadded_rank_tensors = _slice_unpadded_rank_tensors( + rank_tensors, _PADDING_K_UNPADDED + ) + + with set_current_vllm_config(vllm_config): + ref_out = reference_moe_impl( + padded_config, unpadded_weights, unpadded_rank_tensors + ) + + # ref_out's magnitude here (~1e-3-1e-2) is below atol=3e-2, so a zeroed-out + # mk_out would still pass check_accuracy below -- guard against that. + ref_scale = ref_out.abs().mean() + mk_scale = mk_out.abs().mean() + assert mk_scale > 0.5 * ref_scale, ( + f"AiterExperts output magnitude (mean |mk_out|={mk_scale:.6f}) looks " + f"degenerate/zeroed vs. reference (mean |ref_out|={ref_scale:.6f})." + ) + + # Lenient check for every scheme, including unquantized: a strict + # assert_close(atol=rtol=3e-2) spuriously fails ~0.1-0.2% of elements + # on MI300 even for unquantized AiterExperts (confirmed on hardware). + check_accuracy(ref_out, mk_out, atol=3e-2, rtol=3e-2, percent=0.9) + + +@require_aiter_moe +@pytest.mark.parametrize("quant_config", _PADDING_QUANT_CONFIGS, ids=_PADDING_QUANT_IDS) +@pytest.mark.parametrize("mode", list(_PADDING_MODES)) +def test_aiter_moe_padding_matrix(mode: str, quant_config: TestMoEQuantConfig | None): + """See https://github.com/vllm-project/vllm/issues/54966 ("Test padding"). + + Exercises AiterExperts's hidden_pad/intermediate_pad (rocm_aiter_moe.py) + across {no padding, hidden-only, intermediate-only, both} x the + unquantized + 4 fp8 quant schemes it supports. Compares against an + unpadded reference and verifies the hidden_pad/intermediate_pad + AiterExperts.apply() forwards to rocm_aiter_ops.fused_moe. + """ + from vllm.model_executor.layers.fused_moe.prepare_finalize import ( + MoEPrepareAndFinalizeNoDPEPModular, + ) + + pad_hidden, pad_intermediate = _PADDING_MODES[mode] + + if ( + quant_config is not None + and quant_config.block_shape is not None + and not pad_hidden + ): + pytest.skip( + "AITER's block-quantized (128x128) CK GEMM kernel does not " + f"support the raw (unpadded) hidden_dim={_PADDING_K_UNPADDED} " + "shape on this hardware -- it raises 'wrong! device_gemm with " + "the specified compilation parameters does not support this " + "GEMM problem'. Only hidden_pad-forcing modes (raw " + f"hidden_dim={_PADDING_K_PADDED}) are exercised for this quant " + "scheme." + ) + + k = _PADDING_K_PADDED if pad_hidden else _PADDING_K_UNPADDED + n = _PADDING_N_PADDED if pad_intermediate else _PADDING_N_UNPADDED + + config = Config( + Ms=_PADDING_M, + K=k, + N=n, + E=_PADDING_E, + topks=_PADDING_TOPK, + dtype=torch.bfloat16, + quant_config=quant_config, + prepare_finalize_type=MoEPrepareAndFinalizeNoDPEPModular, + fused_experts_type=AiterExperts, + world_size=1, + hidden_dim_unpadded=_PADDING_K_UNPADDED if pad_hidden else None, + intermediate_size_per_partition_unpadded=( + _PADDING_N_UNPADDED if pad_intermediate else None + ), + ) + assert config.is_valid()[0] + assert config.fe_supports_quant_scheme(), ( + f"AiterExperts does not support quant scheme {quant_config}." + ) + + weights = _make_padding_matrix_weights(config) + vllm_config, env_dict = config.make_env_data() + + hidden_pad_expected = _PADDING_HIDDEN_PAD_EXPECTED if pad_hidden else 0 + intermediate_pad_expected = ( + _PADDING_INTERMEDIATE_PAD_EXPECTED if pad_intermediate else 0 + ) + + parallel_launch_with_config( + config.world_size, + _aiter_padding_matrix_worker, + vllm_config, + env_dict, + config, + weights, + False, + hidden_pad_expected, + intermediate_pad_expected, + ) + + +@pytest.mark.skip( + reason="MXFP4 AiterExperts padding requires MI350/gfx950 hardware. " + 'See https://github.com/vllm-project/vllm/issues/54966 ("Test padding").' +) +def test_aiter_moe_padding_matrix_mxfp4(): + """Stub for the {no padding, hidden-only, intermediate-only, both} MXFP4 + padding matrix once MI350/gfx950 hardware is available in CI.""" + + if __name__ == "__main__": # Ability to test individual PrepareAndFinalize and FusedExperts combination from .modular_kernel_tools.cli_args import make_config, make_config_arg_parser diff --git a/tests/kernels/moe/test_moe_fused_mul_sum.py b/tests/kernels/moe/test_moe_fused_mul_sum.py index 5affc6ea16..69f7a49947 100644 --- a/tests/kernels/moe/test_moe_fused_mul_sum.py +++ b/tests/kernels/moe/test_moe_fused_mul_sum.py @@ -16,10 +16,8 @@ from vllm.utils.torch_utils import set_random_seed NUM_TOKENS = [1, 17, 256] TOP_KS = [2, 8] HIDDEN_SIZE = 512 -# One BLOCK_K tile, a partial trailing tile, and Kimi-K3's real dims -# (moe_intermediate_size 3072, hidden_size 7168) -- exercises the per-tile loop -# and its tail mask (BLOCK_K caps at 512 for <=2-byte dtypes). -HIDDEN_SIZES = [512, 1000, 3072, 7168] +# Include partial tiles and hidden sizes above the CUDA 4096-element tile cap. +HIDDEN_SIZES = [512, 1000, 3072, 4097, 7168] DTYPES = [torch.bfloat16, torch.float32] NUM_EXPERTS = 16 @@ -233,7 +231,7 @@ def test_nonlocal_and_padding_adjacent(): @pytest.mark.parametrize("num_tokens", NUM_TOKENS) @pytest.mark.parametrize("top_k", TOP_KS) @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) -@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("dtype", [*DTYPES, torch.float16]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_weighted_sum_matches_reference( num_tokens: int, top_k: int, hidden_size: int, dtype: torch.dtype @@ -268,7 +266,11 @@ def test_weighted_sum_matches_reference( assert not out.isnan().any(), "invalid slots leaked into the reduction" ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map=None) - atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4) + atol, rtol = { + torch.bfloat16: (1e-2, 1e-2), + torch.float16: (1e-3, 1e-3), + torch.float32: (1e-4, 1e-4), + }[dtype] torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol) @@ -360,7 +362,7 @@ def test_num_valid_tokens_bounds_stale_padding(dtype: torch.dtype): """ set_random_seed(0) device = "cuda" - num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 512 + num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 7168 inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device) topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device) diff --git a/tests/kernels/moe/test_ocp_mx_moe.py b/tests/kernels/moe/test_ocp_mx_moe.py index f574c4d5d1..5a7053a8b1 100644 --- a/tests/kernels/moe/test_ocp_mx_moe.py +++ b/tests/kernels/moe/test_ocp_mx_moe.py @@ -781,6 +781,7 @@ def test_trtllm_mxfp4_deferred_finalize( global_num_experts=num_experts, expert_map=None, apply_router_weight_on_input=False, + routing_sink=None, ), ) return @@ -1780,6 +1781,7 @@ def test_rocm_mxfp4_moe_oracle( global_num_experts=num_experts, expert_map=None, apply_router_weight_on_input=False, + routing_sink=None, ) else: # Modular impl uses topk_weights and topk_ids diff --git a/tests/kernels/moe/test_routed_experts_capture_monolithic.py b/tests/kernels/moe/test_routed_experts_capture_monolithic.py index eec6180550..f15d8ab77f 100644 --- a/tests/kernels/moe/test_routed_experts_capture_monolithic.py +++ b/tests/kernels/moe/test_routed_experts_capture_monolithic.py @@ -1,22 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""End-to-end tests for routed-expert capture on the monolithic MoE path. +"""Routed-expert capture on the monolithic (fused router + experts) MoE path. -These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the -expert IDs picked by FlashInfer's fused router-and-experts kernels (the -"monolithic" path). When ``set_capture_fn`` is installed on -a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call -should: - - * allocate an int16 ``(num_tokens, top_k)`` buffer, - * pass it to FlashInfer as ``routing_replay_out``, - * invoke the callback after the kernel returns. +A monolithic kernel that ``supports_routing_replay_capture`` must write, into +the leading rows of ``routing_replay_out``, exactly the experts it routed each +token to. The router logits plant an unambiguous top-k per token, so the test +checks the reported experts against that answer and not just their range. """ from __future__ import annotations -from types import SimpleNamespace -from unittest.mock import patch +from collections.abc import Callable import pytest import torch @@ -38,6 +32,9 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( TrtLlmNvFp4ExpertsMonolithic, ) +from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEExpertsMonolithic, +) from vllm.platforms import current_platform try: @@ -57,822 +54,250 @@ if not current_platform.is_device_capability_family(100): allow_module_level=True, ) +NUM_EXPERTS = 32 +NUM_GROUPS = 4 # DeepSeekV3 grouped routing +TOPK_GROUPS = 2 +HIDDEN = 1024 +INTERMEDIATE = 1024 +DEVICE = torch.device("cuda:0") -def _shuffle_bf16_weights_block_major_k( - w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128 -) -> torch.Tensor: - """Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by - ``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert - row shuffle. +# A kernel under test, and its call: (router_logits, routing_replay_out) -> None. +Kernel = tuple[FusedMoEExpertsMonolithic, Callable[[torch.Tensor, torch.Tensor], None]] + + +def _planted_routing(num_tokens: int, top_k: int) -> tuple[torch.Tensor, torch.Tensor]: + """Router logits with one unambiguous answer under every routing method. + + Each token gives +4 to ``top_k`` experts drawn from at most two expert + groups and -4 to the rest, so softmax top-k, sigmoid top-k and DeepSeekV3 + grouped top-k all pick the same experts, with margins no rounding can flip. + Returns the logits and the planted experts of each token, sorted. """ - from flashinfer import shuffle_matrix_a - from flashinfer.fused_moe import convert_to_block_layout - - num_experts = w.shape[0] - shuffled = [] - for i in range(num_experts): - t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m) - shuffled.append(convert_to_block_layout(t, block_k)) - return torch.stack(shuffled).view(torch.bfloat16) + group_size = NUM_EXPERTS // NUM_GROUPS + planted = [] + for _ in range(num_tokens): + groups = torch.randperm(NUM_GROUPS)[:TOPK_GROUPS] + candidates = (groups[:, None] * group_size + torch.arange(group_size)).flatten() + planted.append(candidates[torch.randperm(len(candidates))[:top_k]]) + experts = torch.stack(planted) + logits = torch.full((num_tokens, NUM_EXPERTS), -4.0).scatter_(1, experts, 4.0) + return logits.to(DEVICE), experts.sort(dim=-1).values.to(DEVICE) -def _make_bf16_monolithic_experts( - num_experts: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - routing_method: RoutingMethodType, - device: torch.device, -) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]: - """Construct the monolithic BF16 experts plus the BlockMajorK weights - expected by ``trtllm_bf16_moe``. - """ - parallel_cfg = FusedMoEParallelConfig.make_no_parallel() - moe_config = FusedMoEConfig( - num_experts=num_experts, +def _routing_kwargs(routing_method: RoutingMethodType) -> dict: + if routing_method != RoutingMethodType.DeepSeekV3: + return {} + return dict( + num_expert_group=NUM_GROUPS, + topk_group=TOPK_GROUPS, + e_score_correction_bias=torch.zeros( + NUM_EXPERTS, device=DEVICE, dtype=torch.bfloat16 + ), + routed_scaling_factor=1.0, + ) + + +def _moe_config(top_k: int, routing_method: RoutingMethodType) -> FusedMoEConfig: + return FusedMoEConfig( + num_experts=NUM_EXPERTS, experts_per_token=top_k, - hidden_dim=hidden_size, - intermediate_size=intermediate_size, - num_local_experts=num_experts, - num_logical_experts=num_experts, - moe_parallel_config=parallel_cfg, + hidden_dim=HIDDEN, + intermediate_size=INTERMEDIATE, + num_local_experts=NUM_EXPERTS, + num_logical_experts=NUM_EXPERTS, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), in_dtype=torch.bfloat16, activation=MoEActivation.SILU, - device=device, + device=DEVICE, routing_method=routing_method, max_num_tokens=16, ) - quant_config = FusedMoEQuantConfig.make( - quant_dtype=None, - per_act_token_quant=False, - per_out_ch_quant=False, - block_shape=None, - ) + + +def _bf16_kernel(top_k: int, routing_method: RoutingMethodType) -> Kernel: + from flashinfer import shuffle_matrix_a + from flashinfer.fused_moe import convert_to_block_layout + + def block_major_k(w: torch.Tensor) -> torch.Tensor: + """(E, M, K) -> the BlockMajorK layout ``trtllm_bf16_moe`` expects.""" + return torch.stack( + [ + convert_to_block_layout(shuffle_matrix_a(e.view(torch.uint8), 64), 128) + for e in w + ] + ).view(torch.bfloat16) experts = TrtLlmBf16ExpertsMonolithic( - moe_config=moe_config, quant_config=quant_config + moe_config=_moe_config(top_k, routing_method), + quant_config=FusedMoEQuantConfig.make( + quant_dtype=None, + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + ), ) - - gemm1 = ( - torch.randn( - num_experts, - 2 * intermediate_size, - hidden_size, - device=device, - dtype=torch.bfloat16, - ) + w13 = block_major_k( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, device=DEVICE).bfloat16() * 0.1 ) - gemm2 = ( - torch.randn( - num_experts, - hidden_size, - intermediate_size, - device=device, - dtype=torch.bfloat16, - ) - * 0.1 - ) - w13 = _shuffle_bf16_weights_block_major_k(gemm1) - w2 = _shuffle_bf16_weights_block_major_k(gemm2) - return experts, w13, w2 - - -def _run_bf16_monolithic( - experts: TrtLlmBf16ExpertsMonolithic, - hidden_states: torch.Tensor, - w13: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - num_experts: int, - *, - n_group: int | None = None, - topk_group: int | None = None, - routed_scaling_factor: float | None = None, - e_score_correction_bias: torch.Tensor | None = None, -) -> torch.Tensor: - return experts.apply( - hidden_states=hidden_states, - w1=w13, - w2=w2, - router_logits=router_logits, - activation=MoEActivation.SILU, - global_num_experts=num_experts, - expert_map=None, - a1q_scale=None, - apply_router_weight_on_input=False, - num_expert_group=n_group, - topk_group=topk_group, - e_score_correction_bias=e_score_correction_bias, - routed_scaling_factor=routed_scaling_factor, + w2 = block_major_k( + torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE, device=DEVICE).bfloat16() * 0.1 ) - -_DSV3_NUM_EXPERTS = 32 -_DSV3_N_GROUP = 4 -_DSV3_TOPK_GROUP = 2 - - -def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor: - return torch.randn(num_experts, device=device, dtype=torch.bfloat16) - - -@pytest.mark.parametrize("num_tokens", [2, 7, 16]) -@pytest.mark.parametrize("top_k", [2, 4]) -def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts( - num_tokens: int, - top_k: int, -) -> None: - """The capture callback should receive the int16 routed-expert IDs the - kernel actually used, the values should be valid expert indices, and - each token should pick ``top_k`` distinct experts.""" - if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: - pytest.skip( - f"DSV3 requires top_k <= n_group * topk_group " - f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})" - ) - torch.manual_seed(0) - device = torch.device("cuda:0") - - num_experts = _DSV3_NUM_EXPERTS - hidden_size = 1024 - intermediate_size = 1024 - - experts, w13, w2 = _make_bf16_monolithic_experts( - num_experts=num_experts, - top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - routing_method=RoutingMethodType.DeepSeekV3, - device=device, - ) - - captured: list[torch.Tensor] = [] - - def capture_fn(replay_out: torch.Tensor) -> None: - captured.append(replay_out.clone()) - - assert experts.supports_routing_replay_capture() - experts.set_capture_fn(capture_fn) - - hidden_states = ( - torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 - ) - router_logits = torch.rand( - num_tokens, num_experts, device=device, dtype=torch.float32 - ) - routing_bias = _make_dsv3_routing_bias(num_experts, device) - - _ = _run_bf16_monolithic( - experts, - hidden_states=hidden_states, - w13=w13, - w2=w2, - router_logits=router_logits, - num_experts=num_experts, - n_group=_DSV3_N_GROUP, - topk_group=_DSV3_TOPK_GROUP, - routed_scaling_factor=1.0, - e_score_correction_bias=routing_bias, - ) - - assert len(captured) == 1 - replay = captured[0] - assert replay.dtype == torch.int16 - assert replay.shape == (num_tokens, top_k) - assert (replay >= 0).all(), f"got out-of-range values: {replay}" - assert (replay < num_experts).all(), f"got out-of-range values: {replay}" - - for t in range(num_tokens): - unique = replay[t].unique() - assert unique.numel() == top_k, ( - f"token {t}: expected {top_k} distinct experts, " - f"got {unique.numel()} ({replay[t].tolist()})" - ) - - -@pytest.mark.parametrize("num_tokens", [2, 7, 16]) -@pytest.mark.parametrize( - "routing_method", - [ - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ], -) -def test_trtllm_bf16_monolithic_routing_replay_non_dsv3( - num_tokens: int, - routing_method: RoutingMethodType, -) -> None: - """Routing replay works for non-DeepSeekV3 routing methods too. - FlashInfer's ``routing_replay_out`` is routing-method-agnostic.""" - torch.manual_seed(0) - device = torch.device("cuda:0") - - num_experts = 8 - top_k = 2 - hidden_size = 1024 - intermediate_size = 1024 - - experts, w13, w2 = _make_bf16_monolithic_experts( - num_experts=num_experts, - top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - routing_method=routing_method, - device=device, - ) - - captured: list[torch.Tensor] = [] - experts.set_capture_fn(lambda r: captured.append(r.clone())) - - hidden_states = ( - torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 - ) - router_logits = torch.rand( - num_tokens, num_experts, device=device, dtype=torch.float32 - ) - - _ = _run_bf16_monolithic( - experts, - hidden_states=hidden_states, - w13=w13, - w2=w2, - router_logits=router_logits, - num_experts=num_experts, - ) - - assert len(captured) == 1 - replay = captured[0] - assert replay.dtype == torch.int16 - assert replay.shape == (num_tokens, top_k) - assert (replay >= 0).all(), f"got out-of-range values: {replay}" - assert (replay < num_experts).all(), f"got out-of-range values: {replay}" - for t in range(num_tokens): - unique = replay[t].unique() - assert unique.numel() == top_k, ( - f"token {t}: expected {top_k} distinct experts, " - f"got {unique.numel()} ({replay[t].tolist()})" - ) - - -def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None: - """With no callback installed the kernel should not see a - ``routing_replay_out`` tensor — verify the helper short-circuits.""" - torch.manual_seed(0) - device = torch.device("cuda:0") - experts, _, _ = _make_bf16_monolithic_experts( - num_experts=_DSV3_NUM_EXPERTS, - top_k=2, - hidden_size=1024, - intermediate_size=1024, - routing_method=RoutingMethodType.DeepSeekV3, - device=device, - ) - # No callback installed. - buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device) - assert buf is None - - # Dispatch is also a no-op. - experts._maybe_dispatch_routing_replay(buf, num_tokens=4) - - -def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None: - """FlashInfer's ``routing_replay_out`` is supported by all routing - methods, so ``supports_routing_replay_capture`` should be True - regardless of routing method.""" - device = torch.device("cuda:0") - for routing_method in ( - RoutingMethodType.DeepSeekV3, - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ): - experts, _, _ = _make_bf16_monolithic_experts( - num_experts=_DSV3_NUM_EXPERTS, - top_k=2, - hidden_size=1024, - intermediate_size=1024, - routing_method=routing_method, - device=device, - ) - assert experts.supports_routing_replay_capture() is True, ( - f"{routing_method!r} should support routing replay capture" - ) - - -def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None: - """When capture is installed, the allocated buffer is int16 and shaped - ``(num_tokens, experts_per_token)``.""" - device = torch.device("cuda:0") - experts, _, _ = _make_bf16_monolithic_experts( - num_experts=_DSV3_NUM_EXPERTS, - top_k=4, - hidden_size=1024, - intermediate_size=1024, - routing_method=RoutingMethodType.DeepSeekV3, - device=device, - ) - experts.set_capture_fn(lambda r: None) - buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device) - assert buf is not None - assert buf.dtype == torch.int16 - assert buf.shape[0] >= 11 - assert buf.shape[1] == 4 - assert buf.device.type == "cuda" - - -def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None: - """End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback - on the monolithic experts and verify the captured rows land in the - capturer's device buffer at the correct layer slot. - - Mirrors the monolithic path in ``bind_routed_experts_capturer``: a single - closure is installed on ``fused_experts`` and routes per-layer based on - the closed-over ``layer_id``. - """ - from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsCapturer, - ) - - torch.manual_seed(7) - device = torch.device("cuda:0") - num_tokens = 4 - top_k = 2 - num_experts = _DSV3_NUM_EXPERTS - hidden_size = 1024 - intermediate_size = 1024 - - experts, w13, w2 = _make_bf16_monolithic_experts( - num_experts=num_experts, - top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - routing_method=RoutingMethodType.DeepSeekV3, - device=device, - ) - - num_layers = 3 - layer_id = 1 - capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) - capturer.dp_rank = 0 - capturer.tp_size = 1 - capturer.device_buffer = torch.full( - (num_tokens + 4, num_layers, top_k), - -1, - dtype=torch.int32, - device=device, - ) - - def capture_fn(replay_out: torch.Tensor) -> None: - capturer.capture(layer_id, replay_out) - - experts.set_capture_fn(capture_fn) - - hidden_states = ( - torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 - ) - router_logits = torch.rand( - num_tokens, num_experts, device=device, dtype=torch.float32 - ) - routing_bias = _make_dsv3_routing_bias(num_experts, device) - - # Patch get_forward_context to return a dp_metadata=None context so the - # capturer takes the single-DP branch. - import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec - - with patch.object( - rec, - "get_forward_context", - return_value=SimpleNamespace(dp_metadata=None), - ): - _ = _run_bf16_monolithic( - experts, - hidden_states=hidden_states, - w13=w13, + def forward(router_logits, routing_replay_out): + hidden_states = torch.randn(len(router_logits), HIDDEN, device=DEVICE) * 0.1 + experts.apply( + hidden_states=hidden_states.bfloat16(), + w1=w13, w2=w2, router_logits=router_logits, - num_experts=num_experts, - n_group=_DSV3_N_GROUP, - topk_group=_DSV3_TOPK_GROUP, - routed_scaling_factor=1.0, - e_score_correction_bias=routing_bias, + activation=MoEActivation.SILU, + global_num_experts=NUM_EXPERTS, + expert_map=None, + a1q_scale=None, + apply_router_weight_on_input=False, + routing_replay_out=routing_replay_out, + **_routing_kwargs(routing_method), ) - captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu() - # Valid expert IDs at this layer. - assert (captured >= 0).all() - assert (captured < num_experts).all() - for t in range(num_tokens): - unique = captured[t].unique() - assert unique.numel() == top_k, ( - f"token {t}: expected {top_k} distinct experts at layer " - f"{layer_id}, got {unique.numel()}" - ) - - # Other layers / trailing token rows untouched. - for other_layer in range(num_layers): - if other_layer == layer_id: - continue - assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), ( - f"layer {other_layer} should be untouched, got writes" - ) - assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), ( - "tail rows beyond num_tokens should remain sentinel" - ) + return experts, forward -# ---------------------------------------------------------------------------- -# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic`` -# ---------------------------------------------------------------------------- - - -def _make_fp8_block_scale_monolithic_experts( - num_experts: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - device: torch.device, -) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]: - """Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale - code path with DSV3 routing. - - Weights are shuffled into the BlockMajorK layout the kernel expects - (same helper the vLLM weight loader uses for DeepSeek-FP8 models). - """ +def _fp8_block_kernel(top_k: int, routing_method: RoutingMethodType) -> Kernel: from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( _shuffle_deepseek_fp8_moe_weights, ) - block_k = 128 - parallel_cfg = FusedMoEParallelConfig.make_no_parallel() - moe_config = FusedMoEConfig( - num_experts=num_experts, - experts_per_token=top_k, - hidden_dim=hidden_size, - intermediate_size=intermediate_size, - num_local_experts=num_experts, - num_logical_experts=num_experts, - moe_parallel_config=parallel_cfg, - in_dtype=torch.bfloat16, - activation=MoEActivation.SILU, - device=device, - routing_method=RoutingMethodType.DeepSeekV3, - max_num_tokens=16, + block = 128 + w13, w2 = _shuffle_deepseek_fp8_moe_weights( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, device=DEVICE).to( + torch.float8_e4m3fn + ), + torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE, device=DEVICE).to( + torch.float8_e4m3fn + ), ) - - # Random fp8 weights + ones-block scales (the kernel decoder only cares - # that the per-block scales are present and finite for routing/replay). - gemm1 = torch.randn( - num_experts, 2 * intermediate_size, hidden_size, device=device - ).to(torch.float8_e4m3fn) - gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to( - torch.float8_e4m3fn - ) - w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2) - - w1_scale = torch.ones( - num_experts, - 2 * intermediate_size // block_k, - hidden_size // block_k, - device=device, - dtype=torch.float32, - ) - w2_scale = torch.ones( - num_experts, - hidden_size // block_k, - intermediate_size // block_k, - device=device, - dtype=torch.float32, - ) - quant_config = fp8_w8a8_moe_quant_config( - w1_scale=w1_scale, - w2_scale=w2_scale, - block_shape=[block_k, block_k], - per_act_token_quant=False, - ) - experts = TrtLlmFp8ExpertsMonolithic( - moe_config=moe_config, quant_config=quant_config - ) - return experts, w13_shuffled, w2_shuffled - - -def _run_fp8_block_scale_monolithic( - experts: TrtLlmFp8ExpertsMonolithic, - hidden_states_fp8: torch.Tensor, - hidden_states_scale: torch.Tensor, - w13: torch.Tensor, - w2: torch.Tensor, - router_logits: torch.Tensor, - num_experts: int, - routing_bias: torch.Tensor, -) -> torch.Tensor: - return experts.apply( - hidden_states=hidden_states_fp8, - w1=w13, - w2=w2, - router_logits=router_logits, - activation=MoEActivation.SILU, - global_num_experts=num_experts, - expert_map=None, - # The block-scale apply path reads ``a1q_scale`` and transposes it - # to ``(hidden_size/128, num_tokens)`` for the kernel call. - a1q_scale=hidden_states_scale, - apply_router_weight_on_input=False, - num_expert_group=_DSV3_N_GROUP, - topk_group=_DSV3_TOPK_GROUP, - e_score_correction_bias=routing_bias, - routed_scaling_factor=1.0, + moe_config=_moe_config(top_k, routing_method), + quant_config=fp8_w8a8_moe_quant_config( + w1_scale=torch.ones( + NUM_EXPERTS, 2 * INTERMEDIATE // block, HIDDEN // block, device=DEVICE + ), + w2_scale=torch.ones( + NUM_EXPERTS, HIDDEN // block, INTERMEDIATE // block, device=DEVICE + ), + block_shape=[block, block], + per_act_token_quant=False, + ), ) - -@pytest.mark.parametrize("num_tokens", [2, 7, 16]) -@pytest.mark.parametrize("top_k", [2, 4]) -def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts( - num_tokens: int, - top_k: int, -) -> None: - """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale - path, DSV3 routing) captures valid expert IDs.""" - if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: - pytest.skip( - f"DSV3 requires top_k <= n_group * topk_group " - f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})" - ) - torch.manual_seed(0) - device = torch.device("cuda:0") - - num_experts = _DSV3_NUM_EXPERTS - hidden_size = 1024 - intermediate_size = 1024 - - experts, w13, w2 = _make_fp8_block_scale_monolithic_experts( - num_experts=num_experts, - top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - device=device, - ) - assert experts.supports_routing_replay_capture() - - captured: list[torch.Tensor] = [] - experts.set_capture_fn(lambda r: captured.append(r.clone())) - - # Per-token / per-block hidden scales (ones is fine for the routing - # path; the GEMM output isn't being asserted on). - hidden_states = ( - torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 - ).to(torch.float8_e4m3fn) - hidden_states_scale = torch.ones( - num_tokens, hidden_size // 128, device=device, dtype=torch.float32 - ) - router_logits = torch.rand( - num_tokens, num_experts, device=device, dtype=torch.float32 - ) - routing_bias = _make_dsv3_routing_bias(num_experts, device) - - _ = _run_fp8_block_scale_monolithic( - experts, - hidden_states_fp8=hidden_states, - hidden_states_scale=hidden_states_scale, - w13=w13, - w2=w2, - router_logits=router_logits, - num_experts=num_experts, - routing_bias=routing_bias, - ) - - assert len(captured) == 1 - replay = captured[0] - assert replay.dtype == torch.int16 - assert replay.shape == (num_tokens, top_k) - assert (replay >= 0).all(), f"got out-of-range values: {replay}" - assert (replay < num_experts).all(), f"got out-of-range values: {replay}" - for t in range(num_tokens): - unique = replay[t].unique() - assert unique.numel() == top_k, ( - f"token {t}: expected {top_k} distinct experts, " - f"got {unique.numel()} ({replay[t].tolist()})" + def forward(router_logits, routing_replay_out): + num_tokens = len(router_logits) + hidden_states = torch.randn(num_tokens, HIDDEN, device=DEVICE) * 0.1 + experts.apply( + hidden_states=hidden_states.to(torch.float8_e4m3fn), + w1=w13, + w2=w2, + router_logits=router_logits, + activation=MoEActivation.SILU, + global_num_experts=NUM_EXPERTS, + expert_map=None, + # Per-block activation scales, (num_tokens, hidden / 128). + a1q_scale=torch.ones(num_tokens, HIDDEN // block, device=DEVICE), + apply_router_weight_on_input=False, + routing_replay_out=routing_replay_out, + **_routing_kwargs(routing_method), ) - -# ---------------------------------------------------------------------------- -# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic`` -# ---------------------------------------------------------------------------- + return experts, forward -def _make_nvfp4_monolithic_experts( - num_experts: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - device: torch.device, -) -> tuple[ - TrtLlmNvFp4ExpertsMonolithic, - torch.Tensor, # w13 (packed nvfp4 uint8) - torch.Tensor, # w13 block-scale (fp8) - torch.Tensor, # w2 (packed nvfp4 uint8) - torch.Tensor, # w2 block-scale (fp8) - torch.Tensor, # input global scale (per-tensor float32) -]: - """Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights - and DSV3 routing. - - NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor - "global" scale. We follow the layout in ``test_ocp_mx_moe.py`` / - ``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``: - * weights: uint8 (packed fp4) ``(E, M, K//2)`` - * weight scales: fp8 ``(E, M, K//16)`` - * hidden states: uint8 (packed fp4) ``(N, K//2)`` - * hidden state scales: fp8 ``(N, K//16)`` - """ +def _nvfp4_kernel(top_k: int, routing_method: RoutingMethodType) -> Kernel: from flashinfer import fp4_quantize - block_size = 16 - parallel_cfg = FusedMoEParallelConfig.make_no_parallel() - moe_config = FusedMoEConfig( - num_experts=num_experts, - experts_per_token=top_k, - hidden_dim=hidden_size, - intermediate_size=intermediate_size, - num_local_experts=num_experts, - num_logical_experts=num_experts, - moe_parallel_config=parallel_cfg, - in_dtype=torch.bfloat16, - activation=MoEActivation.SILU, - device=device, - routing_method=RoutingMethodType.DeepSeekV3, - max_num_tokens=16, - ) + block = 16 + global_scale = torch.tensor(1.0, device=DEVICE) - gemm1 = torch.randn( - num_experts, - 2 * intermediate_size, - hidden_size, - device=device, - dtype=torch.bfloat16, - ) - gemm2 = torch.randn( - num_experts, - hidden_size, - intermediate_size, - device=device, - dtype=torch.bfloat16, - ) - # Per-tensor weight scaling factor (used to build ``g1_alphas`` / - # ``g2_alphas`` below). - w_global_scale = torch.tensor(1.0, device=device) - # Per-tensor input scaling factor. - a_global_scale = torch.tensor(1.0, device=device) + def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Packed fp4 values and their fp8 per-16 block scales, (E, M, K/16).""" + q, scale = fp4_quantize( + w, global_scale, block, sf_use_ue8m0=False, is_sf_swizzled_layout=False + ) + return q, scale.view(torch.float8_e4m3fn).reshape(*w.shape[:2], -1) - w13_q, w13_scale = fp4_quantize( - gemm1, - w_global_scale, - block_size, - sf_use_ue8m0=False, - is_sf_swizzled_layout=False, + w13, w13_scale = quantize( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, device=DEVICE).bfloat16() ) - w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape( - num_experts, 2 * intermediate_size, hidden_size // block_size + w2, w2_scale = quantize( + torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE, device=DEVICE).bfloat16() ) - w2_q, w2_scale = fp4_quantize( - gemm2, - w_global_scale, - block_size, - sf_use_ue8m0=False, - is_sf_swizzled_layout=False, - ) - w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape( - num_experts, hidden_size, intermediate_size // block_size - ) - - # NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2, - # g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these. - g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32) - a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32) - - quant_config = FusedMoEQuantConfig.make( - quant_dtype="nvfp4", - per_act_token_quant=False, - per_out_ch_quant=False, - block_shape=None, - w1_scale=w13_scale, - w2_scale=w2_scale, - g1_alphas=g_alphas, - g2_alphas=g_alphas, - a1_gscale=a_global_scale, - a2_gscale=a2_gscale, - ) - + ones = torch.ones(NUM_EXPERTS, device=DEVICE) experts = TrtLlmNvFp4ExpertsMonolithic( - moe_config=moe_config, quant_config=quant_config - ) - return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale - - -def _run_nvfp4_monolithic( - experts: TrtLlmNvFp4ExpertsMonolithic, - hidden_states_q: torch.Tensor, - hidden_states_scale: torch.Tensor, - router_logits: torch.Tensor, - num_experts: int, - routing_bias: torch.Tensor, -) -> torch.Tensor: - """The monolithic NVFP4 apply expects packed fp4 hidden states + the - matching fp8 per-block scale stored in the ``a1q_scale`` slot.""" - # Stash the weight tensors on the experts in the locations the apply() - # implementation reads from (it pulls them from quant_config / scales - # already; w1/w2 come in as args). - return experts.apply( - hidden_states=hidden_states_q, - w1=experts._w13_packed, - w2=experts._w2_packed, - router_logits=router_logits, - activation=MoEActivation.SILU, - global_num_experts=num_experts, - expert_map=None, - a1q_scale=hidden_states_scale, - apply_router_weight_on_input=False, - num_expert_group=_DSV3_N_GROUP, - topk_group=_DSV3_TOPK_GROUP, - e_score_correction_bias=routing_bias, - routed_scaling_factor=1.0, + moe_config=_moe_config(top_k, routing_method), + quant_config=FusedMoEQuantConfig.make( + quant_dtype="nvfp4", + per_act_token_quant=False, + per_out_ch_quant=False, + block_shape=None, + w1_scale=w13_scale, + w2_scale=w2_scale, + g1_alphas=ones, + g2_alphas=ones, + a1_gscale=global_scale, + a2_gscale=ones, + ), ) + def forward(router_logits, routing_replay_out): + hidden_states = torch.randn(len(router_logits), HIDDEN, device=DEVICE) * 0.1 + # apply() reinterprets the packed uint8 activation scales itself. + hidden_q, hidden_scale = fp4_quantize( + hidden_states.bfloat16(), + global_scale, + block, + sf_use_ue8m0=False, + is_sf_swizzled_layout=False, + ) + experts.apply( + hidden_states=hidden_q, + w1=w13, + w2=w2, + router_logits=router_logits, + activation=MoEActivation.SILU, + global_num_experts=NUM_EXPERTS, + expert_map=None, + a1q_scale=hidden_scale, + apply_router_weight_on_input=False, + routing_replay_out=routing_replay_out, + **_routing_kwargs(routing_method), + ) -@pytest.mark.parametrize("num_tokens", [2, 7, 16]) + return experts, forward + + +KERNELS = {"bf16": _bf16_kernel, "fp8_block": _fp8_block_kernel, "nvfp4": _nvfp4_kernel} + + +@pytest.mark.parametrize("kernel", list(KERNELS)) +@pytest.mark.parametrize( + "routing_method", [RoutingMethodType.Renormalize, RoutingMethodType.DeepSeekV3] +) @pytest.mark.parametrize("top_k", [2, 4]) -def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts( - num_tokens: int, - top_k: int, +@pytest.mark.parametrize("num_tokens", [2, 7, 16]) +def test_routing_replay_reports_the_routed_experts( + kernel: str, routing_method: RoutingMethodType, top_k: int, num_tokens: int ) -> None: - """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs - on the DSV3 routing path.""" - if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP: - pytest.skip( - f"DSV3 requires top_k <= n_group * topk_group " - f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})" - ) - from flashinfer import fp4_quantize - torch.manual_seed(0) - device = torch.device("cuda:0") - - num_experts = _DSV3_NUM_EXPERTS - hidden_size = 1024 - intermediate_size = 1024 - block_size = 16 - - experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts( - num_experts=num_experts, - top_k=top_k, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - device=device, - ) - # The apply() reads w1/w2 from its args, but we keep them on the experts - # for convenience of the helper. - experts._w13_packed = w13_q - experts._w2_packed = w2_q - + experts, forward = KERNELS[kernel](top_k, routing_method) assert experts.supports_routing_replay_capture() - captured: list[torch.Tensor] = [] - experts.set_capture_fn(lambda r: captured.append(r.clone())) - - hidden_states = ( - torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1 - ) - hidden_states_q, hidden_states_scale = fp4_quantize( - hidden_states, - a_gs, - block_size, - sf_use_ue8m0=False, - is_sf_swizzled_layout=False, - ) - # The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave - # ``hidden_states_scale`` in its native (uint8 packed) form. - router_logits = torch.rand( - num_tokens, num_experts, device=device, dtype=torch.float32 - ) - routing_bias = _make_dsv3_routing_bias(num_experts, device) - - _ = _run_nvfp4_monolithic( - experts, - hidden_states_q=hidden_states_q, - hidden_states_scale=hidden_states_scale, - router_logits=router_logits, - num_experts=num_experts, - routing_bias=routing_bias, + router_logits, planted = _planted_routing(num_tokens, top_k) + # One spare row: the kernel must write the batch's rows and nothing else. + routing_replay_out = torch.full( + (num_tokens + 1, top_k), -1, dtype=torch.int16, device=DEVICE ) - assert len(captured) == 1 - replay = captured[0] - assert replay.dtype == torch.int16 - assert replay.shape == (num_tokens, top_k) - assert (replay >= 0).all(), f"got out-of-range values: {replay}" - assert (replay < num_experts).all(), f"got out-of-range values: {replay}" - for t in range(num_tokens): - unique = replay[t].unique() - assert unique.numel() == top_k, ( - f"token {t}: expected {top_k} distinct experts, " - f"got {unique.numel()} ({replay[t].tolist()})" - ) + forward(router_logits, routing_replay_out) + + reported = routing_replay_out[:num_tokens].long().sort(dim=-1).values + assert torch.equal(reported, planted), (reported, planted) + assert (routing_replay_out[num_tokens:] == -1).all() diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 412e772a2b..7cdaf2b495 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -91,6 +91,12 @@ QWEN4_EXP_SM100_CASES = [ for num_tokens, config in plans.items() ] +QWEN4_EXP_SM121_CASES = [ + (n, k, num_tokens, config) + for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS.items() + for num_tokens, config in plans.items() +] + EXPECTED_CUTE_CONFIGS = { (3072, 7168, 1): (224, 3, 4, 8), (3072, 7168, 2): (128, 3, 2, 8), @@ -725,6 +731,7 @@ def test_qwen4_exp_hopper_plans_are_valid() -> None: ((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS), ((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS), ((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS), + ((12, 1), qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS), ((8, 0), {}), ], ) @@ -939,6 +946,29 @@ def test_qwen4_exp_sm100_selected_shapes( assert cosine > 0.999 +@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM121_CASES) +def test_qwen4_exp_sm121_selected_shapes( + n: int, + k: int, + num_tokens: int, + config: SkinnyGemmConfig, +) -> None: + _require_capability_and_cute((12, 1)) + torch.manual_seed(42 + num_tokens) + x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + + selected = qwen4_exp_gemm.QWEN4_EXP_SM121_GEMM_PLANS[(n, k)][num_tokens] + assert selected == config + output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight) + + reference = torch.nn.functional.linear(x, weight) + cosine = torch.nn.functional.cosine_similarity( + output.float().flatten(), reference.float().flatten(), dim=0 + ).item() + assert cosine > 0.999 + + def test_glm52_q_b_nonpacked_single_row_falls_back() -> None: _require_sm103_and_cute() spec = glm52_gemm.GLM52_Q_B_PROJECTION diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index b6c469181f..05fda51b7c 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -636,8 +636,8 @@ def test_top_k_per_row_decode_gfx950_k512_masked_graph_replay() -> None: @pytest.mark.skipif(not current_platform.is_rocm(), reason="This test requires ROCm") @torch.inference_mode() def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: + from vllm._aiter_ops import rocm_aiter_ops from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _get_aiter_top_k_kernel, _launch_aiter_top_k_per_row_prefill, ) @@ -647,17 +647,14 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: top_k = 4 indices = torch.empty((3, top_k), dtype=torch.int32, device="cuda") - aiter_topk_kernel = _get_aiter_top_k_kernel( + if not rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=True, compress_ratio=4, num_rows=logits.shape[0], - on_gfx950=True, - ) - if aiter_topk_kernel is None: + ): pytest.skip("AITER top-k is unavailable") assert ( _launch_aiter_top_k_per_row_prefill( - aiter_topk_kernel, logits, row_starts, row_ends, @@ -683,10 +680,7 @@ def test_aiter_c4a_prefill_topk_returns_sequence_local_indices() -> None: def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( num_speculative_tokens: int, ) -> None: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - _get_aiter_top_k_kernel, - _launch_aiter_top_k_per_row_decode, - ) + from vllm._aiter_ops import rocm_aiter_ops query_tokens = num_speculative_tokens + 1 offsets = torch.arange(query_tokens, dtype=torch.int32, device="cuda") @@ -701,20 +695,18 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( top_k = 4 indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") - aiter_topk_kernel = _get_aiter_top_k_kernel( + if not rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=False, compress_ratio=4, num_rows=num_rows, max_valid_seq_len=seq_lens.max().item(), - on_gfx950=True, - ) - if aiter_topk_kernel is None: + ): pytest.skip("AITER top-k is unavailable") assert ( - _launch_aiter_top_k_per_row_decode( - aiter_topk_kernel, + rocm_aiter_ops.indexer_top_k_decode( logits, - seq_lens, + 1, + seq_lens.reshape(-1), indices, top_k, ) @@ -729,130 +721,107 @@ def test_aiter_c4a_decode_topk_uses_exact_mtp_lengths( assert torch.all(indices[row_idx, num_valid:] == -1) -def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: - from vllm.v1.attention.ops import rocm_aiter_mla_sparse +@requires_gfx950 +@torch.inference_mode() +def test_aiter_decode_topk_matches_per_row_under_mtp() -> None: + """The aiter backend expands (B, 1) seq_lens into per-row ends before + launching, so next_n > 1 must agree with the in-tree kernel, including + the rows that clamp to a zero-length prefix. + """ + from vllm.model_executor.layers.indexer_topk import get_indexer_topk - def prefill_kernel(*args, **kwargs) -> None: - pass + next_n = 4 + top_k = 512 + num_columns = 1024 + seq_lens = torch.tensor([[2], [513], [700]], dtype=torch.int32, device="cuda") + num_rows = seq_lens.shape[0] * next_n - def decode_kernel(*args, **kwargs) -> None: - pass + torch.manual_seed(0) + logits = torch.randn(num_rows, num_columns, dtype=torch.float32, device="cuda") + + outputs = {} + for backend in ("aiter", "per_row"): + indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda") + get_indexer_topk(backend)(logits, seq_lens, next_n, indices, top_k, num_columns) + outputs[backend] = indices + + row_ends = ( + (seq_lens - next_n + 1 + torch.arange(next_n, device="cuda")) + .clamp_(min=0) + .reshape(-1) + .tolist() + ) + assert row_ends[:next_n] == [0, 0, 1, 2] + + for row_idx, row_end in enumerate(row_ends): + num_valid = min(top_k, row_end) + aiter_row = outputs["aiter"][row_idx] + assert set(aiter_row[:num_valid].tolist()) == set( + outputs["per_row"][row_idx][:num_valid].tolist() + ) + assert torch.all(aiter_row[num_valid:] == -1) + + +def _enable_aiter_topk(monkeypatch, enabled: bool = True) -> None: + from vllm._aiter_ops import rocm_aiter_ops monkeypatch.setattr( - rocm_aiter_mla_sparse, - "_get_aiter_topk_ops", - lambda: (prefill_kernel, decode_kernel), + rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: enabled, raising=True ) - get_kernel = rocm_aiter_mla_sparse._get_aiter_top_k_kernel + + +def test_aiter_c4a_topk_kernel_selection(monkeypatch) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + _enable_aiter_topk(monkeypatch) + is_supported = rocm_aiter_ops.is_indexer_top_k_supported eligible_cases = [ - ( - dict(is_prefill=True, compress_ratio=4, num_rows=1, on_gfx950=True), - prefill_kernel, + dict(is_prefill=True, compress_ratio=4, num_rows=1), + dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_536), + dict( + is_prefill=False, + compress_ratio=4, + num_rows=257, + max_valid_seq_len=250_000, ), - ( - dict( - is_prefill=False, - compress_ratio=4, - num_rows=1, - max_valid_seq_len=65_536, - on_gfx950=True, - ), - decode_kernel, - ), - ( - dict( - is_prefill=False, - compress_ratio=4, - num_rows=257, - max_valid_seq_len=250_000, - num_columns=524_288, - on_gfx950=True, - ), - decode_kernel, - ), - ( - dict( - is_prefill=False, - compress_ratio=2, - num_rows=385, - max_valid_seq_len=250_000, - num_columns=524_288, - topk_tokens=512, - on_gfx950=True, - ), - decode_kernel, - ), - ( - dict( - is_prefill=False, - compress_ratio=2, - num_rows=384, - max_valid_seq_len=50_000, - num_columns=1_048_577, - topk_tokens=512, - on_gfx950=True, - ), - decode_kernel, - ), - ] - for kwargs, expected_kernel in eligible_cases: - assert get_kernel(**kwargs) is expected_kernel - - def fail_if_imported() -> None: - pytest.fail("ineligible shapes must not import AITER top-k") - - monkeypatch.setattr( - rocm_aiter_mla_sparse, - "_get_aiter_topk_ops", - fail_if_imported, - ) - ineligible_cases = [ - dict(is_prefill=True, compress_ratio=1, num_rows=1, on_gfx950=True), dict( is_prefill=False, compress_ratio=2, num_rows=384, max_valid_seq_len=1_048_576, - num_columns=1_048_576, - topk_tokens=512, - on_gfx950=True, - ), - dict( - is_prefill=False, - compress_ratio=4, - num_rows=1, - max_valid_seq_len=65_537, - on_gfx950=True, ), + ] + for kwargs in eligible_cases: + assert is_supported(**kwargs) is True + + ineligible_cases = [ + dict(is_prefill=True, compress_ratio=1, num_rows=1), + dict(is_prefill=False, compress_ratio=4, num_rows=1, max_valid_seq_len=65_537), dict( is_prefill=False, compress_ratio=4, num_rows=256, max_valid_seq_len=125_000, - on_gfx950=True, - ), - dict( - is_prefill=False, - compress_ratio=4, - num_rows=320, - max_valid_seq_len=2_500, - on_gfx950=False, ), ] for kwargs in ineligible_cases: - assert get_kernel(**kwargs) is None + assert is_supported(**kwargs) is False - monkeypatch.setattr(rocm_aiter_mla_sparse, "_get_aiter_topk_ops", lambda: None) - assert ( - get_kernel( - is_prefill=True, - compress_ratio=4, - num_rows=1, - on_gfx950=True, - ) - is None - ) + _enable_aiter_topk(monkeypatch, enabled=False) + assert is_supported(is_prefill=True, compress_ratio=4, num_rows=1) is False + + +def test_dsv4_native_top_k_window() -> None: + """The in-tree decode kernel's measured advantage band over AITER, which + applies to the DSV4 indexer only.""" + from vllm._aiter_ops import rocm_aiter_ops + + prefers_native = rocm_aiter_ops.dsv4_indexer_prefers_native_top_k + assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=512) is True + assert prefers_native(num_rows=385, num_columns=524_288, topk_tokens=512) is False + assert prefers_native(num_rows=128, num_columns=1_048_577, topk_tokens=512) is False + assert prefers_native(num_rows=128, num_columns=524_288, topk_tokens=1024) is False @pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") @@ -1927,3 +1896,141 @@ def test_sparse_indexer_topk_backend_resolution() -> None: if not is_sm100: with pytest.raises(RuntimeError, match="SM100a/SM103a"): resolve("deep_select") + + +def _patch_aiter_topk(monkeypatch, decode_kernel, enabled: bool = True) -> None: + from vllm._aiter_ops import rocm_aiter_ops + + _enable_aiter_topk(monkeypatch, enabled=enabled) + monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: enabled, raising=True) + monkeypatch.setattr(rocm_aiter_ops, "indexer_top_k_decode", decode_kernel) + + +def _make_topk(backend: str): + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.indexer_topk import SparseIndexerTopk + + cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) + with set_current_vllm_config(cfg): + return SparseIndexerTopk() + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +def test_sparse_indexer_topk_aiter_requires_rocm() -> None: + """Explicitly requesting the ROCm-only backend must fail fast elsewhere.""" + logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda") + with pytest.raises(RuntimeError, match="ROCm"): + _make_topk("aiter").resolve_backend(logits, 2048, 8) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +def test_sparse_indexer_topk_backend_resolution_aiter(monkeypatch) -> None: + """Aiter is opt-in only: "auto" never selects it (the shape-dependent perf + gate lives at the call site), and an explicit request validates hard + capability only.""" + + def decode_kernel(*args, **kwargs) -> None: + pass + + _patch_aiter_topk(monkeypatch, decode_kernel) + logits = torch.randn(8, 4096, dtype=torch.float32, device="cuda") + + assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row" + assert _make_topk("aiter").resolve_backend(logits, 2048, 8) == "aiter" + # Shapes the call-site perf gate would decline still resolve. + assert _make_topk("aiter").resolve_backend(logits, 512, 8) == "aiter" + with pytest.raises(RuntimeError, match="topk_tokens must be in"): + _make_topk("aiter").resolve_backend(logits, 3000, 8) + + _enable_aiter_topk(monkeypatch, enabled=False) + with pytest.raises(RuntimeError, match="no tuned indexer top-k kernels"): + _make_topk("aiter").resolve_backend(logits, 2048, 8) + + _patch_aiter_topk(monkeypatch, decode_kernel, enabled=False) + assert _make_topk("auto").resolve_backend(logits, 2048, 8) == "per_row" + with pytest.raises(RuntimeError, match="VLLM_ROCM_USE_AITER"): + _make_topk("aiter").resolve_backend(logits, 2048, 8) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +@torch.inference_mode() +def test_sparse_indexer_topk_aiter_applies_no_shape_gate(monkeypatch) -> None: + """Shape-dependent perf gating belongs to the call site: once "aiter" is + requested the dispatcher must run it for every supported k, including the + k=512 band the in-tree kernel wins on DSV4's compressed-KV logits.""" + from vllm.model_executor.layers import indexer_topk + + calls: list[tuple] = [] + _patch_aiter_topk(monkeypatch, lambda *args: calls.append(args)) + monkeypatch.setattr( + indexer_topk.ops, + "top_k_per_row_decode", + lambda *args: pytest.fail("an explicit aiter request must not fall back"), + ) + + num_rows = 8 + logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda") + seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda") + indices = torch.empty((num_rows, 512), dtype=torch.int32, device="cuda") + op = _make_topk("aiter") + + op(logits, seq_lens, 1, indices, 512, 4096) + # Past the compressed-context boundary the call-site gate declines on. + op(logits, seq_lens, 1, indices, 512, 4 * (64 * 1024 + 1)) + assert len(calls) == 2 + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +@torch.inference_mode() +def test_sparse_indexer_topk_auto_stays_on_per_row(monkeypatch) -> None: + """The "auto" chain must not reach AITER on its own: the decode top-k runs + through top_k_per_row_decode unless the caller asks for the aiter backend.""" + from vllm.model_executor.layers import indexer_topk + + def decode_kernel(*args, **kwargs) -> None: + pytest.fail("AITER is opt-in only") + + _patch_aiter_topk(monkeypatch, decode_kernel) + calls: list[tuple] = [] + monkeypatch.setattr( + indexer_topk.ops, "top_k_per_row_decode", lambda *args: calls.append(args) + ) + + num_rows = 8 + logits = torch.randn(num_rows, 4096, dtype=torch.float32, device="cuda") + seq_lens = torch.full((num_rows, 1), 4096, dtype=torch.int32, device="cuda") + indices = torch.empty((num_rows, 1024), dtype=torch.int32, device="cuda") + + _make_topk("auto")(logits, seq_lens, 1, indices, 1024, 4096) + assert len(calls) == 1 + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="requires ROCm") +@torch.inference_mode() +def test_sparse_indexer_topk_aiter_expands_mtp_row_ends(monkeypatch) -> None: + """AITER's decode kernel only implements the next_n == 1 form, so spec + decode batches must be handed the equivalent expanded per-row ends.""" + from vllm.model_executor.layers import indexer_topk + + recorded: dict = {} + + def decode_kernel(logits, next_n, seq_lens, *args, **kwargs) -> None: + recorded["next_n"] = next_n + recorded["row_ends"] = seq_lens.tolist() + + _patch_aiter_topk(monkeypatch, decode_kernel) + monkeypatch.setattr( + indexer_topk.ops, + "top_k_per_row_decode", + lambda *args: pytest.fail("an explicit aiter request must not fall back"), + ) + + next_n = 2 + logits = torch.randn(4, 4096, dtype=torch.float32, device="cuda") + seq_lens = torch.tensor([[100], [200]], dtype=torch.int32, device="cuda") + indices = torch.empty((4, 1024), dtype=torch.int32, device="cuda") + + _make_topk("aiter")(logits, seq_lens, next_n, indices, 1024, 800) + + assert recorded["next_n"] == 1 + assert recorded["row_ends"] == [99, 100, 199, 200] diff --git a/tests/lora/test_fused_moe_lora_kernel.py b/tests/lora/test_fused_moe_lora_kernel.py index 8b5b3d577a..11e4d31f87 100644 --- a/tests/lora/test_fused_moe_lora_kernel.py +++ b/tests/lora/test_fused_moe_lora_kernel.py @@ -362,6 +362,83 @@ def test_fused_moe_lora_kernel( torch.testing.assert_close(output, output2, atol=1e-2, rtol=1e-2) +@pytest.mark.parametrize("max_lora_rank", [128, 256]) +@pytest.mark.parametrize("add_inputs", [True, False]) +@pytest.mark.parametrize("num_slices", [1, 2]) +@pytest.mark.parametrize("device", DEVICES) +def test_fused_moe_lora_kernel_large_rank_fallback( + max_lora_rank, + add_inputs, + num_slices, + device, +): + """Ranks above the one-shot limit fall back to shrink/expand kernels.""" + torch.set_default_device(device) + set_random_seed(42) + + num_tokens = 8 + num_sequences = 4 + top_k_num = 2 + num_experts = 4 + max_loras = 2 + N = 256 + K = 256 + block_size = 16 + dtype = torch.bfloat16 + + topk_ids, topk_weights, token_lora_mapping, lora_ids = sample_data( + num_tokens, num_sequences, max_loras, num_experts, top_k_num + ) + lora_a_stacked = [ + torch.rand( + (max_loras, num_experts, max_lora_rank, K), + dtype=dtype, + ) + for _ in range(num_slices) + ] + lora_b_stacked = [ + torch.rand( + (max_loras, num_experts, N // num_slices, max_lora_rank), + dtype=dtype, + ) + for _ in range(num_slices) + ] + hidden_states = torch.rand((num_tokens, K), dtype=dtype) + + residual = torch.randn((num_tokens, top_k_num, N), dtype=dtype) + output = residual.clone() + use_fused_moe_lora_kernel( + topk_ids, + topk_weights, + token_lora_mapping, + max_lora_rank, + top_k_num, + lora_ids, + lora_a_stacked, + lora_b_stacked, + hidden_states, + output, + max_loras, + num_experts, + block_size, + add_inputs=add_inputs, + ) + + expected = use_torch( + hidden_states, + token_lora_mapping, + topk_ids, + lora_a_stacked, + lora_b_stacked, + top_k_num, + num_slices, + ) + if add_inputs: + expected += residual + + torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2) + + def use_fused_moe_lora_kernel_naive( topk_ids, topk_weights, diff --git a/tests/model_executor/test_qwen3_omni_mrope.py b/tests/model_executor/test_qwen3_omni_mrope.py new file mode 100644 index 0000000000..bb8de63dc5 --- /dev/null +++ b/tests/model_executor/test_qwen3_omni_mrope.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses +from dataclasses import dataclass + +import pytest +import torch + +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, +) +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) + + +@pytest.fixture(autouse=True, scope="module") +def _force_cpu_default_device(): + # get_mrope_input_positions returns CPU tensors (via torch.from_numpy). + # Ensure the default device is CPU so the rest of the test tensors match. + original = torch.get_default_device() + torch.set_default_device("cpu") + yield + torch.set_default_device(original) + + +TEXT_TOKEN_ID = 1 +VISION_START_TOKEN_ID = 100 +VISION_END_TOKEN_ID = 101 +IMAGE_TOKEN_ID = 200 +VIDEO_TOKEN_ID = 201 +AUDIO_TOKEN_ID = 202 +AUDIO_START_TOKEN_ID = 300 +AUDIO_END_TOKEN_ID = 301 + + +@dataclass +class DummyVisionConfig: + spatial_merge_size: int = 1 + + +@dataclass +class DummyConfig: + position_id_per_seconds: int = 25 + vision_config: DummyVisionConfig = dataclasses.field( + default_factory=DummyVisionConfig + ) + + +def make_model() -> Qwen3OmniMoeThinkerForConditionalGeneration: + model = Qwen3OmniMoeThinkerForConditionalGeneration.__new__( + Qwen3OmniMoeThinkerForConditionalGeneration + ) + model.config = DummyConfig() + return model + + +def _elem(data): + return MultiModalFieldElem( + data=torch.tensor(data), + field=None, # HACK. + ) + + +def make_image_feature(offset: int, grid_thw: tuple[int, int, int]): + return MultiModalFeatureSpec( + data=MultiModalKwargsItem({"image_grid_thw": _elem(grid_thw)}), + modality="image", + identifier="DUMMY", + mm_position=PlaceholderRange( + offset=offset, length=int(torch.tensor(grid_thw).prod()) + ), + ) + + +def make_video_feature( + offset: int, grid_thw: tuple[int, int, int], use_audio_in_video: bool = False +): + data = {"video_grid_thw": _elem(grid_thw)} + if use_audio_in_video: + data["use_audio_in_video"] = _elem(True) + return MultiModalFeatureSpec( + data=MultiModalKwargsItem(data), + modality="video", + identifier="DUMMY", + mm_position=PlaceholderRange( + offset=offset, length=int(torch.tensor(grid_thw).prod()) + ), + ) + + +def make_audio_feature(offset: int, audio_feature_length: int): + return MultiModalFeatureSpec( + data=MultiModalKwargsItem( + {"audio_feature_lengths": _elem(audio_feature_length)} + ), + modality="audio", + identifier="DUMMY", + mm_position=PlaceholderRange(offset=offset, length=0), + ) + + +def test_mrope_positions_image_at_prompt_end(): + """A prompt ending with <|vision_end|> must still produce seq_len positions.""" + model = make_model() + # [<|vision_start|>, <|image_pad|> x 4, <|vision_end|>] + input_tokens = ( + [VISION_START_TOKEN_ID] + [IMAGE_TOKEN_ID] * 4 + [VISION_END_TOKEN_ID] + ) + mm_features = [make_image_feature(offset=1, grid_thw=(1, 2, 2))] + + positions, _ = model.get_mrope_input_positions(input_tokens, mm_features) + + assert positions.shape == (3, len(input_tokens)) + + +def test_mrope_positions_audio_at_prompt_end(): + model = make_model() + # audio_feature_length=16 -> 2 audio tokens. + # [<|audio_start|>, <|audio_pad|> x 2, <|audio_end|>] + input_tokens = [AUDIO_START_TOKEN_ID] + [AUDIO_TOKEN_ID] * 2 + [AUDIO_END_TOKEN_ID] + mm_features = [make_audio_feature(offset=1, audio_feature_length=16)] + + positions, _ = model.get_mrope_input_positions(input_tokens, mm_features) + + assert positions.shape == (3, len(input_tokens)) + + +def test_mrope_positions_video_at_prompt_end(): + model = make_model() + # [T, <|vision_start|>, <|video_pad|> x 8, <|vision_end|>] + input_tokens = ( + [TEXT_TOKEN_ID] + + [VISION_START_TOKEN_ID] + + [VIDEO_TOKEN_ID] * 8 + + [VISION_END_TOKEN_ID] + ) + mm_features = [make_video_feature(offset=2, grid_thw=(2, 2, 2))] + + positions, _ = model.get_mrope_input_positions(input_tokens, mm_features) + + assert positions.shape == (3, len(input_tokens)) + + +def test_mrope_positions_audio_in_video_at_prompt_end(): + model = make_model() + # video grid (1, 2, 2) -> 4 video tokens; audio_feature_length=16 -> 2 + # audio tokens, interleaved after the video tokens. + # [T, T, <|vision_start|>, <|audio_start|>, V x 4, A x 2, + # <|audio_end|>, <|vision_end|>] + input_tokens = ( + [TEXT_TOKEN_ID] * 2 + + [VISION_START_TOKEN_ID, AUDIO_START_TOKEN_ID] + + [VIDEO_TOKEN_ID] * 4 + + [AUDIO_TOKEN_ID] * 2 + + [AUDIO_END_TOKEN_ID, VISION_END_TOKEN_ID] + ) + mm_features = [ + make_video_feature(offset=3, grid_thw=(1, 2, 2), use_audio_in_video=True), + make_audio_feature(offset=3, audio_feature_length=16), + ] + + positions, _ = model.get_mrope_input_positions(input_tokens, mm_features) + + assert positions.shape == (3, len(input_tokens)) + + +def test_mrope_positions_image_then_text(): + """Media followed by text: positions must match the expected layout exactly. + + Before the fix this case passed the length guard but every position from + the first <|image_pad|> onward was silently shifted by one slot. + """ + model = make_model() + # [T, T, <|vision_start|>, <|image_pad|> x 4, <|vision_end|>, T, T, T] + input_tokens = ( + [TEXT_TOKEN_ID] * 2 + + [VISION_START_TOKEN_ID] + + [IMAGE_TOKEN_ID] * 4 + + [VISION_END_TOKEN_ID] + + [TEXT_TOKEN_ID] * 3 + ) + mm_features = [make_image_feature(offset=3, grid_thw=(1, 2, 2))] + + positions, delta = model.get_mrope_input_positions(input_tokens, mm_features) + + expected = torch.tensor( + [ + [0, 1, 2, 3, 3, 3, 3, 5, 6, 7, 8], + [0, 1, 2, 3, 3, 4, 4, 5, 6, 7, 8], + [0, 1, 2, 3, 4, 3, 4, 5, 6, 7, 8], + ] + ) + torch.testing.assert_close(positions, expected) + assert delta == 9 - len(input_tokens) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index ff683ae44c..2d57236f8e 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -11,8 +11,16 @@ from vllm.config import VllmConfig from vllm.config.compilation import CompilationMode from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import CPUExpertsInt4 +from vllm.model_executor.layers.fused_moe.modular_kernel import ( + FusedMoEKernelMonolithicImpl, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPMonolithic, +) from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, + RoutedExpertsSink, bind_routed_experts_capturer, ) from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter @@ -174,45 +182,139 @@ def test_public_binding_supports_direct_capture_source(): assert calls == [(3, topk_ids)] -def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch): - class DummyFusedMoE: - def __init__(self): - self.layer_id = 3 - self.router = _make_router() - # Use a concrete monolithic expert and override its capability - # instead of instantiating the abstract base class directly. - from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( - CPUExpertsFp8, - ) +_SINK_CONFIG = SimpleNamespace( + use_ep=False, + dp_size=1, + ep_size=1, + max_num_tokens=4, + experts_per_token=2, + device="cpu", +) - fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8) - self.routed_experts = types.SimpleNamespace( - quant_method=types.SimpleNamespace( - is_monolithic=True, - moe_kernel=types.SimpleNamespace( - impl=types.SimpleNamespace(fused_experts=fused_experts) - ), - ) - ) - self._quant_method = self.routed_experts.quant_method - self._quant_method.moe_kernel.impl.fused_experts = fused_experts - fused_experts.supports_routing_replay_capture = lambda: False - class DummyCapturer: - def capture(self, layer_id, topk_ids): - pass +class _ToyMonolithicExperts(CPUExpertsInt4): + """A monolithic kernel instance, as a weight reload rebuilds it: routes + token t to experts (t + offset, t + offset + 1).""" - dummy_module = DummyFusedMoE() + moe_config = _SINK_CONFIG + quant_config = None + + def __init__(self, offset: int = 0, supports_capture: bool = True): + self.offset = offset + self.supports_capture = supports_capture + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def supports_routing_replay_capture(self) -> bool: + return self.supports_capture + + def apply(self, hidden_states, *args, routing_replay_out=None, **kwargs): + if routing_replay_out is not None: + t = torch.arange(hidden_states.shape[0]) + routing_replay_out[: len(t)] = torch.stack([t, t + 1], 1) + self.offset + return hidden_states + + +class _MonolithicMoE: + """The parts of a MoE layer that capture binding reads.""" + + layer_id = 5 + + def __init__(self, experts: _ToyMonolithicExperts): + self.router = _make_router() + kernel = SimpleNamespace(impl=SimpleNamespace(fused_experts=experts)) + self._quant_method = SimpleNamespace(is_monolithic=True, moe_kernel=kernel) + self.routed_experts = SimpleNamespace( + quant_method=self._quant_method, routing_sink=None + ) + + +def _bind_monolithic( + monkeypatch, experts: _ToyMonolithicExperts, capture +) -> _MonolithicMoE: import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer - monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE) + monkeypatch.setattr(fused_moe_layer, "MoERunner", _MonolithicMoE) + layer = _MonolithicMoE(experts) + bind_routed_experts_capturer( + SimpleNamespace(modules=lambda: [layer]), SimpleNamespace(capture=capture) + ) + return layer + +def test_public_binding_gives_monolithic_layers_a_sink(monkeypatch): + """A monolithic kernel routes internally, so binding gives the layer, which + outlives kernel rebuilds, a sink that captures under its layer id.""" + calls = [] + layer = _bind_monolithic( + monkeypatch, _ToyMonolithicExperts(), lambda *args: calls.append(args) + ) + ids = torch.tensor([[1, 2]]) + layer.routed_experts.routing_sink.capture_fn(ids) + assert calls == [(_MonolithicMoE.layer_id, ids)] + + +def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch): with pytest.raises(ValueError, match="monolithic MoE kernel"): - bind_routed_experts_capturer( - types.SimpleNamespace(modules=lambda: [dummy_module]), DummyCapturer() + _bind_monolithic( + monkeypatch, _ToyMonolithicExperts(supports_capture=False), Mock() ) +@pytest.mark.parametrize( + ("use_ep", "dp_size", "ep_size", "rows"), + [(False, 1, 1, 4), (False, 2, 4, 8), (True, 2, 4, 16)], + ids=["single-rank", "dp", "ep"], +) +def test_routed_experts_sink_holds_a_dispatched_batch(use_ep, dp_size, ep_size, rows): + """One int16 buffer, allocated once, large enough for the batch the kernel + sees after a DP or EP gather.""" + config = SimpleNamespace( + **{ + **vars(_SINK_CONFIG), + "use_ep": use_ep, + "dp_size": dp_size, + "ep_size": ep_size, + } + ) + sink = RoutedExpertsSink(config, Mock()) + + assert sink.buffer.shape == (rows, 2) and sink.buffer.dtype == torch.int16 + + +def test_monolithic_capture_survives_kernel_rebuilds(): + """The bug in #59449: a weight reload rebuilds the kernel. Through the + production kernel path, the layer's sink captures what the current kernel + routed, before and after a rebuild, and a rebuilt kernel that cannot + capture is refused instead of leaving stale rows.""" + captured = [] + sink = RoutedExpertsSink(_SINK_CONFIG, lambda ids: captured.append(ids.tolist())) + + def forward(experts: _ToyMonolithicExperts) -> None: + kernel = FusedMoEKernelMonolithicImpl( + MoEPrepareAndFinalizeNoDPEPMonolithic(), experts + ) + kernel.apply( + torch.zeros(3, 8), + w1=None, + w2=None, + router_logits=torch.zeros(3, 16), + activation=None, + global_num_experts=16, + expert_map=None, + apply_router_weight_on_input=False, + routing_sink=sink, + ) + + for offset in (0, 7): # the kernel before and after a reload + forward(_ToyMonolithicExperts(offset)) + assert captured[-1] == [[t + offset, t + offset + 1] for t in range(3)] + with pytest.raises(ValueError, match="not supported"): + forward(_ToyMonolithicExperts(supports_capture=False)) + + def test_routed_experts_capturer_single_dp_no_metadata(): """dp_metadata is None: capture writes the full topk_ids rows.""" capturer = _capturer_with_buffer(dp_rank=0) diff --git a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py index 9f94da459c..8b4a09f42f 100644 --- a/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py +++ b/tests/models/glm5next/test_sparse_indexer_topk_dispatch.py @@ -8,10 +8,6 @@ import torch from vllm.config import VllmConfig, set_current_vllm_config from vllm.platforms import current_platform -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda(), reason="CUDA-only dispatch" -) - def _require_deep_gemm() -> None: from vllm.utils.deep_gemm import has_deep_gemm @@ -20,17 +16,10 @@ def _require_deep_gemm() -> None: pytest.skip("kpool indexer requires DeepGEMM") -@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) -def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: - """The kpool indexer must read kernel_config.sparse_indexer_topk_backend - and hand it to the shared SparseIndexerTopk dispatcher, rather than - hard-coding its own cooperative/persistent/per_row choice.""" - _require_deep_gemm() - from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool - +def _build(indexer_cls, backend: str): cfg = VllmConfig(kernel_config={"sparse_indexer_topk_backend": backend}) with set_current_vllm_config(cfg): - op = SparseAttnIndexerKpool( + return indexer_cls( k_cache=None, quant_block_size=128, scale_fmt="ue8m0", @@ -40,4 +29,73 @@ def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> N max_total_seq_len=8192, topk_indices_buffer=torch.empty(8, 2176, dtype=torch.int32, device="cuda"), ) - assert op.topk_backend == backend + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only dispatch") +@pytest.mark.parametrize("backend", ["auto", "persistent", "cooperative", "torch"]) +def test_kpool_indexer_dispatches_through_shared_topk_backend(backend: str) -> None: + """The kpool indexer must read kernel_config.sparse_indexer_topk_backend + and hand it to the shared SparseIndexerTopk dispatcher, rather than + hard-coding its own cooperative/persistent/per_row choice.""" + _require_deep_gemm() + from vllm.models.glm5next.nvidia.sparse_indexer import SparseAttnIndexerKpool + + assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +@pytest.mark.parametrize("backend", ["auto", "aiter", "per_row", "torch"]) +def test_kpool_indexer_dispatches_through_shared_topk_backend_rocm( + backend: str, +) -> None: + """The AMD kpool indexer must go through the same dispatcher, so the + AITER decode top-k is reachable from kernel_config.""" + from vllm.models.glm5next.amd.sparse_indexer import SparseAttnIndexerKpool + + assert _build(SparseAttnIndexerKpool, backend).topk_backend == backend + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only dispatch") +@pytest.mark.parametrize( + "configured,full_cudagraph,expected", + [ + ("auto", False, "aiter"), + ("auto", True, "auto"), + ("aiter", True, "aiter"), + ("per_row", False, "per_row"), + ], +) +def test_kpool_decode_topk_backend_never_auto_selects_aiter_under_full_cudagraph( + monkeypatch: pytest.MonkeyPatch, + configured: str, + full_cudagraph: bool, + expected: str, +) -> None: + """A FULL cudagraph bakes in one kernel for every replay, so "auto" cannot + use the context length to pick AITER and must stay on the in-tree kernel. + Explicit backends keep working, and outside FULL the long-context heuristic + still narrows to AITER.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.models.glm5next.amd import sparse_indexer + + monkeypatch.setattr( + rocm_aiter_ops, "is_indexer_top_k_enabled", lambda: True, raising=False + ) + monkeypatch.setattr( + rocm_aiter_ops, + "is_indexer_top_k_supported", + lambda **kwargs: True, + raising=False, + ) + + assert ( + sparse_indexer._kpool_decode_topk_backend( + configured, + num_rows=8, + max_valid_seq_len=64 * 1024, + select_k=512, + index_kpool=4, + full_cudagraph=full_cudagraph, + ) + == expected + ) diff --git a/tests/models/kimi_k3/test_gfx942_int4.py b/tests/models/kimi_k3/test_gfx942_int4.py new file mode 100644 index 0000000000..20bef41eb8 --- /dev/null +++ b/tests/models/kimi_k3/test_gfx942_int4.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +from collections.abc import Callable +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.fused_moe import RoutedExperts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.quantization.mxfp4 import ( + _moe_weight_override_is_int4, + _use_k3_situ_int4_gfx942, +) +from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, +) + +pytestmark = pytest.mark.cpu_test + + +class _FakeRocmModule(ModuleType): + on_gfx942: Callable[[], bool] + + +class _FakeAiterOpsModule(ModuleType): + rocm_aiter_ops: SimpleNamespace + + +def _vllm_config(moe_weight: str | None): + quantization_config = ( + None + if moe_weight is None + else QuantizationConfigArgs(moe={"weight": moe_weight}) + ) + return SimpleNamespace( + model_config=SimpleNamespace(quantization_config=quantization_config) + ) + + +@pytest.mark.parametrize( + ("moe_weight", "expected"), + [ + (None, False), + ("mxfp4", False), + ("int4_per_group_32", True), + ], +) +def test_moe_weight_override_is_int4(monkeypatch, moe_weight, expected): + monkeypatch.setattr( + "vllm.config.get_current_vllm_config", + lambda: _vllm_config(moe_weight), + ) + assert _moe_weight_override_is_int4() is expected + + +@pytest.mark.parametrize("int4_requested", [False, True]) +def test_k3_situ_int4_gfx942_requires_explicit_override(int4_requested): + moe = SimpleNamespace( + activation=MoEActivation.SITU, + activation_situ_linear_beta=25.0, + ) + fake_rocm = _FakeRocmModule("vllm.platforms.rocm") + fake_rocm.on_gfx942 = lambda: True + fake_aiter_ops = _FakeAiterOpsModule("vllm._aiter_ops") + fake_aiter_ops.rocm_aiter_ops = SimpleNamespace(is_fused_moe_enabled=lambda: True) + with ( + patch("vllm.platforms.current_platform.is_rocm", return_value=True), + patch.dict( + sys.modules, + { + "vllm.platforms.rocm": fake_rocm, + "vllm._aiter_ops": fake_aiter_ops, + }, + ), + patch( + "vllm.model_executor.layers.quantization.mxfp4." + "_moe_weight_override_is_int4", + return_value=int4_requested, + ), + ): + assert _use_k3_situ_int4_gfx942(moe) is int4_requested + + +def test_int4_per_group_32_is_not_an_online_moe_method(): + # Load-time gfx942 requant. Must not match the online MoE table. + config = OnlineQuantizationConfig( + QuantizationConfigArgs(moe={"weight": "int4_per_group_32"}) + ) + assert ( + config.resolve_quant_method_cls( + Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" + ) + is None + ) + + +def test_unknown_online_moe_weight_still_raises(): + config = OnlineQuantizationConfig( + QuantizationConfigArgs(moe={"weight": "fp8_per_token"}) + ) + with pytest.raises(ValueError, match="online quantization"): + config.resolve_quant_method_cls( + Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" + ) diff --git a/tests/models/kimi_k3/test_prefix_cache.py b/tests/models/kimi_k3/test_prefix_cache.py new file mode 100644 index 0000000000..3525550fe9 --- /dev/null +++ b/tests/models/kimi_k3/test_prefix_cache.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 multi-turn prefix cache reuse with KV offload, P/D and DCP.""" + +import contextlib +import json +import random +from collections.abc import Iterator +from dataclasses import dataclass +from math import lcm + +import pytest +import requests +from prometheus_client.parser import text_string_to_metric_families + +from tests.models.utils import check_logprobs_close +from tests.utils import RemoteOpenAIServer, multi_gpu_marks +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port + +# Official Kimi-K3 config with 8 layers and no weights. +MODEL = "riverclouds/Kimi-K3-8L-dummy" +# The recipe's DSpark draft, reading target layers that exist in MODEL. +DRAFT = "riverclouds/Kimi-K3-DSpark-for-8L-dummy" + +# Conversation: the first prompt spans FIRST_PROMPT_BLOCKS blocks plus an +# unaligned tail, and each later turn appends TURN_TOKENS random tokens. +NUM_TURNS = 3 +FIRST_PROMPT_BLOCKS = 3 +UNALIGNED_TAIL = 37 +TURN_TOKENS = 777 +MAX_TOKENS = 8 +NUM_LOGPROBS = 5 +STARTUP_TIMEOUT = 1200 + +BASE_ARGS = [ + "--trust-remote-code", + "--language-model-only", + "--load-format", + "dummy", + "--no-enable-flashinfer-autotune", + "--enable-prompt-tokens-details", + # Default max_num_seqs OOMs one GPU during warmup. + "--max-num-seqs", + "8", +] + +pytestmark = pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Kimi-K3 NVIDIA kernels require the SM100 family", +) + +SPEC_CONFIG = { + "method": "dspark", + "model": DRAFT, + "num_speculative_tokens": 7, + "draft_sample_method": "probabilistic", + "rejection_sample_method": "block", +} + +NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"} +# SimpleCPUOffloadConnector as in the Kimi-K3 recipe. +OFFLOAD = { + "kv_connector": "SimpleCPUOffloadConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "cpu_bytes_to_use_per_rank": 32 << 30, + "lazy_offload": False, + }, +} +NIXL_OFFLOAD = { + "kv_connector": "MultiConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": {"connectors": [NIXL, OFFLOAD]}, +} + + +@dataclass(frozen=True) +class Mode: + # A finer prefix_match_unit enables partial hits inside a Mamba block. + prefix_match_unit: int | None + spec: bool + + @property + def name(self) -> str: + name = "block" if self.prefix_match_unit is None else "partial" + return f"{name}-dspark" if self.spec else name + + +MODES = [Mode(unit, spec) for spec in (False, True) for unit in (None, 128)] + + +@dataclass(frozen=True) +class Instance: + tp: int = 1 + dp: int = 1 + dcp: int = 1 + kv_config: dict | None = None + prefix_caching: bool = True + + @property + def num_gpus(self) -> int: + return self.tp * self.dp + + def args(self, mode: Mode) -> list[str]: + args = BASE_ARGS + ["--tensor-parallel-size", str(self.tp)] + if mode.spec: + args += ["--speculative-config", json.dumps(SPEC_CONFIG)] + if self.dp > 1: + args += ["--data-parallel-size", str(self.dp), "--enable-expert-parallel"] + if self.dcp > 1: + args += ["--decode-context-parallel-size", str(self.dcp)] + if self.kv_config is not None: + args += ["--kv-transfer-config", json.dumps(self.kv_config)] + if not self.prefix_caching: + return args + ["--no-enable-prefix-caching"] + args.append("--enable-prefix-caching") + if mode.prefix_match_unit is not None: + args += ["--prefix-match-unit", str(mode.prefix_match_unit)] + return args + + +@dataclass(frozen=True) +class Deployment: + decode: Instance + prefill: Instance | None = None + # Reset GPU prefix caches each turn so hits must come from CPU. + offload: bool = False + + @property + def instances(self) -> tuple[Instance, ...]: + return (self.prefill, self.decode) if self.prefill else (self.decode,) + + @property + def num_gpus(self) -> int: + return sum(i.num_gpus for i in self.instances) + + +DEPLOYMENTS = { + "plain": Deployment(Instance()), + "offload": Deployment(Instance(kv_config=OFFLOAD), offload=True), + "dcp2": Deployment(Instance(tp=2, dcp=2)), + "dcp2-offload": Deployment(Instance(tp=2, dcp=2, kv_config=OFFLOAD), offload=True), + "pd": Deployment(Instance(kv_config=NIXL), prefill=Instance(kv_config=NIXL)), + # NIXL rejects prefix caching on a hybrid decoder with a different TP. + "pd-tp2-dep2": Deployment( + Instance(dp=2, kv_config=NIXL, prefix_caching=False), + prefill=Instance(tp=2, kv_config=NIXL), + ), + "pd-offload": Deployment( + Instance(kv_config=NIXL_OFFLOAD), + prefill=Instance(kv_config=NIXL_OFFLOAD), + offload=True, + ), +} + + +def _known_failure(name: str, mode: Mode) -> str | None: + deployment = DEPLOYMENTS[name] + if mode.spec and deployment.decode.dcp > 1: + return "FlashInfer MLA DCP decode breaks on ragged spec batches (#59392)" + return None + + +def _case(name: str, mode: Mode): + marks = [] + if (num_gpus := DEPLOYMENTS[name].num_gpus) > 1: + marks += multi_gpu_marks(num_gpus=num_gpus) + if reason := _known_failure(name, mode): + marks.append(pytest.mark.xfail(reason=reason, strict=True)) + return pytest.param(name, mode, marks=marks, id=f"{name}-{mode.name}") + + +def _block_size(url: str) -> int: + response = requests.get(f"{url}/metrics", timeout=30) + response.raise_for_status() + for family in text_string_to_metric_families(response.text): + for sample in family.samples: + if sample.name == "vllm:cache_config_info": + return lcm( + *( + int(value) + for key in ("block_size", "mamba_block_size") + if (value := sample.labels.get(key, "None")) != "None" + ) + ) + raise AssertionError("missing vllm:cache_config_info metric") + + +def _complete(url: str, prompt: list[int], salt: str, **extra) -> dict: + body = { + "model": MODEL, + "prompt": prompt, + "max_tokens": MAX_TOKENS, + "temperature": 0, + "logprobs": NUM_LOGPROBS, + "return_tokens_as_token_ids": True, + "cache_salt": salt, + **extra, + } + response = requests.post(f"{url}/v1/completions", json=body, timeout=600) + response.raise_for_status() + return response.json() + + +def _pd_complete( + prefill_url: str, decode_url: str, prompt: list[int], salt: str +) -> tuple[dict, dict]: + remote_decode = { + "do_remote_decode": True, + "do_remote_prefill": False, + "remote_engine_id": None, + "remote_block_ids": None, + "remote_host": None, + "remote_port": None, + } + prefill = _complete( + prefill_url, prompt, salt, max_tokens=1, kv_transfer_params=remote_decode + ) + decode = _complete( + decode_url, prompt, salt, kv_transfer_params=prefill["kv_transfer_params"] + ) + return prefill, decode + + +def _cached(response: dict) -> int: + return response["usage"]["prompt_tokens_details"]["cached_tokens"] + + +def _token_id(token: str) -> int: + return int(token.removeprefix("token_id:")) + + +def _tokens_text_logprobs(response: dict): + choice = response["choices"][0] + ids = [_token_id(t) for t in choice["logprobs"]["tokens"]] + top = [ + {_token_id(t): lp for t, lp in d.items()} + for d in choice["logprobs"]["top_logprobs"] + ] + return ids, choice["text"], top + + +def _reset_gpu_prefix_cache(url: str) -> None: + response = requests.post( + f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60 + ) + response.raise_for_status() + + +@dataclass(frozen=True) +class Turn: + prompt_len: int + cached: int + expected_cached: int + recompute_cached: int + output: tuple + recompute: tuple + + @property + def matches_recompute(self) -> bool: + return self.output[0] == self.recompute[0] + + +def _expected_cached(prev_prompt_len: int, hit_unit: int, spec: bool) -> int: + """A turn must reuse the previous prompt up to its last checkpoint. + + EAGLE-style drafts drop the trailing unit, so the checkpoint is one earlier. + """ + if not prev_prompt_len: + return 0 + checkpoint = (prev_prompt_len - 1) // hit_unit * hit_unit + return max(checkpoint - hit_unit, 0) if spec else checkpoint + + +@contextlib.contextmanager +def _serve(deployment: Deployment, mode: Mode) -> Iterator[list[RemoteOpenAIServer]]: + with contextlib.ExitStack() as stack: + # Sequential shutdown waits on memory still held by sibling servers. + servers: list[RemoteOpenAIServer] = [] + stack.callback(RemoteOpenAIServer.shutdown_many, servers) + next_gpu = 0 + for instance in deployment.instances: + env = {} + if deployment.offload: + # Enables /reset_prefix_cache. + env["VLLM_SERVER_DEV_MODE"] = "1" + if deployment.prefill is not None: + gpus = range(next_gpu, next_gpu + instance.num_gpus) + next_gpu += instance.num_gpus + env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus)) + env["VLLM_SSM_CONV_STATE_LAYOUT"] = "DS" + env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(get_open_port()) + servers.append( + RemoteOpenAIServer( + MODEL, + instance.args(mode), + env_dict=env, + max_wait_seconds=STARTUP_TIMEOUT, + ) + ) + yield servers + + +def _run_conversation( + deployment: Deployment, + servers: list[RemoteOpenAIServer], + mode: Mode, +) -> list[Turn]: + # Under P/D the prefiller computes the prompt, so hits are read from it. + compute_url, decode_url = servers[0].url_root, servers[-1].url_root + block_size = _block_size(compute_url) + hit_unit = mode.prefix_match_unit or block_size + rng = random.Random(0) + + def new_tokens(n: int) -> list[int]: + return [rng.randint(1000, 150000) for _ in range(n)] + + prompt = new_tokens(FIRST_PROMPT_BLOCKS * block_size + UNALIGNED_TAIL) + prev_len = 0 + turns = [] + for turn in range(NUM_TURNS): + if deployment.offload: + for server in servers: + _reset_gpu_prefix_cache(server.url_root) + if deployment.prefill is not None: + computed, output = _pd_complete(compute_url, decode_url, prompt, "conv") + else: + computed = output = _complete(decode_url, prompt, "conv") + recompute = _complete(decode_url, prompt, f"recompute-{turn}") + turns.append( + Turn( + prompt_len=len(prompt), + cached=_cached(computed), + expected_cached=_expected_cached(prev_len, hit_unit, mode.spec), + recompute_cached=_cached(recompute), + output=_tokens_text_logprobs(output), + recompute=_tokens_text_logprobs(recompute), + ) + ) + prev_len = len(prompt) + prompt = prompt + new_tokens(TURN_TOKENS) + return turns + + +def _format(turns: list[Turn]) -> str: + rows = ["turn prompt cached expected matches_recompute"] + rows += [ + f"{i:>4} {t.prompt_len:>6} {t.cached:>6} {t.expected_cached:>8} " + f"{t.matches_recompute}" + for i, t in enumerate(turns) + ] + return "\n".join(rows) + + +@pytest.mark.parametrize( + "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES] +) +def test_turns_reuse_prefix_and_match_recompute(name: str, mode: Mode) -> None: + deployment = DEPLOYMENTS[name] + with _serve(deployment, mode) as servers: + turns = _run_conversation(deployment, servers, mode) + table = _format(turns) + print(f"\n{name}-{mode.name}\n{table}") + + assert all(t.recompute_cached == 0 for t in turns), table + assert all(t.cached >= t.expected_cached for t in turns), table + # Dummy weights differ across TP sizes. + if len({i.tp for i in deployment.instances}) == 1: + check_logprobs_close( + outputs_0_lst=[t.recompute for t in turns], + outputs_1_lst=[t.output for t in turns], + name_0="recompute", + name_1=name, + ) diff --git a/tests/models/language/generation/test_granite.py b/tests/models/language/generation/test_granite.py index 207e798fef..f07839c29f 100644 --- a/tests/models/language/generation/test_granite.py +++ b/tests/models/language/generation/test_granite.py @@ -5,16 +5,15 @@ from transformers import GraniteConfig from vllm.model_executor.models.granite import granite_layer_attn_params -from ...utils import check_logprobs_close, check_transformers_version +from ...utils import check_logprobs_close -# model -> minimum transformers version, or None if unconstrained -MODELS = { +MODELS = [ # TODO(sang): Sliding window should be tested separately. - "ibm/PowerLM-3b": None, - "ibm/PowerMoE-3b": None, - "ibm-granite/granite-swash-2b": "5.15.1", - "ibm-granite/granite-swash-3b-a600m": "5.15.1", -} + "ibm/PowerLM-3b", + "ibm/PowerMoE-3b", + "ibm-granite/granite-swash-2b", + "ibm-granite/granite-swash-3b-a600m", +] @pytest.mark.parametrize("model", MODELS) @@ -31,8 +30,6 @@ def test_models( max_tokens: int, num_logprobs: int, ) -> None: - check_transformers_version(model, min_transformers_version=MODELS[model]) - with hf_runner(model, dtype=dtype) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit( example_prompts, max_tokens, num_logprobs diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 269dad05f4..32a9adda88 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -799,7 +799,6 @@ VLM_TEST_SETTINGS = { dtype="half", num_logprobs=10, patch_hf_runner=model_utils.ovis2_5_patch_hf_runner, - hf_model_kwargs={"revision": "refs/pr/5"}, ), "paddleocr_vl": VLMTestInfo( models=["PaddlePaddle/PaddleOCR-VL"], diff --git a/tests/models/multimodal/generation/test_nemotron_parse.py b/tests/models/multimodal/generation/test_nemotron_parse.py index 57e40430c6..69e41b21e3 100644 --- a/tests/models/multimodal/generation/test_nemotron_parse.py +++ b/tests/models/multimodal/generation/test_nemotron_parse.py @@ -104,16 +104,7 @@ def run_test( ) -@pytest.mark.parametrize( - "model", - [ - "nvidia/NVIDIA-Nemotron-Parse-v1.2", - # v2.0's checkpoint ties lm_head to the decoder's input embeddings - # instead of shipping a separate lm_head.weight tensor; this - # exercises that path (see tie_weights() in nemotron_parse.py). - "nvidia/NVIDIA-Nemotron-Parse-2.0", - ], -) +@pytest.mark.parametrize("model", ["nvidia/NVIDIA-Nemotron-Parse-v1.2"]) @pytest.mark.parametrize("dtype", ["bfloat16"]) @pytest.mark.parametrize("num_logprobs", [5]) def test_models( diff --git a/tests/models/multimodal/processing/test_deepseek_ocr.py b/tests/models/multimodal/processing/test_deepseek_ocr.py index 2bcf5e05d3..0523c34de0 100644 --- a/tests/models/multimodal/processing/test_deepseek_ocr.py +++ b/tests/models/multimodal/processing/test_deepseek_ocr.py @@ -12,12 +12,22 @@ Run with: pytest tests/models/multimodal/processing/test_deepseek_ocr.py -v """ +from types import SimpleNamespace + import pytest +import torch from PIL import Image from transformers import AutoTokenizer -from vllm.model_executor.models.deepseek_ocr import DeepseekOCRImagePixelInputs -from vllm.transformers_utils.processors.deepseek_ocr import DeepseekOCRProcessor +from vllm.model_executor.models.deepseek_ocr import ( + DeepseekOCRForCausalLM, + DeepseekOCRImagePixelInputs, +) +from vllm.model_executor.models.deepseek_ocr2 import DeepseekOCR2ForCausalLM +from vllm.transformers_utils.processors.deepseek_ocr import ( + DeepseekOCRProcessor, + ImageTransform, +) MODEL_ID = "deepseek-ai/DeepSeek-OCR" @@ -155,3 +165,80 @@ class TestDeepseekOCRInputValidation: def test_missing_images_raises_value_error(self, processor): with pytest.raises(ValueError, match="prompt and images"): processor(prompt="\nDescribe this image.", images=None) + + +def _balanced_normalized_pixels( + batch: int = 1, channels: int = 3, height: int = 64, width: int = 64 +) -> torch.Tensor: + """Build a pixel tensor whose values cancel to an exact sum of 0. + + Matches the normalized half-black / half-white case from the DeepSeek-OCR + processor (Normalize(0.5, 0.5) maps 0→-1 and 255→+1). + """ + assert width % 2 == 0 + pixel_values = torch.ones(batch, channels, height, width) + pixel_values[..., : width // 2] = -1.0 + assert pixel_values.sum().item() == 0.0 + return pixel_values + + +class TestDeepseekOCRZeroSumPixelsAccepted: + """Zero-sum normalized pixels must still produce validated image inputs. + + Returning None from ``_parse_and_validate_image_input`` makes + ``embed_multimodal`` return None and trips the worker encoder-output + sanity check, killing EngineCore. The zero-sum shortcut is therefore + removed; balanced images must parse like any other valid tensor. + """ + + def test_ocr_accepts_balanced_normalized_pixels(self): + pixel_values = _balanced_normalized_pixels() + images_crop = torch.zeros(0, 3, 40, 40) + images_spatial_crop = torch.tensor([[1, 1]]) + + image_input = DeepseekOCRForCausalLM._parse_and_validate_image_input( + None, + pixel_values=pixel_values, + images_crop=images_crop, + images_spatial_crop=images_spatial_crop, + ) + + assert image_input is not None + assert image_input.data is pixel_values + + def test_ocr2_accepts_balanced_normalized_pixels(self): + pixel_values = _balanced_normalized_pixels() + images_crop = torch.zeros(0, 3, 40, 40) + images_spatial_crop = torch.tensor([[1, 1]]) + model = SimpleNamespace(vision_config=SimpleNamespace(image_size=64)) + + image_input = DeepseekOCR2ForCausalLM._parse_and_validate_image_input( + model, + pixel_values=pixel_values, + images_crop=images_crop, + images_spatial_crop=images_spatial_crop, + ) + + assert image_input is not None + assert image_input.data is pixel_values + + def test_image_transform_half_black_half_white_sums_to_zero(self): + """A 1024×1024 half-black/half-white PNG through ImageTransform + (ToTensor + Normalize(0.5, 0.5)) yields an exact zero sum, and parse + must still accept it. + """ + image = Image.new("RGB", (1024, 1024), (0, 0, 0)) + image.paste(Image.new("RGB", (512, 1024), (255, 255, 255)), (512, 0)) + pixel_values = ImageTransform()(image).unsqueeze(0) + assert pixel_values.sum().item() == 0.0 + + images_crop = torch.zeros(0, 3, 640, 640) + images_spatial_crop = torch.tensor([[1, 1]]) + image_input = DeepseekOCRForCausalLM._parse_and_validate_image_input( + None, + pixel_values=pixel_values, + images_crop=images_crop, + images_spatial_crop=images_spatial_crop, + ) + assert image_input is not None + assert torch.equal(image_input.data, pixel_values) diff --git a/tests/models/multimodal/processing/test_glm5next.py b/tests/models/multimodal/processing/test_glm5next.py index 719154a2b4..111fb982cf 100644 --- a/tests/models/multimodal/processing/test_glm5next.py +++ b/tests/models/multimodal/processing/test_glm5next.py @@ -188,3 +188,54 @@ def test_mm_device_do_normalize(): torch.testing.assert_close( output, normalized_values.to(device), rtol=1e-5, atol=1e-6 ) + + +def _image_info(**kwargs): + ctx = build_model_context( + "zai-org/GLM-5.3-Flash", + limit_mm_per_prompt={"image": 1}, + **kwargs, + ) + return MULTIMODAL_REGISTRY.create_processor( + ctx.model_config, + tokenizer=ctx.tokenizer, + ).info + + +def test_image_encoder_cache_covers_full_token_budget(): + """The most-features probe must reach the processor's token ceiling. + + The inherited square probe refits to 89x89 = 7921 tokens under the + max_image_tokens=8000 budget, so the encoder cache came up short and + ordinary non-square images in 7922-8000 tokens were refused with + HTTP 400 (#59539). + """ + info = _image_info() + assert info.get_max_image_tokens() == 8000 + + # The shapes from the issue, with the token counts the processor + # actually produces; every one must fit the cache. + for width, height, expected_tokens in [ + (4032, 3024, 7931), # phone photo, 4:3 + (3840, 2160, 7973), # 4K frame, 16:9 + (3508, 2480, 7950), # A4 at 300 dpi + (2600, 2400, 7998), # 13:12 + (3000, 3000, 7921), # square worst case before the fix + ]: + num_tokens = info.get_num_image_tokens(image_width=width, image_height=height) + assert num_tokens == expected_tokens + assert num_tokens <= info.get_max_image_tokens() + + # The profiling dummy covers the real worst case, not the square one. + size = info.get_image_size_with_most_features() + assert size.width * size.height == 2240 * 2800 + + +def test_image_encoder_cache_follows_max_pixels_override(): + info = _image_info(mm_processor_kwargs={"max_pixels": 1568 * 100}) + assert info.get_max_image_tokens() == 100 + size = info.get_image_size_with_most_features() + assert ( + info.get_num_image_tokens(image_width=size.width, image_height=size.height) + == 100 + ) diff --git a/tests/models/multimodal/processing/test_transformers_audio.py b/tests/models/multimodal/processing/test_transformers_audio.py index 61517ea0af..1c1204d8e5 100644 --- a/tests/models/multimodal/processing/test_transformers_audio.py +++ b/tests/models/multimodal/processing/test_transformers_audio.py @@ -4,11 +4,8 @@ import numpy as np import pytest from vllm.config import ModelConfig -from vllm.model_executor.models.transformers.multimodal import LegacyMultiModalProcessor from vllm.multimodal import MULTIMODAL_REGISTRY -from .transformers_backend import PROCESSOR_CLASSES, create_processor - AUDIO_MODEL_SETTINGS = { "ibm-granite/granite-speech-3.3-2b": { "prompt": ( @@ -56,7 +53,6 @@ AUDIO_MODEL_SETTINGS = { } -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) @pytest.mark.parametrize( "model_id", [ @@ -78,10 +74,12 @@ AUDIO_MODEL_SETTINGS = { "zai-org/GLM-ASR-Nano-2512", ], ) -def test_audio_multimodal_processor(model_id, processor_cls): +def test_audio_multimodal_processor(model_id): settings = AUDIO_MODEL_SETTINGS[model_id] - mm_processor = create_processor(model_id, processor_cls) + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) audio = np.zeros(16000, dtype=np.float32) mm_data = {"audio": (audio, 16000)} @@ -112,13 +110,14 @@ def test_audio_multimodal_processor(model_id, processor_cls): ) -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) @pytest.mark.parametrize("separator", [" and ", ""]) -def test_audio_multiple_inputs(separator, processor_cls): +def test_audio_multiple_inputs(separator): """Multiple audios per prompt are each detected as a separate placeholder and multi-modal item by the Transformers modelling backend.""" model_id = "ibm-granite/granite-speech-3.3-2b" - mm_processor = create_processor(model_id, processor_cls) + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) audio_token = mm_processor.info.get_hf_processor().audio_token # One token per audio; the processor expands each to its placeholder run. @@ -128,21 +127,11 @@ def test_audio_multiple_inputs(separator, processor_cls): ) audios = [np.zeros(16000, dtype=np.float32), np.zeros(24000, dtype=np.float32)] - def process(): - return mm_processor( - prompt=prompt, - mm_items=mm_processor.info.parse_mm_data({"audio": audios}), - hf_processor_mm_kwargs={}, - ) - - # The legacy path reads placeholders off contiguous runs of the audio token, so - # it cannot tell adjacent ones apart and says so instead of merging them - if processor_cls is LegacyMultiModalProcessor and not separator: - with pytest.raises(ValueError, match="Separate them in the prompt"): - process() - return - - result = process() + result = mm_processor( + prompt=prompt, + mm_items=mm_processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) assert len(result["mm_placeholders"]["audio"]) == 2 assert len(result["mm_kwargs"]["audio"]) == 2 diff --git a/tests/models/multimodal/processing/test_transformers_image.py b/tests/models/multimodal/processing/test_transformers_image.py index 8f5a47be7b..26a3a592ca 100644 --- a/tests/models/multimodal/processing/test_transformers_image.py +++ b/tests/models/multimodal/processing/test_transformers_image.py @@ -5,23 +5,16 @@ from unittest.mock import patch import pytest from vllm.assets.image import ImageAsset -from vllm.model_executor.models.transformers.multimodal import ( - LegacyMultiModalProcessor, - OffsetsMultiModalProcessor, -) - -from .transformers_backend import ( - PROCESSOR_CLASSES, - create_cached_processor, - create_processor, - offsets_only, -) +from vllm.config import ModelConfig +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.cache import MultiModalProcessorOnlyCache -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) @pytest.mark.parametrize("model_id", ["llava-hf/llava-onevision-qwen2-0.5b-ov-hf"]) -def test_multimodal_processor(model_id, processor_cls): - mm_processor = create_processor(model_id, processor_cls) +def test_multimodal_processor(model_id): + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) image_pil = ImageAsset("cherry_blossom").pil_image mm_data = {"image": image_pil} @@ -63,9 +56,11 @@ def test_multimodal_processor(model_id, processor_cls): ) -def _process_two_images(processor_cls, separator: str): +def _process_two_images(separator: str): model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" - mm_processor = create_processor(model_id, processor_cls) + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) image = ImageAsset("cherry_blossom").pil_image prompt = ( @@ -80,30 +75,29 @@ def _process_two_images(processor_cls, separator: str): ) -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) -def test_image_multiple_inputs(processor_cls): +def test_image_multiple_inputs(): """Multiple images per prompt are each detected as a separate placeholder and multi-modal item by the Transformers modelling backend.""" - result = _process_two_images(processor_cls, separator="\n and ") + result = _process_two_images(separator="\n and ") assert len(result["mm_placeholders"]["image"]) == 2 assert len(result["mm_kwargs"]["image"]) == 2 -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) -def test_image_adjacent_inputs(processor_cls): +def test_image_adjacent_inputs(): """Adjacent images stay separate placeholders rather than merging into one.""" - result = _process_two_images(processor_cls, separator="") + result = _process_two_images(separator="") assert len(result["mm_placeholders"]["image"]) == 2 assert len(result["mm_kwargs"]["image"]) == 2 -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) -def test_batch_padding_removed_from_image_items(processor_cls): +def test_batch_padding_removed_from_image_items(): """Emu3 pads every image up to the largest in the batch, which would leave an item's data dependent on what it was processed with and so uncacheable.""" - mm_processor = create_processor("BAAI/Emu3-Chat-hf", processor_cls) + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model="BAAI/Emu3-Chat-hf", model_impl="transformers") + ) image_token = mm_processor.info.get_hf_processor().image_token images = [ @@ -128,53 +122,39 @@ def test_batch_padding_removed_from_image_items(processor_cls): assert len(shapes) == 2 -def _process_one_gemma3_image(processor_cls): - mm_processor = create_processor("google/gemma-3-4b-it", processor_cls) - hf_processor = mm_processor.info.get_hf_processor() - result = mm_processor( - prompt=f"{hf_processor.boi_token} What is this?", +def _process_one_gemma3_image(): + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model="google/gemma-3-4b-it", model_impl="transformers") + ) + boi_token = mm_processor.info.get_hf_processor().boi_token + return mm_processor( + prompt=f"{boi_token} What is this?", mm_items=mm_processor.info.parse_mm_data( {"image": ImageAsset("cherry_blossom").pil_image} ), hf_processor_mm_kwargs={}, ) - return hf_processor, result -@offsets_only def test_non_embedding_tokens_excluded_from_placeholders(): """Gemma3 wraps each image in text that carries no embeddings, which must be inside the placeholder range but masked out of it.""" - _, result = _process_one_gemma3_image(OffsetsMultiModalProcessor) + result = _process_one_gemma3_image() (placeholder,) = result["mm_placeholders"]["image"] assert placeholder.is_embed is not None assert 0 < int(placeholder.is_embed.sum()) < placeholder.length -def test_legacy_placeholders_hold_only_image_tokens(): - """The legacy path spans whatever `mm_token_type_ids` attributes to the image, - which for Gemma3 excludes the text wrapping it, unlike the replacement the offsets - path spans. Gemma3 is also the sharp case for the mask: its `image_token_id` is the - marker in the unexpanded prompt, not the token the expansion repeats.""" - hf_processor, result = _process_one_gemma3_image(LegacyMultiModalProcessor) - - (placeholder,) = result["mm_placeholders"]["image"] - assert placeholder.length == hf_processor.image_seq_length - prompt_ids = result["prompt_token_ids"] - covered = prompt_ids[placeholder.offset : placeholder.offset + placeholder.length] - assert set(covered) == {hf_processor.tokenizer.image_token_id} - - -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) -def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): +def test_tokens_structuring_an_image_are_masked_not_dropped(): """SmolVLM splits each image into tiles introduced by tokens carrying no embeddings. Those belong inside the placeholder and masked out, because the token count the processor reports is over the whole span. Idefics3 also refuses a prompt - holding `` when no images are passed, which is how the offsets path has to - tokenize it before splicing in the expansion.""" - mm_processor = create_processor( - "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls + holding `` when no images are passed, which is how the prompt has to be + tokenized before splicing in the expansion.""" + model_id = "HuggingFaceTB/SmolVLM-256M-Instruct" + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") ) result = mm_processor( prompt="What is this?", @@ -189,12 +169,13 @@ def test_tokens_structuring_an_image_are_masked_not_dropped(processor_cls): assert 0 < int(placeholder.is_embed.sum()) < placeholder.length -@offsets_only def test_missing_replacement_offsets_names_the_processor(): """A processor that reports no replacement offsets cannot be served, which must be said plainly rather than surfacing later as a field config mismatch.""" model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" - mm_processor = create_processor(model_id, OffsetsMultiModalProcessor) + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) hf_processor_cls = type(mm_processor.info.get_hf_processor()) hf_call = hf_processor_cls.__call__ @@ -216,11 +197,12 @@ def test_missing_replacement_offsets_names_the_processor(): ) -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) -def test_text_only_prompt(processor_cls): +def test_text_only_prompt(): """An image model still accepts a prompt with no images.""" model_id = "llava-hf/llava-onevision-qwen2-0.5b-ov-hf" - mm_processor = create_processor(model_id, processor_cls) + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) result = mm_processor( prompt="<|im_start|>user Hello!<|im_end|><|im_start|>assistant\n", @@ -232,12 +214,14 @@ def test_text_only_prompt(processor_cls): assert not result["mm_placeholders"] -@offsets_only def test_repeated_image_hits_the_processor_cache(): """Check that mm caching is actually working.""" - mm_processor, cache = create_cached_processor( - "llava-hf/llava-onevision-qwen2-0.5b-ov-hf", OffsetsMultiModalProcessor + model_config = ModelConfig( + model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf", model_impl="transformers" ) + model_config.multimodal_config.mm_processor_cache_gb = 4 + mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config) + cache = MultiModalProcessorOnlyCache(model_config) image = ImageAsset("cherry_blossom").pil_image def process(): @@ -255,7 +239,6 @@ def test_repeated_image_hits_the_processor_cache(): assert first["mm_hashes"] == second["mm_hashes"] -@offsets_only @pytest.mark.parametrize( ("model_id", "prompt"), [ @@ -267,7 +250,7 @@ def test_repeated_image_hits_the_processor_cache(): " and more text", marks=pytest.mark.xfail( reason="Emu3Processor prepends its BOS token only when images are " - "passed, so the unexpanded prompt the offsets path tokenizes never " + "passed, so the unexpanded prompt vLLM tokenizes never " "gets one. Fixed by huggingface/transformers#47924, unreleased.", strict=False, ), @@ -275,33 +258,40 @@ def test_repeated_image_hits_the_processor_cache(): ], ) def test_spliced_prompt_matches_hf_expansion(model_id, prompt): - """The offsets path splices the expansion into a prompt tokenized without any - multi-modal data, so its token ids have to come out the same as the ones the HF - processor produces itself, which is what the legacy path returns.""" - prompt_ids = [] - for processor_cls in (LegacyMultiModalProcessor, OffsetsMultiModalProcessor): - mm_processor = create_processor(model_id, processor_cls) - prompt_ids.append( - mm_processor( - prompt=prompt, - mm_items=mm_processor.info.parse_mm_data( - {"image": ImageAsset("cherry_blossom").pil_image} - ), - hf_processor_mm_kwargs={}, - )["prompt_token_ids"] - ) + """The prompt is tokenized without any multi-modal data and the expansion spliced + in, so its token ids have to come out the same as the ones the HF processor + produces itself.""" + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") + ) + info = mm_processor.info + image = ImageAsset("cherry_blossom").pil_image - legacy_ids, offsets_ids = prompt_ids - assert legacy_ids == offsets_ids + hf_processor = info.get_hf_processor() + prompt_ids = info.get_tokenizer().encode( + prompt, **info.default_tok_params.get_encode_kwargs() + ) + hf_ids = info.ctx.call_hf_processor( + hf_processor, + dict(text=hf_processor.decode(prompt_ids), images=[image]), + dict(truncation=False, add_special_tokens=False), + )["input_ids"][0].tolist() + + result = mm_processor( + prompt=prompt, + mm_items=info.parse_mm_data({"image": image}), + hf_processor_mm_kwargs={}, + ) + assert result["prompt_token_ids"] == hf_ids -@pytest.mark.parametrize("processor_cls", PROCESSOR_CLASSES) -def test_nested_image_fields_split_per_image(processor_cls): +def test_nested_image_fields_split_per_image(): """Idefics3 returns image fields with a leading batch dimension, putting the rows belonging to each image one dimension further in. Slicing the batch dimension instead handed the first image every row and the second an empty tensor.""" - mm_processor = create_processor( - "HuggingFaceTB/SmolVLM-256M-Instruct", processor_cls + model_id = "HuggingFaceTB/SmolVLM-256M-Instruct" + mm_processor = MULTIMODAL_REGISTRY.create_processor( + ModelConfig(model=model_id, model_impl="transformers") ) image = ImageAsset("cherry_blossom").pil_image result = mm_processor( diff --git a/tests/models/multimodal/processing/transformers_backend.py b/tests/models/multimodal/processing/transformers_backend.py deleted file mode 100644 index 869afe9456..0000000000 --- a/tests/models/multimodal/processing/transformers_backend.py +++ /dev/null @@ -1,51 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Helpers for testing both processors of the Transformers modelling backend.""" - -import pytest - -from vllm.config import ModelConfig -from vllm.model_executor.models.transformers.multimodal import ( - LegacyMultiModalProcessor, - MultiModalDummyInputsBuilder, - MultiModalProcessingInfo, - MultiModalProcessor, - OffsetsMultiModalProcessor, -) -from vllm.multimodal.cache import MultiModalProcessorOnlyCache -from vllm.multimodal.processing import InputProcessingContext -from vllm.tokenizers.registry import cached_tokenizer_from_config - -offsets_only = pytest.mark.skipif( - MultiModalProcessor is not OffsetsMultiModalProcessor, - reason="Replacement offsets are only used from transformers 5.15.0 onwards", -) - -PROCESSOR_CLASSES = [ - pytest.param(LegacyMultiModalProcessor, id="legacy"), - pytest.param(OffsetsMultiModalProcessor, id="offsets", marks=offsets_only), -] - - -def create_processor(model_id: str, processor_cls): - """Build a processor directly, because the registry only ever builds the one the - installed transformers version selects, leaving the other path untested.""" - model_config = ModelConfig(model=model_id, model_impl="transformers") - ctx = InputProcessingContext( - model_config, cached_tokenizer_from_config(model_config) - ) - info = MultiModalProcessingInfo(ctx) - return processor_cls(info, MultiModalDummyInputsBuilder(info)) - - -def create_cached_processor(model_id: str, processor_cls): - """Build a processor backed by a real multi-modal processor cache, and hand the - cache back so a test can check it was actually used.""" - model_config = ModelConfig(model=model_id, model_impl="transformers") - model_config.multimodal_config.mm_processor_cache_gb = 4 - ctx = InputProcessingContext( - model_config, cached_tokenizer_from_config(model_config) - ) - info = MultiModalProcessingInfo(ctx) - cache = MultiModalProcessorOnlyCache(model_config) - return processor_cls(info, MultiModalDummyInputsBuilder(info)), cache diff --git a/tests/models/registry.py b/tests/models/registry.py index 45d1a0a4f5..80ca856615 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -233,9 +233,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "CohereLabs/North-Mini-Code", trust_remote_code=True, is_available_online=False, - min_transformers_version="5.9.0", ), - "CwmForCausalLM": _HfExamplesInfo("facebook/cwm", min_transformers_version="4.58"), + "CwmForCausalLM": _HfExamplesInfo("facebook/cwm"), # FIXME: databricks/dbrx-instruct has been deleted "DbrxForCausalLM": _HfExamplesInfo( "databricks/dbrx-instruct", is_available_online=False @@ -266,9 +265,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct", trust_remote_code=True ), "Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"), - "ExaoneMoeForCausalLM": _HfExamplesInfo( - "LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.1.0" - ), + "ExaoneMoeForCausalLM": _HfExamplesInfo("LGAI-EXAONE/K-EXAONE-236B-A23B"), "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"), "FalconH1ForCausalLM": _HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"), "FlexOlmoForCausalLM": _HfExamplesInfo("allenai/Flex-reddit-2x7B-1T"), @@ -279,7 +276,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"), "Gemma4ForCausalLM": _HfExamplesInfo( "google/gemma-4-E2B-it", - min_transformers_version="5.0.0", ), "Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it"), "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"), @@ -287,11 +283,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Glm4MoeForCausalLM": _HfExamplesInfo("zai-org/GLM-4.5"), "Glm4MoeLiteForCausalLM": _HfExamplesInfo( "zai-org/GLM-4.7-Flash", - min_transformers_version="5.0.0", - ), - "GlmMoeDsaForCausalLM": _HfExamplesInfo( - "zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False ), + "GlmMoeDsaForCausalLM": _HfExamplesInfo("zai-org/GLM-5", is_available_online=False), "Glm5NextForCausalLM": _HfExamplesInfo("zai-org/GLM-5.3-Flash"), "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2"), "GPTBigCodeForCausalLM": _HfExamplesInfo( @@ -316,15 +309,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "GraniteMoeSharedForCausalLM": _HfExamplesInfo( "ibm-research/moe-7b-1b-active-shared-experts" ), - "GraniteMoeSWAForCausalLM": _HfExamplesInfo( - "ibm-granite/granite-swash-3b-a600m", min_transformers_version="5.15.1" - ), - "GraniteSWAForCausalLM": _HfExamplesInfo( - "ibm-granite/granite-swash-2b", min_transformers_version="5.15.1" - ), + "GraniteMoeSWAForCausalLM": _HfExamplesInfo("ibm-granite/granite-swash-3b-a600m"), + "GraniteSWAForCausalLM": _HfExamplesInfo("ibm-granite/granite-swash-2b"), "HrmTextForCausalLM": _HfExamplesInfo( "sapientinc/HRM-Text-1B", - min_transformers_version="5.9.0", ), "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"), "HunYuanMoEV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-A13B-Instruct"), @@ -333,7 +321,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "HYV4ForCausalLM": _HfExamplesInfo("tencent/Hy4-preview", trust_remote_code=True), "HyperCLOVAXForCausalLM": _HfExamplesInfo( "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", - min_transformers_version="5.9.0", ), "InklingForCausalLM": _HfExamplesInfo( "thinkingmachines/Inkling-NVFP4", @@ -353,9 +340,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "IQuestLoopCoderForCausalLM": _HfExamplesInfo( "IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct", trust_remote_code=True ), - "Jais2ForCausalLM": _HfExamplesInfo( - "inceptionai/Jais-2-8B-Chat", min_transformers_version="4.58" - ), + "Jais2ForCausalLM": _HfExamplesInfo("inceptionai/Jais-2-8B-Chat"), "JambaForCausalLM": _HfExamplesInfo( "ai21labs/AI21-Jamba-1.5-Mini", extras={ @@ -375,7 +360,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Lfm2ForCausalLM": _HfExamplesInfo("LiquidAI/LFM2-1.2B"), "Lfm2MoeForCausalLM": _HfExamplesInfo( "LiquidAI/LFM2-8B-A1B", - min_transformers_version="5.0.0", use_original_num_layers=True, # Initialize at least one MoE layer hf_overrides={"num_hidden_layers": 4}, @@ -422,7 +406,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { ), "MiniCPM4ForCausalLM": _HfExamplesInfo( "openbmb/MiniCPM4.1-8B", - min_transformers_version="4.56", max_transformers_version="4.57", transformers_version_reason={ "hf": "HF remote code imports removed `is_torch_fx_available`; " @@ -523,7 +506,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "Qwen3NextForCausalLM": _HfExamplesInfo( "Qwen/Qwen3-Next-80B-A3B-Instruct", extras={"tiny-random": "tiny-random/qwen3-next-moe"}, - min_transformers_version="4.56.3", ), "Rnj1ForCausalLM": _HfExamplesInfo( "EssentialAI/rnj-1-instruct", @@ -722,7 +704,6 @@ _TOKEN_CLASSIFICATION_EXAMPLE_MODELS = { ), "OpenAIPrivacyFilterForTokenClassification": _HfExamplesInfo( "openai/privacy-filter", - min_transformers_version="5.6.0.dev0", ), "RobertaForTokenClassification": _HfExamplesInfo( "Jean-Baptiste/roberta-large-ner-english" @@ -794,7 +775,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"), "AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo( "nvidia/audio-flamingo-3-hf", - min_transformers_version="5.3.0", transformers_version_reason={ "vllm": "Needs https://github.com/huggingface/transformers/pull/43538" }, @@ -823,22 +803,18 @@ _MULTIMODAL_EXAMPLE_MODELS = { "Cohere2VisionForConditionalGeneration": _HfExamplesInfo( "CohereLabs/command-a-vision-07-2025", extras={"command-a-plus": "CohereLabs/command-a-plus-05-2026-bf16"}, - min_transformers_version="5.9.0", ), "CohereCompassForConditionalGeneration": _HfExamplesInfo( "CohereLabs/North-Micro-Vision-Instruct", - min_transformers_version="5.16.0", ), "Cosmos3ForConditionalGeneration": _HfExamplesInfo( "nvidia/Cosmos3-Nano", extras={"super": "nvidia/Cosmos3-Super"}, max_model_len=4096, - min_transformers_version="4.57", ), "Cosmos3EdgeForConditionalGeneration": _HfExamplesInfo( "nvidia/Cosmos3-Edge", max_model_len=4096, - min_transformers_version="5.15", use_original_num_layers=True, hf_overrides={ "text_config": { @@ -887,7 +863,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "Exaone4_5_ForConditionalGeneration": _HfExamplesInfo( "LGAI-EXAONE/EXAONE-4.5-33B", - min_transformers_version="5.6.0", ), "FireRedASR2ForConditionalGeneration": _HfExamplesInfo( "allendou/FireRedASR2-LLM-vllm", @@ -918,19 +893,14 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "Gemma4ForConditionalGeneration": _HfExamplesInfo( "google/gemma-4-E2B-it", - min_transformers_version="5.5.0", ), - # TODO: update min_transformers_version when Gemma4 Unified lands in - # a stable transformers release. "Gemma4UnifiedForConditionalGeneration": _HfExamplesInfo( "google/gemma-4-12B-it", - min_transformers_version="5.8.0", is_available_online=False, ), "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it"), "GlmAsrForConditionalGeneration": _HfExamplesInfo( "zai-org/GLM-ASR-Nano-2512", - min_transformers_version="5.0.0", ), "Granite4VisionForConditionalGeneration": _HfExamplesInfo( "ibm-granite/granite-vision-4.1-4b", @@ -943,7 +913,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "GraniteSpeechPlusForConditionalGeneration": _HfExamplesInfo( "ibm-granite/granite-speech-4.1-2b-plus", - min_transformers_version="5.8.0", ), "GLM4VForCausalLM": _HfExamplesInfo( "zai-org/glm-4v-9b", @@ -957,7 +926,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { "Glm4vMoeForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.5V"), "GlmOcrForConditionalGeneration": _HfExamplesInfo( "zai-org/GLM-OCR", - min_transformers_version="5.1.0", ), "Glm5NextForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-5.3-Flash"), "H2OVLChatModel": _HfExamplesInfo( @@ -974,7 +942,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "HunYuanVLForConditionalGeneration": _HfExamplesInfo( "tencent/HunyuanOCR", - min_transformers_version="5.13", ), "Idefics3ForConditionalGeneration": _HfExamplesInfo( "HuggingFaceM4/Idefics3-8B-Llama3", @@ -1087,7 +1054,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "Lfm2VlForConditionalGeneration": _HfExamplesInfo( "LiquidAI/LFM2-VL-450M", - min_transformers_version="5.0.0", ), "Llama4ForConditionalGeneration": _HfExamplesInfo( "meta-llama/Llama-4-Scout-17B-16E-Instruct", @@ -1154,7 +1120,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "MiniCPMV4_6ForConditionalGeneration": _HfExamplesInfo( "openbmb/MiniCPM-V-4_6", - min_transformers_version="5.7.0", ), "MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo( "MiniMaxAI/MiniMax-M3", @@ -1201,7 +1166,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { "Molmo2ForConditionalGeneration": _HfExamplesInfo( "allenai/Molmo2-8B", extras={"olmo": "allenai/Molmo2-O-7B"}, - min_transformers_version="4.51", trust_remote_code=True, # required by current PrefixLM implementation max_num_batched_tokens=31872, @@ -1296,19 +1260,15 @@ _MULTIMODAL_EXAMPLE_MODELS = { trust_remote_code=True, max_transformers_version="4.57", transformers_version_reason={ - "vllm": "Custom processor code is not compatible with Transformers v5." + "hf": "HF remote code reads `is_parallelizable`, which was removed " + "in Transformers v5." }, ), "Ovis2_6ForCausalLM": _HfExamplesInfo( "AIDC-AI/Ovis2.6-2B", is_available_online=False, trust_remote_code=True ), "Ovis2_6_MoeForCausalLM": _HfExamplesInfo( - "AIDC-AI/Ovis2.6-30B-A3B", - trust_remote_code=True, - max_transformers_version="4.57", - transformers_version_reason={ - "vllm": "Custom processor code is not compatible with Transformers v5." - }, + "AIDC-AI/Ovis2.6-30B-A3B", trust_remote_code=True ), "PaddleOCRVLForConditionalGeneration": _HfExamplesInfo( "PaddlePaddle/PaddleOCR-VL", @@ -1353,7 +1313,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "QianfanOCRForConditionalGeneration": _HfExamplesInfo( "baidu/Qianfan-OCR", - min_transformers_version="5.6.0", ), "Qwen2AudioForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen2-Audio-7B-Instruct" @@ -1368,12 +1327,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { "Qwen3VLForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3-VL-4B-Instruct", max_model_len=4096, - min_transformers_version="4.57", ), "Qwen3VLMoeForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3-VL-30B-A3B-Instruct", max_model_len=4096, - min_transformers_version="4.57", ), "Qwen3_5ForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3.5-0.8B", @@ -1391,23 +1348,19 @@ _MULTIMODAL_EXAMPLE_MODELS = { "Qwen3OmniMoeForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3-Omni-30B-A3B-Instruct", max_model_len=4096, - min_transformers_version="4.57", ), "Qwen3ASRForConditionalGeneration": _HfExamplesInfo( "Qwen/Qwen3-ASR-0.6B", max_model_len=4096, - min_transformers_version="4.57", ), "Qwen3ASRRealtimeGeneration": _HfExamplesInfo( "Qwen/Qwen3-ASR-0.6B", max_model_len=4096, - min_transformers_version="4.57", hf_overrides={"architectures": ["Qwen3ASRRealtimeGeneration"]}, ), "Qwen3ASRForcedAlignerForTokenClassification": _HfExamplesInfo( "Qwen/Qwen3-ForcedAligner-0.6B", max_model_len=4096, - min_transformers_version="4.57", hf_overrides={"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}, ), "RForConditionalGeneration": _HfExamplesInfo("YannQi/R-4B", trust_remote_code=True), @@ -1456,7 +1409,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { ), "VibeVoiceAsrForConditionalGeneration": _HfExamplesInfo( "microsoft/VibeVoice-ASR-HF", - min_transformers_version="5.13.0", ), "VoxtralForConditionalGeneration": _HfExamplesInfo( "mistralai/Voxtral-Mini-3B-2507", @@ -1530,7 +1482,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { use_original_num_layers=True, # DFlash requires all layers max_model_len=8192, # Reduce for CI max_num_seqs=32, - min_transformers_version="4.56.3", # Required for Qwen3Next ), "MuseGlimmerAssistantModel": _HfExamplesInfo( "meta-models/Muse-Glimmer-30B", @@ -1576,7 +1527,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { speculative_model="Qwen/Qwen3-Omni-30B-A3B-DSpark", is_available_online=False, use_original_num_layers=True, - min_transformers_version="4.57", ), "Gemma4DSparkModel": _HfExamplesInfo( "google/gemma-4-12B-it", @@ -1737,7 +1687,6 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "google/gemma-4-E4B-it", speculative_model="google/gemma-4-E4B-it-assistant", trust_remote_code=True, - min_transformers_version="5.8.0", ), "ErnieMTPModel": _HfExamplesInfo( "baidu/ERNIE-4.5-21B-A3B-PT", @@ -1747,13 +1696,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "ExaoneMoeMTP": _HfExamplesInfo( "LGAI-EXAONE/K-EXAONE-236B-A23B", speculative_model="LGAI-EXAONE/K-EXAONE-236B-A23B", - min_transformers_version="5.1.0", enable_prefix_caching=False, ), "Exaone4_5_MTP": _HfExamplesInfo( "LGAI-EXAONE/EXAONE-4.5-33B", speculative_model="LGAI-EXAONE/EXAONE-4.5-33B", - min_transformers_version="5.6.0", ), "ExtractHiddenStatesModel": _HfExamplesInfo( "Qwen/Qwen3-8B", @@ -1766,12 +1713,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { "Glm4MoeLiteMTPModel": _HfExamplesInfo( "zai-org/GLM-4.7-Flash", speculative_model="zai-org/GLM-4.7-Flash", - min_transformers_version="5.0.0", ), "GlmOcrMTPModel": _HfExamplesInfo( "zai-org/GLM-OCR", speculative_model="zai-org/GLM-OCR", - min_transformers_version="5.1.0", ), "Glm5NextMTPModel": _HfExamplesInfo( "zai-org/GLM-5.3-Flash", @@ -1842,9 +1787,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { trust_remote_code=True, is_available_online=False, ), - "Qwen3NextMTP": _HfExamplesInfo( - "Qwen/Qwen3-Next-80B-A3B-Instruct", min_transformers_version="4.56.3" - ), + "Qwen3NextMTP": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct"), "Qwen3_5MTP": _HfExamplesInfo( "Qwen/Qwen3.5-0.8B", speculative_model="Qwen/Qwen3.5-0.8B", @@ -1865,29 +1808,20 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { } _TRANSFORMERS_BACKEND_MODELS = { - "TransformersEmbeddingModel": _HfExamplesInfo( - "BAAI/bge-base-en-v1.5", min_transformers_version="5.0.0" - ), + "TransformersEmbeddingModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"), "TransformersForSequenceClassification": _HfExamplesInfo( "papluca/xlm-roberta-base-language-detection", - min_transformers_version="5.0.0", ), "TransformersForCausalLM": _HfExamplesInfo( "hmellor/Ilama-3.2-1B", trust_remote_code=True ), "TransformersMultiModalForCausalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"), - "TransformersMoEForCausalLM": _HfExamplesInfo( - "allenai/OLMoE-1B-7B-0924", min_transformers_version="5.0.0" - ), + "TransformersMoEForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924"), "TransformersMultiModalMoEForCausalLM": _HfExamplesInfo( - "Qwen/Qwen3-VL-30B-A3B-Instruct", min_transformers_version="5.0.0" - ), - "TransformersMoEEmbeddingModel": _HfExamplesInfo( - "Qwen/Qwen3-30B-A3B", min_transformers_version="5.0.0" - ), - "TransformersMoEForSequenceClassification": _HfExamplesInfo( - "Qwen/Qwen3-30B-A3B", min_transformers_version="5.0.0" + "Qwen/Qwen3-VL-30B-A3B-Instruct" ), + "TransformersMoEEmbeddingModel": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"), + "TransformersMoEForSequenceClassification": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"), "TransformersMultiModalEmbeddingModel": _HfExamplesInfo("google/gemma-3-4b-it"), "TransformersMultiModalForSequenceClassification": _HfExamplesInfo( "google/gemma-3-4b-it" diff --git a/tests/models/test_deepseek_v4_fi_moe_ep.py b/tests/models/test_deepseek_v4_fi_moe_ep.py index 99d661e755..7687a6786f 100644 --- a/tests/models/test_deepseek_v4_fi_moe_ep.py +++ b/tests/models/test_deepseek_v4_fi_moe_ep.py @@ -1,116 +1,90 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the flashinfer moe_ep backend plumbing. -Everything here runs without a GPU or a flashinfer install: the flashinfer -modules the helpers import lazily are replaced with capture fakes. -""" - -import sys -from dataclasses import dataclass, field -from types import ModuleType, SimpleNamespace -from typing import Any +from dataclasses import replace +from types import SimpleNamespace import pytest import torch +from tests.kernels.moe.utils import make_dummy_moe_config from vllm.config.kernel import ( FLASHINFER_MOE_EP_BACKENDS, + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, MEGA_MOE_BACKENDS, + NATIVE_MEGA_MOE_BACKENDS, + PASSTHROUGH_ALL2ALL_BACKEND, + bind_passthrough_all2all_backend, validate_flashinfer_moe_ep_model, ) -from vllm.utils.flashinfer_moe_ep import ( - _E2M1_LUT, - FI_MOE_EP_BACKEND_SPECS, - _dequant_fp4_ue8m0_gran32, - build_fi_mega_config, - fi_moe_ep_backend_spec, - make_fi_moe_ep_bootstrap, - megakernel_runtime_requirements, +from vllm.engine.arg_utils import EngineArgs +from vllm.model_executor.layers.fused_moe import flashinfer_moe_ep as fi_ep +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, ) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + RoutingMethodType, + nvfp4_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( + FlashInferMoeEpExperts, + epilogue_from_quant_config, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + FLASHINFER_MOE_EP_MXFP4_BACKENDS, + Mxfp4MoeBackend, + make_mxfp4_moe_quant_config, + map_mxfp4_backend, + mxfp4_round_up_hidden_size_and_intermediate_size, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( + backend_to_kernel_cls as mxfp4_backend_to_kernel_cls, +) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + convert_to_nvfp4_moe_kernel_format, + map_nvfp4_backend, +) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + backend_to_kernel_cls as nvfp4_backend_to_kernel_cls, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import ( + PassThroughPrepareAndFinalize, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Static, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.utils.argparse_utils import FlexibleArgumentParser -@dataclass -class _FakeBootstrapConfig: - world_size: int - rank: int - process_group: Any = None - auto_bootstrap: bool = True - device: int | None = field(default=None, kw_only=True) +def test_flashinfer_backends_are_megakernels_outside_the_native_model_path(): + assert FLASHINFER_MOE_EP_BACKENDS <= MEGA_MOE_BACKENDS + assert FLASHINFER_MOE_EP_BACKENDS.isdisjoint(NATIVE_MEGA_MOE_BACKENDS) -@dataclass -class _FakeDeepGemmMegaMoeConfig: - intermediate_size: int - top_k: int - activation_clamp: float | None - fast_math: bool - - -@dataclass -class _FakeNvfp4CutedslMegaMoeConfig: - intermediate_size: int - top_k: int - activation_clamp: float | None - fast_math: bool - - -@dataclass -class _FakeMegaConfig: - megakernel: Any - preprocess_weights: bool - quantize_input: bool - - -@pytest.fixture -def fake_flashinfer(monkeypatch): - """Install a minimal fake flashinfer.moe_ep for the lazy imports.""" - moe_ep = ModuleType("flashinfer.moe_ep") - core = ModuleType("flashinfer.moe_ep.core") - runtime = ModuleType("flashinfer.moe_ep.core.runtime") - flashinfer = ModuleType("flashinfer") - fake_attrs: dict[ModuleType, dict[str, Any]] = { - moe_ep: { - "BootstrapConfig": _FakeBootstrapConfig, - "DeepGemmMegaMoeConfig": _FakeDeepGemmMegaMoeConfig, - "Nvfp4CutedslMegaMoeConfig": _FakeNvfp4CutedslMegaMoeConfig, - "MegaConfig": _FakeMegaConfig, - "core": core, - }, - runtime: {"TORCH_DIST": "torch_dist", "NVSHMEM": "nvshmem"}, - flashinfer: {"moe_ep": moe_ep}, - core: {"runtime": runtime}, - } - for mod, attrs in fake_attrs.items(): - for attr, value in attrs.items(): - setattr(mod, attr, value) - - for name, mod in { - "flashinfer": flashinfer, - "flashinfer.moe_ep": moe_ep, - "flashinfer.moe_ep.core": core, - "flashinfer.moe_ep.core.runtime": runtime, - }.items(): - monkeypatch.setitem(sys.modules, name, mod) - return moe_ep - - -def test_fi_backend_strings_are_registered_mega_moe_backends(): - assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS - assert FLASHINFER_MOE_EP_BACKENDS < MEGA_MOE_BACKENDS - - -@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS)) -def test_fi_moe_ep_backend_rejected_for_non_dsv4(moe_backend): - """An FI moe_ep backend with a non-DSv4 model must fail at config time - instead of silently falling through to the generic FusedMoE path.""" +def test_only_deep_gemm_backend_is_dsv4_specific(): + validate_flashinfer_moe_ep_model( + FLASHINFER_MOE_EP_CUTEDSL, + ["MixtralForCausalLM"], + ) with pytest.raises(ValueError, match="only supported for DeepSeek-V4"): - validate_flashinfer_moe_ep_model(moe_backend, ["MixtralForCausalLM"]) - - -@pytest.mark.parametrize("moe_backend", sorted(FLASHINFER_MOE_EP_BACKENDS)) -def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend): - validate_flashinfer_moe_ep_model(moe_backend, ["DeepseekV4ForCausalLM"]) + validate_flashinfer_moe_ep_model( + FLASHINFER_MOE_EP_DEEP_GEMM, + ["MixtralForCausalLM"], + ) + validate_flashinfer_moe_ep_model( + FLASHINFER_MOE_EP_DEEP_GEMM, + ["DeepseekV4ForCausalLM"], + ) @pytest.mark.parametrize( @@ -118,8 +92,7 @@ def test_fi_moe_ep_backend_accepted_for_dsv4(moe_backend): [["KimiK3ForConditionalGeneration"], ["MixtralForCausalLM"]], ) def test_native_deep_gemm_mega_moe_not_arch_gated(architectures): - """VLLM's own deep_gemm mega path is not DSv4-only (Kimi K3 uses it); - models validate their own constraints at construction time.""" + """Models validate native deep_gemm mega constraints at construction time.""" validate_flashinfer_moe_ep_model("deep_gemm_mega_moe", architectures) @@ -128,11 +101,7 @@ def test_non_fi_backend_ignores_architectures(): @pytest.mark.parametrize("moe_backend", sorted(MEGA_MOE_BACKENDS)) -def test_all_mega_backends_get_sequence_parallel_moe(moe_backend): - """Every mega backend must qualify for sequence-parallel MoE at - TP>1/EP: the predicate once matched only the native backend string, - which silently ran the fi backends full-batch with an all-reduce on - every rank — 0.42-0.65x native e2e at TP8.""" +def test_token_sharding_backends_enable_dsv4_sequence_parallel(moe_backend: str): from vllm.models.deepseek_v4.nvidia.model import _use_sequence_parallel vllm_config = SimpleNamespace( @@ -147,113 +116,262 @@ def test_all_mega_backends_get_sequence_parallel_moe(moe_backend): assert _use_sequence_parallel(vllm_config) -def test_fi_moe_ep_backend_spec_kernel_and_nvshmem_contract(): - dg = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm") - assert dg.megakernel == "deep_gemm_mega" - assert not dg.needs_nvshmem - - cd = fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl") - assert cd.megakernel == "nvfp4_cutedsl" - assert cd.needs_nvshmem - - with pytest.raises(ValueError, match="not a flashinfer moe_ep backend"): - fi_moe_ep_backend_spec("deep_gemm_mega_moe") +def test_dsv4_requests_pre_fc2_router_weight_placement(): + dsv4 = SimpleNamespace(routing_method=RoutingMethodType.DeepseekV4) + default = SimpleNamespace(routing_method=RoutingMethodType.Default) + assert fi_ep.apply_topk_in_fc1(dsv4) + assert not fi_ep.apply_topk_in_fc1(default) -def test_megakernel_runtime_requirements(fake_flashinfer): - dg = megakernel_runtime_requirements( - fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_deep_gemm") +def test_backend_specs_preserve_weight_format_contracts(): + cutedsl = fi_ep.flashinfer_moe_ep_backend_spec(FLASHINFER_MOE_EP_CUTEDSL) + assert cutedsl.kernel == "cutedsl" + assert cutedsl.weight_formats == frozenset({"nvfp4", "mxfp4"}) + + deep_gemm = fi_ep.flashinfer_moe_ep_backend_spec(FLASHINFER_MOE_EP_DEEP_GEMM) + assert deep_gemm.kernel == "deep_gemm" + assert deep_gemm.weight_formats == frozenset({"mxfp4"}) + + +def _megakernel_moe(moe_backend: str, **overrides) -> SimpleNamespace: + fields = dict( + moe_backend=moe_backend, + moe_parallel_config=SimpleNamespace( + use_ep=True, use_batched_activation_format=False + ), + is_act_and_mul=True, + activation=MoEActivation.SILU, + has_hash_routing=False, + routing_method=RoutingMethodType.DeepseekV4, + router_logits_dtype=torch.float32, + hidden_dim=256, + is_lora_enabled=False, + skip_final_all_reduce=False, + in_dtype=torch.bfloat16, + has_bias=False, + swiglu_alpha=None, + swiglu_beta=None, ) - assert dg == frozenset({"torch_dist"}) + fields.update(overrides) + return SimpleNamespace(**fields) - cd = megakernel_runtime_requirements( - fi_moe_ep_backend_spec("flashinfer_moe_ep_mega_cutedsl") + +def _is_supported(moe, weight_key, activation_key) -> tuple[bool, str]: + supported, reason = FlashInferMoeEpExperts.is_supported_config( + FlashInferMoeEpExperts, + moe, + weight_key, + activation_key, + mk.FusedMoEActivationFormat.Standard, ) - assert cd == frozenset({"torch_dist", "nvshmem"}) + return supported, reason or "" -def test_bootstrap_pins_the_device_vllm_bound(fake_flashinfer, monkeypatch): - """The runtime must not rederive the device from LOCAL_RANK/rank: under a - remapped CUDA_VISIBLE_DEVICES that ordinal points at the wrong GPU - (CUDA_ERROR_ILLEGAL_ADDRESS in the weight transforms). vLLM passes the - device it already bound via BootstrapConfig.device.""" - import vllm.utils.flashinfer_moe_ep as mod - - pg = object() +@pytest.fixture +def megakernel_host(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + """A vLLM config without DBO/EPLB/weight transfer, on a supported device.""" + config = SimpleNamespace( + weight_transfer_config=None, + parallel_config=SimpleNamespace(enable_dbo=False, enable_eplb=False), + ) + monkeypatch.setattr(fi_ep, "get_current_vllm_config", lambda: config) monkeypatch.setattr( - mod, - "get_ep_group", - lambda: SimpleNamespace(world_size=4, rank_in_group=2, device_group=pg), + FlashInferMoeEpExperts, "_supports_current_device", staticmethod(lambda: True) ) - monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3) - - bootstrap = make_fi_moe_ep_bootstrap() - - assert bootstrap.world_size == 4 - assert bootstrap.rank == 2 - assert bootstrap.process_group is pg - assert bootstrap.auto_bootstrap is False - assert bootstrap.device == 3 + return config -def test_build_fi_mega_config_selects_kernel_config(fake_flashinfer): - dg = build_fi_mega_config( - intermediate_size=2048, - top_k=8, - activation_clamp=7.0, - megakernel="deep_gemm_mega", +def test_megakernel_requires_expert_parallel(megakernel_host): + """The oracle rejects the megakernel through the experts' support check.""" + moe = _megakernel_moe(FLASHINFER_MOE_EP_CUTEDSL) + assert _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) == (True, "") + + moe.moe_parallel_config.use_ep = False + supported, reason = _is_supported(moe, kNvfp4Static, kNvfp4Dynamic) + assert not supported and "parallel config" in reason + + +def test_megakernel_reports_its_own_constraints(megakernel_host): + """Constraints the framework does not know about surface as reasons.""" + _, reason = _is_supported( + _megakernel_moe(FLASHINFER_MOE_EP_CUTEDSL), kNvfp4Static, None ) - assert isinstance(dg.megakernel, _FakeDeepGemmMegaMoeConfig) - assert dg.megakernel.intermediate_size == 2048 - assert dg.megakernel.top_k == 8 - assert dg.megakernel.activation_clamp == 7.0 - assert dg.preprocess_weights and dg.quantize_input + assert "A16 activations" in reason - cd = build_fi_mega_config( - intermediate_size=2048, - top_k=8, - activation_clamp=None, - megakernel="nvfp4_cutedsl", + deep_gemm = _megakernel_moe( + FLASHINFER_MOE_EP_DEEP_GEMM, routing_method=RoutingMethodType.Renormalize ) - assert isinstance(cd.megakernel, _FakeNvfp4CutedslMegaMoeConfig) + _, reason = _is_supported(deep_gemm, kMxfp4Static, None) + assert "routing method" in reason - with pytest.raises(ValueError, match="Unsupported fi_moe_ep megakernel"): - build_fi_mega_config( - intermediate_size=2048, - top_k=8, - activation_clamp=None, - megakernel="deep_gemm", + megakernel_host.parallel_config.enable_eplb = True + _, reason = _is_supported( + _megakernel_moe(FLASHINFER_MOE_EP_DEEP_GEMM), kMxfp4Static, None + ) + assert reason.endswith("EPLB") + + +def test_megakernels_are_oracle_backends(): + """Quant methods reach the megakernel through the NVFP4/MXFP4 oracles.""" + nvfp4 = map_nvfp4_backend(FLASHINFER_MOE_EP_CUTEDSL) + assert nvfp4 is NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL + assert nvfp4_backend_to_kernel_cls(nvfp4) == [FlashInferMoeEpExperts] + + assert map_mxfp4_backend(FLASHINFER_MOE_EP_DEEP_GEMM) == [ + Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM + ] + for backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + assert mxfp4_backend_to_kernel_cls(backend) == [FlashInferMoeEpExperts] + # The megakernel takes the model's own dimensions. + assert mxfp4_round_up_hidden_size_and_intermediate_size( + backend, 2880, 2880 + ) == (2880, 2880) + + quant_config = make_mxfp4_moe_quant_config( + Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, w1_scale=None, w2_scale=None + ) + assert quant_config is not None and quant_config.use_nvfp4_w4a4 + + +def test_megakernel_is_a_modular_kernel_with_pass_through_stages(): + """The megakernel plugs into ``FusedMoEKernel`` like any other experts impl. + + Routing stays outside (top-k ids in), dispatch and combine happen inside, + so prepare passes tokens through untouched and finalize has nothing to do. + """ + moe = make_dummy_moe_config( + num_experts=8, experts_per_token=2, hidden_dim=256, intermediate_size=128 + ) + quant_config = FusedMoEQuantConfig.make("nvfp4", weight_dtype="nvfp4") + kernel = mk.FusedMoEKernel( + PassThroughPrepareAndFinalize(), + FlashInferMoeEpExperts(moe, quant_config), + ) + + assert not kernel.is_monolithic + assert kernel.prepare_finalize.output_is_reduced() + # Generic stages: the int32 top-k ids the megakernel wants are its concern. + assert kernel.prepare_finalize.topk_indices_dtype() is None + assert kernel.fused_experts.expects_unquantized_inputs + assert isinstance( + kernel.fused_experts.finalize_weight_and_reduce_impl(), TopKWeightAndReduceNoOP + ) + assert kernel.fused_experts.workspace_shapes( + 16, 256, 256, 2, 8, 8, None, MoEActivation.SILU + ) == ((0,), (0,), (16, 256)) + + tokens = torch.zeros(4, 256, dtype=torch.bfloat16) + prepared = kernel.prepare_finalize.prepare( + tokens, + torch.ones(4, 2), + torch.zeros(4, 2, dtype=torch.int32), + 8, + None, + False, + quant_config, + ) + assert prepared[0] is tokens and all(item is None for item in prepared[1:]) + + +def test_epilogue_takes_the_weight_global_scales_from_the_quant_config(): + """The canonical ``g1/g2_alphas`` become the per-expert fc1/fc2 alphas and + keep aliasing the layer's tensors, so EPLB permutations reach the kernel.""" + g1_alphas = torch.tensor([1.0, 2.0]) + g2_alphas = torch.tensor([3.0, 4.0]) + quant_config = nvfp4_moe_quant_config( + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, + a1_gscale=torch.ones(2), + a2_gscale=torch.ones(2), + w1_scale=torch.ones(2, 4, 1), + w2_scale=torch.ones(2, 2, 1), + ) + + epilogue = epilogue_from_quant_config(quant_config) + + assert epilogue.fc1_alpha is not None and epilogue.fc2_alpha is not None + assert epilogue.fc1_alpha.data_ptr() == g1_alphas.data_ptr() + assert epilogue.fc2_alpha.data_ptr() == g2_alphas.data_ptr() + # Activations are quantized dynamically: no norm constants. + assert epilogue.input_norm_const == 1.0 and epilogue.fc1_norm_const is None + + +def test_epilogue_rejects_unloaded_weight_global_scales(): + quant_config = nvfp4_moe_quant_config( + g1_alphas=torch.tensor([1.0, float("nan")]), + g2_alphas=torch.ones(2), + a1_gscale=torch.ones(2), + a2_gscale=torch.ones(2), + w1_scale=torch.ones(2, 4, 1), + w2_scale=torch.ones(2, 2, 1), + ) + with pytest.raises(ValueError, match="g1_alphas"): + epilogue_from_quant_config(quant_config) + + +def test_mxfp4_checkpoints_keep_the_default_epilogue(): + """MXFP4 has no per-tensor global scales; the megakernel needs no alphas.""" + quant_config = FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4") + assert epilogue_from_quant_config(quant_config) == fi_ep.FlashInferMoeEpEpilogue() + + +def test_nvfp4_kernel_format_passes_megakernel_weights_through(): + """The megakernel takes the canonical tensors as they are; the oracle must + not reject or reshape them.""" + tensors = tuple(torch.arange(4.0).reshape(2, 2) + i for i in range(8)) + converted = convert_to_nvfp4_moe_kernel_format( + NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, + SimpleNamespace(), + *tensors, + is_act_and_mul=True, + ) + assert all(out is src for out, src in zip(converted, tensors)) + + +def test_passthrough_all2all_backend_selects_the_pass_through_stages(): + """``all2all_backend="passthrough"`` is how the oracles get the + pass-through prepare/finalize.""" + moe = make_dummy_moe_config( + num_experts=8, experts_per_token=2, hidden_dim=256, intermediate_size=128 + ) + moe.moe_parallel_config = replace( + moe.moe_parallel_config, all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND + ) + assert not moe.moe_parallel_config.use_all2all_kernels + + prepare_finalize = maybe_make_prepare_finalize( + moe, FusedMoEQuantConfig.make("nvfp4", weight_dtype="nvfp4") + ) + + assert isinstance(prepare_finalize, PassThroughPrepareAndFinalize) + + +def test_passthrough_all2all_backend_is_not_a_cli_choice(): + """Users cannot pick the passthrough all2all backend explicitly.""" + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + with pytest.raises(SystemExit): + parser.parse_args(["--all2all-backend", PASSTHROUGH_ALL2ALL_BACKEND]) + parser.parse_args(["--all2all-backend", "deepep_low_latency"]) + + +def test_megakernel_moe_backend_binds_the_all2all_backend(): + """A megakernel deployment gets its all2all backend by default.""" + parallel = SimpleNamespace(all2all_backend="allgather_reducescatter") + bind_passthrough_all2all_backend(FLASHINFER_MOE_EP_CUTEDSL, parallel) + assert parallel.all2all_backend == PASSTHROUGH_ALL2ALL_BACKEND + bind_passthrough_all2all_backend(FLASHINFER_MOE_EP_CUTEDSL, parallel) + + with pytest.raises(ValueError, match="cannot be paired"): + bind_passthrough_all2all_backend( + FLASHINFER_MOE_EP_DEEP_GEMM, + SimpleNamespace(all2all_backend="deepep_low_latency"), + ) + with pytest.raises(ValueError, match="only valid with"): + bind_passthrough_all2all_backend( + "flashinfer_trtllm", + SimpleNamespace(all2all_backend=PASSTHROUGH_ALL2ALL_BACKEND), ) - -def test_ckpt_uses_nvfp4_experts_reads_moe_quant_algo(): - from vllm.models.deepseek_v4.nvidia.fi_moe import ckpt_uses_nvfp4_experts - - nvfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo="NVFP4")) - assert ckpt_uses_nvfp4_experts(nvfp4) - - mxfp4 = SimpleNamespace(quant_config=SimpleNamespace(moe_quant_algo=None)) - assert not ckpt_uses_nvfp4_experts(mxfp4) - - no_algo = SimpleNamespace(quant_config=SimpleNamespace()) - assert not ckpt_uses_nvfp4_experts(no_algo) - - -def test_dequant_fp4_ue8m0_gran32_decodes_lut_and_scales(): - """One 32-element scale group per row: low nibble is the even element, - high nibble the odd one, ue8m0 scale applies to the whole group.""" - packed = torch.arange(32, dtype=torch.uint8).reshape(2, 16) - sf = torch.tensor([[127], [128]], dtype=torch.uint8) # 2**0, 2**1 - - out = _dequant_fp4_ue8m0_gran32(packed, sf) - - assert out.shape == (2, 32) - assert out.dtype == torch.bfloat16 - expected = torch.empty(2, 32) - for row in range(2): - for col in range(16): - byte = int(packed[row, col]) - expected[row, 2 * col] = _E2M1_LUT[byte & 0x0F] - expected[row, 2 * col + 1] = _E2M1_LUT[byte >> 4] - expected[row] *= 2.0**row - assert torch.equal(out, expected.to(torch.bfloat16)) + untouched = SimpleNamespace(all2all_backend="deepep_low_latency") + bind_passthrough_all2all_backend("flashinfer_trtllm", untouched) + assert untouched.all2all_backend == "deepep_low_latency" diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index f87dd75fb2..9ee0212fba 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -657,7 +657,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( moe = DeepseekV4MoE.__new__(DeepseekV4MoE) torch.nn.Module.__init__(moe) - moe.use_mega_moe = True + moe.use_native_mega_moe = True moe.gate = FakeGate() moe.experts = FakeExperts() moe.shared_experts = FakeSharedExperts() diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 3e70450f2c..f77832d608 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -241,7 +241,7 @@ def test_v41_dspark_loads_linear_scales( linear_scale_name=scale_name, pad_shared_expert=False, model=SimpleNamespace( - layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], + layers=[SimpleNamespace(ffn=SimpleNamespace(use_native_mega_moe=False))], confidence_head=None, ), named_parameters=lambda: [(runtime_name, param)], diff --git a/tests/models/transformers/test_backend.py b/tests/models/transformers/test_backend.py index 2064f01e0c..001ee5ef12 100644 --- a/tests/models/transformers/test_backend.py +++ b/tests/models/transformers/test_backend.py @@ -185,17 +185,6 @@ def test_sinks(hf_runner: type[HfRunner], vllm_runner: type[VllmRunner]) -> None def test_mla(vllm_runner: type[VllmRunner], example_prompts: list[str]) -> None: - import transformers - from packaging.version import Version - - installed = Version(transformers.__version__) - required = Version("5.15.0.dev0") - if installed < required: - pytest.skip( - "MLA models with the Transformers modeling backend require " - f"transformers>={required}, but got {installed}" - ) - model = "hmellor/tiny-random-DeepseekV2ForCausalLM" args = (example_prompts, 32, 5) kwargs: dict[str, Any] = {"max_model_len": 2048, "enforce_eager": True} diff --git a/tests/models/utils.py b/tests/models/utils.py index 003ef1a078..4e98a981d4 100644 --- a/tests/models/utils.py +++ b/tests/models/utils.py @@ -609,17 +609,3 @@ def dummy_hf_overrides( ) return hf_config - - -def check_transformers_version( - model: str, - min_transformers_version: str | None = None, - max_transformers_version: str | None = None, -): - from .registry import _HfExamplesInfo - - return _HfExamplesInfo( - model, - min_transformers_version=min_transformers_version, - max_transformers_version=max_transformers_version, - ).check_transformers_version(on_fail="skip") diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 3723c6be30..1d6ced0d48 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8DynamicTokenSym, kFp8Static128BlockSym, kFp8StaticTensorSym, + kInt4Static32, kInt8StaticChannelSym, kMxfp8Dynamic, ) @@ -46,6 +47,12 @@ def test_quant_spec_string_representation_uses_quantization_name(): assert str(quant_spec(weight="mxfp4")) == "mxfp4" +def test_quant_spec_resolves_groupwise_int4_weight(): + spec = QuantSpec(weight="int4_per_group_32") + assert spec.weight == kInt4Static32 + assert spec.activation is None + + def test_quant_spec_rejects_unknown_name(): with pytest.raises(ValueError, match="unknown quantization name"): quant_spec(weight="not_a_real_format") diff --git a/tests/rocm/test_moe_padding.py b/tests/rocm/test_moe_padding.py new file mode 100644 index 0000000000..178b2eb9f4 --- /dev/null +++ b/tests/rocm/test_moe_padding.py @@ -0,0 +1,426 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the `VLLM_ROCM_MOE_PADDING` weight memory-layout padding trick. + +`VLLM_ROCM_MOE_PADDING` (see `UnquantizedFusedMoEMethod._maybe_pad_weight`) +is a memory-stride trick, not a logical-shape one: when eligible, it enlarges +a weight tensor's storage so per-expert allocations land farther apart in +HBM. Shape and values are unaffected, but `stride(-2)` increases by +`num_pad` elements -- that's the whole mechanism. For the **AITER** path +tested here, the padded storage never persists on `w13_weight`/`w2_weight`: +`convert_to_unquantized_kernel_format` reallocates via +`rocm_aiter_ops.shuffle_weights`. This is AITER-specific -- on TRITON-on-ROCm, +`.contiguous()` is skipped specifically to keep this same padding, so its +persisted parameters *do* stay padded. + +This file verifies: +1. `_maybe_pad_weight` in isolation (any ROCm backend, AITER not required): + grows storage/`stride(-2)` only when eligible (512-byte-aligned) and the + flag is on -- covering no padding, hidden-only, intermediate-only, and + both-dimensions-padded, since w13 and w2 key their alignment gate off + different dims; a true no-op (identical object) otherwise. +2. End to end, AITER only (per #54966's "Test padding" ask): + `AiterExperts.apply()` matches an independent reference regardless of + the flag, and persisted parameter storage size is unaffected by it. + +Logical `hidden_dim_unpadded`/`intermediate_size_per_partition_unpadded` +padding and HIP-graph token-padding are separate mechanisms, covered in +#59333 and #59334. + +See https://github.com/vllm-project/vllm/issues/54966 ("Test padding"). +""" + +import pytest +import torch +import torch.nn.functional as F + +import vllm.envs as envs +from vllm._aiter_ops import is_aiter_found_and_supported +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.fused_moe.expert_map_manager import ( + ExpertMapManager, +) +from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( + UnquantizedFusedMoEMethod, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +aiter_available = is_aiter_found_and_supported() + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="VLLM_ROCM_MOE_PADDING only takes effect on ROCm", +) +requires_aiter = pytest.mark.skipif( + not aiter_available, + reason="requires AITER to exercise AiterExperts.apply()", +) + +DEVICE = current_platform.device_type +DTYPE = torch.bfloat16 +NUM_EXPERTS = 4 +TOPK = 2 +NUM_TOKENS = 16 +# Mirrors `_maybe_pad_weight`'s `num_pad` constant in unquantized_fused_moe_method.py. +PAD_BYTES = 256 +# 256 bf16 elems * 2 bytes = 512 bytes: eligible for `_maybe_pad_weight` on +# both w13 and w2. +HIDDEN_SIZE = 256 +INTERMEDIATE_SIZE = 256 + + +def _set_padding_env(monkeypatch: pytest.MonkeyPatch, padding: bool) -> None: + """Sets the flag only; `_maybe_pad_weight` reads it lazily, no AITER needed.""" + monkeypatch.setenv("VLLM_ROCM_MOE_PADDING", "1" if padding else "0") + + +def _set_padding_env_with_aiter(monkeypatch: pytest.MonkeyPatch, padding: bool) -> None: + """Also enables AITER, for tests routing through RoutedExperts's backend + selection.""" + _set_padding_env(monkeypatch, padding) + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.refresh_env_variables() + + +def _make_moe_config( + hidden_size: int, + intermediate_size: int, + dtype: torch.dtype = DTYPE, +) -> FusedMoEConfig: + return FusedMoEConfig( + num_experts=NUM_EXPERTS, + experts_per_token=TOPK, + hidden_dim=hidden_size, + intermediate_size=intermediate_size, + num_local_experts=NUM_EXPERTS, + num_logical_experts=NUM_EXPERTS, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=dtype, + device=DEVICE, + routing_method=RoutingMethodType.TopK, + max_num_tokens=NUM_TOKENS, + ) + + +def _make_routed_experts(hidden_size: int, intermediate_size: int) -> RoutedExperts: + moe_config = _make_moe_config(hidden_size, intermediate_size) + expert_map_manager = ExpertMapManager( + max_num_batched_tokens=NUM_TOKENS, + top_k=TOPK, + global_num_experts=NUM_EXPERTS, + num_redundant_experts=0, + num_expert_group=None, + moe_parallel_config=moe_config.moe_parallel_config, + placement_strategy="linear", + enable_eplb=False, + ) + return RoutedExperts( + "experts", + DTYPE, + moe_config, + quant_config=None, + expert_map_manager=expert_map_manager, + ) + + +def _expert_weight_iterator(seed: int, hidden_size: int, intermediate_size: int): + set_random_seed(seed) + for expert_id in range(NUM_EXPERTS): + yield ( + f"{expert_id}.gate_proj.weight", + torch.randn( + intermediate_size, + hidden_size, + device=DEVICE, + dtype=DTYPE, + ), + ) + yield ( + f"{expert_id}.up_proj.weight", + torch.randn( + intermediate_size, + hidden_size, + device=DEVICE, + dtype=DTYPE, + ), + ) + yield ( + f"{expert_id}.down_proj.weight", + torch.randn( + hidden_size, + intermediate_size, + device=DEVICE, + dtype=DTYPE, + ), + ) + + +def _load_and_process_weights( + layer: RoutedExperts, seed: int, hidden_size: int, intermediate_size: int +) -> None: + loaded = set( + layer.load_weights( + _expert_weight_iterator(seed, hidden_size, intermediate_size) + ) + ) + assert {"w13_weight", "w2_weight"} <= loaded + layer.quant_method.process_weights_after_loading(layer) + torch.accelerator.synchronize() + + +def _make_static_inputs( + hidden_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + set_random_seed(2026) + x = torch.randn( + NUM_TOKENS, + hidden_size, + device=DEVICE, + dtype=DTYPE, + ) + topk_ids = torch.arange(NUM_TOKENS, device=DEVICE, dtype=torch.int64).unsqueeze(1) + topk_ids = torch.cat( + [topk_ids % NUM_EXPERTS, (topk_ids + 1) % NUM_EXPERTS], + dim=1, + ) + topk_weights = torch.full( + (NUM_TOKENS, TOPK), + 1.0 / TOPK, + device=DEVICE, + dtype=DTYPE, + ) + return x, topk_weights, topk_ids + + +def _forward( + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> torch.Tensor: + return layer.forward_modular( + x, + topk_weights, + topk_ids, + shared_experts=None, + shared_experts_input=None, + ) + + +def _assert_backend_is_aiter(layer: RoutedExperts) -> None: + selected = layer.quant_method.unquantized_backend + assert selected is UnquantizedMoeBackend.AITER, ( + f"expected AiterExperts to be selected, got backend={selected!r}" + ) + + +# --- 1. `_maybe_pad_weight` in isolation ----------------------------------- + + +@pytest.mark.parametrize( + "padding,hidden_size,intermediate_size,dtype,expect_padded", + [ + # eligible shape (512 bytes): padding fires iff the env var is on. + (True, HIDDEN_SIZE, INTERMEDIATE_SIZE, DTYPE, True), + (False, HIDDEN_SIZE, INTERMEDIATE_SIZE, DTYPE, False), + # ineligible shape (400 bytes, not 512-aligned): never fires. + (True, 200, 200, DTYPE, False), + # fp32: alignment threshold and num_pad both scale with + # element_size() (128 fp32 elems * 4 bytes == 512). + (True, 128, 128, torch.float32, True), + ], + ids=["eligible-on", "eligible-off", "ineligible-on", "eligible-on-fp32"], +) +def test_maybe_pad_weight_transparent( + monkeypatch: pytest.MonkeyPatch, + default_vllm_config, + padding: bool, + hidden_size: int, + intermediate_size: int, + dtype: torch.dtype, + expect_padded: bool, +) -> None: + """Shape/values never change; storage only grows (`stride(-2) += num_pad`) + when eligible and the flag is on. Otherwise it's a true no-op (identical + object).""" + assert default_vllm_config is not None + _set_padding_env(monkeypatch, padding) + assert envs.VLLM_ROCM_MOE_PADDING is padding + + moe_config = _make_moe_config(hidden_size, intermediate_size, dtype=dtype) + method = UnquantizedFusedMoEMethod(moe_config) + + original = torch.randn( + NUM_EXPERTS, + 2 * intermediate_size, + hidden_size, + device=DEVICE, + dtype=dtype, + ) + result = method._maybe_pad_weight(original) + + assert result.shape == original.shape + assert torch.equal(result, original), "padding must not alter weight values" + + if expect_padded: + num_pad = PAD_BYTES // original.element_size() + assert result.stride(-1) == 1 + assert result.stride(-2) == original.stride(-2) + num_pad + assert result.untyped_storage().nbytes() > original.untyped_storage().nbytes() + assert result.data_ptr() != original.data_ptr() + else: + assert result is original + + +@pytest.mark.parametrize( + "hidden_size,intermediate_size,expect_w13_padded,expect_w2_padded", + [ + # only hidden_size aligned: only w13 (last dim hidden_size) eligible. + (HIDDEN_SIZE, 200, True, False), + # only intermediate_size aligned: only w2 (last dim + # intermediate_size) eligible. + (200, HIDDEN_SIZE, False, True), + ], + ids=["only-w13-eligible", "only-w2-eligible"], +) +def test_maybe_pad_weight_asymmetric_w13_w2_eligibility( + monkeypatch: pytest.MonkeyPatch, + default_vllm_config, + hidden_size: int, + intermediate_size: int, + expect_w13_padded: bool, + expect_w2_padded: bool, +) -> None: + """w13 and w2 key their alignment gate off different dims (hidden_size + vs. intermediate_size), so they're independently eligible.""" + assert default_vllm_config is not None + _set_padding_env(monkeypatch, True) + + moe_config = _make_moe_config(hidden_size, intermediate_size) + method = UnquantizedFusedMoEMethod(moe_config) + + w13 = torch.randn( + NUM_EXPERTS, 2 * intermediate_size, hidden_size, device=DEVICE, dtype=DTYPE + ) + w2 = torch.randn( + NUM_EXPERTS, hidden_size, intermediate_size, device=DEVICE, dtype=DTYPE + ) + + w13_result = method._maybe_pad_weight(w13) + w2_result = method._maybe_pad_weight(w2) + + assert (w13_result is not w13) == expect_w13_padded + assert (w2_result is not w2) == expect_w2_padded + + +def _reference_moe_forward( + seed: int, + hidden_size: int, + intermediate_size: int, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> torch.Tensor: + """A minimal, independent (no AITER) reference MoE forward, so the E2E + test below can't pass merely because both flag settings are wrong the + same way.""" + weights: dict[int, dict[str, torch.Tensor]] = {} + for name, tensor in _expert_weight_iterator(seed, hidden_size, intermediate_size): + expert_id_str, proj_name = name.split(".", 1) + weights.setdefault(int(expert_id_str), {})[proj_name] = tensor + + output = torch.zeros_like(x) + for expert_id, proj in weights.items(): + gate = F.silu(x @ proj["gate_proj.weight"].T) + up = x @ proj["up_proj.weight"].T + expert_out = (gate * up) @ proj["down_proj.weight"].T + + mask = topk_ids == expert_id + weight = (topk_weights * mask).sum(dim=1, keepdim=True) + output = output + weight * expert_out + return output + + +# --- 2. End-to-end numerical transparency through AiterExperts ------------- + + +@requires_aiter +@torch.inference_mode() +def test_aiter_moe_padding_numerically_transparent( + monkeypatch: pytest.MonkeyPatch, + default_vllm_config, + workspace_init, +) -> None: + """`AiterExperts.apply()` output must match an independent reference + (not just itself across flag settings) for an eligible weight shape, + regardless of `VLLM_ROCM_MOE_PADDING`. Persisted parameter storage size + must also be unaffected by the flag (see module docstring).""" + assert default_vllm_config is not None + assert workspace_init is None + + outputs: dict[bool, torch.Tensor] = {} + storage_nbytes: dict[bool, tuple[int, int]] = {} + for padding in (True, False): + _set_padding_env_with_aiter(monkeypatch, padding) + + with torch.device(DEVICE): + layer = _make_routed_experts(HIDDEN_SIZE, INTERMEDIATE_SIZE) + _load_and_process_weights( + layer, + seed=1, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + ) + _assert_backend_is_aiter(layer) + storage_nbytes[padding] = ( + layer.w13_weight.data.untyped_storage().nbytes(), + layer.w2_weight.data.untyped_storage().nbytes(), + ) + + x, topk_weights, topk_ids = _make_static_inputs(HIDDEN_SIZE) + outputs[padding] = _forward(layer, x, topk_weights, topk_ids).clone() + + assert storage_nbytes[True] == storage_nbytes[False], ( + "VLLM_ROCM_MOE_PADDING must not change the persisted parameters' " + "storage size for the AITER path" + ) + + reference = _reference_moe_forward( + seed=1, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + x=x, + topk_weights=topk_weights, + topk_ids=topk_ids, + ) + # Same AITER kernel, byte-identical weights and inputs either way (only + # a transient weight-loading buffer differs): bitwise-identical results. + torch.testing.assert_close(outputs[True], outputs[False], atol=0, rtol=0) + + # Cosine similarity, not elementwise assert_close, for the reference + # comparison: two independently computed bf16 GEMM chains can have + # near-zero elements where rounding noise blows up elementwise relative + # error. Same pattern as test_b12x.py / test_mxfp4_moe.py. + for padded, output in outputs.items(): + cos_sim = F.cosine_similarity( + output.flatten().float(), reference.flatten().float(), dim=0 + ) + assert cos_sim > 0.99, ( + f"padding={padded} output diverged from reference " + f"(cosine similarity {cos_sim:.4f})" + ) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 7b234884c5..9acc1f7b74 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -2,8 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch -from vllm.outputs import RequestOutput +from vllm.outputs import ( + ClassificationRequestOutput, + EmbeddingRequestOutput, + PoolingOutput, + PoolingRequestOutput, + RequestError, + RequestOutput, + ScoringRequestOutput, +) pytestmark = pytest.mark.cpu_test @@ -19,3 +28,36 @@ def test_request_output_forward_compatible(): example_arg_added_in_new_version="some_value", ) assert output is not None + + +@pytest.mark.parametrize( + ("output_type", "empty_value"), + [ + (EmbeddingRequestOutput, []), + (ClassificationRequestOutput, []), + (ScoringRequestOutput, 0.0), + ], +) +@pytest.mark.skip_global_cleanup +def test_specialized_pooling_outputs_preserve_request_error(output_type, empty_value): + error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + base_output = PoolingRequestOutput( + request_id="request", + outputs=PoolingOutput(torch.empty(0)), + prompt_token_ids=[1, 2], + num_cached_tokens=0, + finished=True, + error=error, + ) + + output = output_type.from_base(base_output) + + assert output.error is error + value = getattr(output.outputs, "embedding", None) + value = getattr(output.outputs, "probs", value) + value = getattr(output.outputs, "score", value) + assert value == empty_value diff --git a/tests/tokenizers_/test_registry.py b/tests/tokenizers_/test_registry.py index a4ee4a755e..312b0c03a9 100644 --- a/tests/tokenizers_/test_registry.py +++ b/tests/tokenizers_/test_registry.py @@ -10,7 +10,6 @@ from transformers import AutoConfig from transformers.models.auto.configuration_auto import CONFIG_MAPPING from vllm.tokenizers import TokenizerLike -from vllm.tokenizers.hf import CachedHfTokenizer from vllm.tokenizers.registry import ( TokenizerRegistry, cached_get_tokenizer, @@ -93,54 +92,6 @@ def test_resolve_tokenizer_args_rejects_hf_mistral_format_true(): ) -@pytest.mark.parametrize( - ("transformers_version", "model_type", "native_filename", "expect_error"), - [ - pytest.param("5.14.0", "mistral", "tekken.json", True), - pytest.param("5.15.0", "mistral", "tekken.json", False), - pytest.param("5.14.0", "mistral3", "tekken.json", True), - pytest.param("5.15.0", "mistral3", "tekken.json", False), - pytest.param("5.14.0", "mixtral", "tokenizer.model.v1", True), - pytest.param("5.15.0", "mixtral", "tokenizer.model.v1", False), - pytest.param("5.14.0", "mistral", None, False), - pytest.param("5.14.0", "qwen3_5_moe", "tekken.json", False), - ], -) -def test_get_tokenizer_handles_hf_mistral_transformers_compatibility( - tmp_path: Path, - transformers_version: str, - model_type: str, - native_filename: str | None, - expect_error: bool, -): - if native_filename is not None: - (tmp_path / native_filename).touch() - - with ( - patch( - "vllm.tokenizers.registry.get_config", - return_value=SimpleNamespace(model_type=model_type), - ), - patch( - "vllm.tokenizers.registry.transformers.__version__", - transformers_version, - ), - patch.object( - CachedHfTokenizer, - "from_pretrained", - return_value=SimpleNamespace(is_fast=True), - ) as from_pretrained, - ): - if expect_error: - with pytest.raises(ValueError, match="requires transformers>=5.15.0"): - get_tokenizer(str(tmp_path), tokenizer_mode="hf") - else: - tokenizer = get_tokenizer(str(tmp_path), tokenizer_mode="hf") - assert tokenizer.is_fast is True - - assert from_pretrained.called == (not expect_error) - - def test_customized_tokenizer(): TokenizerRegistry.register("test_tokenizer", __name__, TestTokenizer.__name__) diff --git a/tests/tool_parsers/test_structural_tag_registry.py b/tests/tool_parsers/test_structural_tag_registry.py index 50f8751f90..d5ab3e4f1a 100644 --- a/tests/tool_parsers/test_structural_tag_registry.py +++ b/tests/tool_parsers/test_structural_tag_registry.py @@ -3,11 +3,14 @@ import json from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock import pytest from openai.types.responses import ToolChoiceFunction +from openai.types.responses.tool_choice_allowed import ToolChoiceAllowed from xgrammar import Grammar, StructuralTag +from xgrammar import get_model_structural_tag as get_xgrammar_builtin_structural_tag from xgrammar.testing import _is_grammar_accept_string import vllm.envs as envs @@ -1311,6 +1314,536 @@ def test_tool_strict_level_from_name(): ToolStrictLevel.from_name("off") +# --------------------------------------------------------------------------- +# GLM-4.7 structural tag: builtin content for strict tools, shallow otherwise +# --------------------------------------------------------------------------- + + +def _glm47_tools() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["city"], + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "run_command", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + }, + }, + ), + ] + + +def _glm47_arg(key: str, value: str) -> str: + return f"{key}{value}" + + +def _glm47_call(name: str, *args: str) -> str: + return f"{name}{''.join(args)}" + + +def _glm47_shallow_grammar( + tool_choice, tools=None, reasoning: bool = False, *, auto: bool = False +): + # auto needs the FUNCTION level to lift the non-strict auto gate. + tag = get_model_structural_tag( + model="glm_4_7", + tools=tools if tools is not None else _glm47_tools(), + tool_choice=tool_choice, + reasoning=reasoning, + strict_level=ToolStrictLevel.FUNCTION if auto else ToolStrictLevel.AUTO, + ) + assert isinstance(tag, StructuralTag) + return Grammar.from_structural_tag(tag) + + +@pytest.mark.parametrize( + "body", + [ + # plain text, no tool call + "Hello there", + # text followed by a call + "Sure." + _glm47_call("get_weather", _glm47_arg("city", "Paris")), + # call only, required arg present + _glm47_call("get_weather", _glm47_arg("city", "Paris")), + # all arguments omitted + _glm47_call("get_weather"), + # repeated argument + _glm47_call( + "get_weather", _glm47_arg("city", "Paris"), _glm47_arg("city", "Rome") + ), + # arguments in reverse declaration order + _glm47_call( + "get_weather", _glm47_arg("days", "3"), _glm47_arg("city", "Paris") + ), + # enum-constrained value + _glm47_call("get_weather", _glm47_arg("unit", "celsius")), + # string values with spaces and a bare '<' + _glm47_call("run_command", _glm47_arg("command", "grep 'ab")), + _glm47_call("get_weather", _glm47_arg("city", "ab")), + # integer-typed value given non-numeric text + _glm47_call("get_weather", _glm47_arg("days", "abc")), + # enum-typed value outside the enum + _glm47_call("get_weather", _glm47_arg("unit", "kelvin")), + # unterminated tool call + "get_weather" + _glm47_arg("city", "Paris"), + ], +) +def test_glm47_non_strict_auto_rejects_invalid(body: str): + grammar = _glm47_shallow_grammar("auto", auto=True) + assert not _is_grammar_accept_string(grammar, body) + + +def test_glm47_non_strict_required_requires_at_least_one_call(): + grammar = _glm47_shallow_grammar("required") + + assert not _is_grammar_accept_string(grammar, "Just answering.") + assert _is_grammar_accept_string( + grammar, _glm47_call("get_weather", _glm47_arg("city", "Paris")) + ) + assert _is_grammar_accept_string( + grammar, + _glm47_call("run_command", _glm47_arg("command", "ls")) + + _glm47_call("get_weather", _glm47_arg("city", "Paris")), + ) + + +def test_glm47_non_strict_forced_emits_single_named_call(): + grammar = _glm47_shallow_grammar( + ChatCompletionNamedToolChoiceParam( + type="function", + function=ChatCompletionNamedFunction(name="get_weather"), + ) + ) + + assert _is_grammar_accept_string( + grammar, _glm47_call("get_weather", _glm47_arg("city", "Paris")) + ) + assert not _is_grammar_accept_string(grammar, "No call here") + assert not _is_grammar_accept_string(grammar, _glm47_call("run_command")) + assert not _is_grammar_accept_string( + grammar, + _glm47_call("get_weather", _glm47_arg("city", "Paris")) + + _glm47_call("run_command"), + ) + + +def test_glm47_non_strict_reasoning_gates_on_think_close(): + grammar = _glm47_shallow_grammar("required", reasoning=True) + call = _glm47_call("get_weather", _glm47_arg("city", "Paris")) + + assert _is_grammar_accept_string(grammar, "thinking..." + call) + assert not _is_grammar_accept_string(grammar, call) + + +@pytest.mark.parametrize( + ("parameters", "key"), + [ + ( + { + "$ref": "#/$defs/args", + "$defs": { + "args": { + "type": "object", + "properties": {"city": {"type": "string"}}, + } + }, + }, + "city", + ), + ( + { + "anyOf": [ + {"type": "object", "properties": {"city": {"type": "string"}}}, + {"type": "null"}, + ] + }, + "city", + ), + ( + {"allOf": [{"type": "object", "properties": {"city": {"type": "string"}}}]}, + "city", + ), + ({"type": "object", "patternProperties": {"^arg": {"type": "string"}}}, "city"), + ({"type": "object"}, "city"), + ({}, "city"), + ( + { + "type": "object", + "properties": {"city": {"type": "string"}}, + "additionalProperties": True, + }, + "undeclared_key", + ), + ], +) +def test_glm47_non_plain_object_schemas_keep_free_arguments(parameters, key): + # Schemas that are not a plain object with declared properties fall back + # to the builtin's non-strict envelope content, which accepts any key. + tools = [ + ChatCompletionToolsParam( + type="function", function={"name": "f", "parameters": parameters} + ) + ] + tag = get_model_structural_tag( + model="glm_4_7", + tools=tools, + tool_choice="required", + reasoning=False, + ) + assert isinstance(tag, StructuralTag) + grammar = Grammar.from_structural_tag(tag) + call = _glm47_call("f", _glm47_arg(key, "Paris")) + assert _is_grammar_accept_string(grammar, call) + + +def _glm47_strict_tools() -> list[ChatCompletionToolsParam]: + return [ + ChatCompletionToolsParam( + type="function", + function={ + "name": "get_weather", + "strict": True, + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["city"], + "additionalProperties": False, + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={ + "name": "run_command", + "strict": True, + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + }, + }, + ), + ChatCompletionToolsParam( + type="function", + function={"name": "no_params", "strict": True, "parameters": None}, + ), + ] + + +@pytest.mark.parametrize( + ("tool_choice", "reasoning"), + [ + ("auto", False), + ("required", False), + ("named", False), + ("required", True), + ("allowed", False), + ], +) +def test_glm47_strict_only_matches_xgrammar_builtin(tool_choice: str, reasoning: bool): + # Strict-only requests stay byte-identical to the xgrammar builtin. + strict_tools = _glm47_strict_tools() + dumped = [ + {"type": "function", "function": t.function.model_dump(exclude_none=True)} + for t in strict_tools + ] + if tool_choice == "named": + choice = ChatCompletionNamedToolChoiceParam( + type="function", + function=ChatCompletionNamedFunction(name="run_command"), + ) + builtin_choice: Any = {"type": "function", "function": {"name": "run_command"}} + elif tool_choice == "allowed": + choice = ToolChoiceAllowed( + type="allowed_tools", + mode="auto", + tools=[{"type": "function", "function": {"name": "get_weather"}}], + ) + builtin_choice = { + "type": "allowed_tools", + "allowed_tools": { + "mode": "auto", + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + }, + } + else: + choice = tool_choice + builtin_choice = tool_choice + + ours = get_model_structural_tag( + model="glm_4_7", + tools=strict_tools, + tool_choice=choice, + reasoning=reasoning, + ) + expected = get_xgrammar_builtin_structural_tag( + model="glm_4_7", + tools=dumped, + tool_choice=builtin_choice, + reasoning=reasoning, + ) + + assert isinstance(ours, StructuralTag) + assert ours.model_dump() == expected.model_dump() + + +def test_glm47_parameter_level_pins_schemas_like_the_builtin( + sample_tools: list[ChatCompletionToolsParam], +): + # The operator floor overrides per-tool strictness: every tool is pinned. + ours = get_model_structural_tag( + model="glm_4_7", + tools=sample_tools, + tool_choice="required", + reasoning=False, + strict_level=ToolStrictLevel.PARAMETER, + ) + expected = get_xgrammar_builtin_structural_tag( + model="glm_4_7", + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "strict": True, + "parameters": sample_tools[0].function.parameters, + }, + } + ], + tool_choice="required", + reasoning=False, + ) + + assert isinstance(ours, StructuralTag) + assert ours.model_dump() == expected.model_dump() + + +def test_glm47_mixed_tools_dispatch_per_tool( + sample_tools_strict: list[ChatCompletionToolsParam], + sample_tools: list[ChatCompletionToolsParam], +): + mixed = [sample_tools_strict[0], sample_tools[0]] + mixed[1] = mixed[1].model_copy(deep=True) + mixed[1].function.name = "run_command" + + tag = get_model_structural_tag( + model="glm_4_7", + tools=mixed, + tool_choice="required", + reasoning=False, + ) + + assert isinstance(tag, StructuralTag) + tags = tag.model_dump()["format"]["tags"] + assert tags[0]["content"]["json_schema"] == ( + sample_tools_strict[0].function.parameters + ) + assert tags[1]["content"]["type"] != "json_schema" + + +def _glm47_parser( + tools, monkeypatch: pytest.MonkeyPatch, *, strict_calling: bool = True +) -> Glm47MoeModelToolParser: + monkeypatch.setattr(envs, "VLLM_ENFORCE_STRICT_TOOL_CALLING", strict_calling) + return Glm47MoeModelToolParser(MagicMock(), tools=tools) + + +@pytest.mark.parametrize( + ("tool_choice", "tools"), + [("auto", "plain"), ("required", "plain"), ("required", "strict")], +) +def test_glm47_get_structural_tag_disabled_when_flag_off( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], + sample_tools_strict: list[ChatCompletionToolsParam], + tool_choice: str, + tools: str, +): + parser = _glm47_parser(sample_tools, monkeypatch, strict_calling=False) + request = ChatCompletionRequest( + messages=[], + model="m", + tools=sample_tools_strict if tools == "strict" else sample_tools, + tool_choice=tool_choice, + ) + + assert parser.get_structural_tag(request) is None + + +def test_glm47_default_strict_level_constrains_auto( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], +): + # --tool-strict-level auto: the parser's FUNCTION default lifts the + # non-strict auto gate. + monkeypatch.setattr(envs, "VLLM_ENFORCE_STRICT_TOOL_CALLING", True) + + class TestParser(DelegatingParser): + tool_parser_cls = Glm47MoeModelToolParser + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=[tool.model_dump(exclude_none=True) for tool in sample_tools], + tool_choice="auto", + ) + parser = TestParser(MagicMock(), tools=sample_tools) + parser._reasoning_parser = MagicMock(adjust_request=lambda request: request) + + out = parser.adjust_request(request) + + assert out.structured_outputs is not None + assert out.structured_outputs.json is None + grammar = Grammar.from_structural_tag(out.structured_outputs.structural_tag) + call = ( + "get_weathercity" + "Paris" + ) + assert _is_grammar_accept_string(grammar, "Sure. " + call) + assert _is_grammar_accept_string(grammar, "Plain answer.") + assert not _is_grammar_accept_string( + grammar, "Sure. " + call.replace("city", "ctiy") + ) + + +def test_glm47_responses_request_attaches_shallow_tag( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(envs, "VLLM_ENFORCE_STRICT_TOOL_CALLING", True) + + class TestParser(DelegatingParser): + tool_parser_cls = Glm47MoeModelToolParser + + request = ResponsesRequest.model_validate( + { + "input": "hi", + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + } + ], + "tool_choice": "required", + } + ) + parser = TestParser(MagicMock(), tools=None) + parser._reasoning_parser = MagicMock(adjust_request=lambda request: request) + + out = parser.adjust_request(request) + + assert out.structured_outputs is not None + assert out.structured_outputs.structural_tag is not None + assert out.text is None + grammar = Grammar.from_structural_tag(out.structured_outputs.structural_tag) + assert _is_grammar_accept_string( + grammar, + "get_weathercity" + "Paris", + ) + + +def test_glm47_operator_level_overrides_parser_default( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], +): + monkeypatch.setattr(envs, "VLLM_ENFORCE_STRICT_TOOL_CALLING", True) + + class TestParser(DelegatingParser): + tool_parser_cls = Glm47MoeModelToolParser + tool_strict_level = ToolStrictLevel.PARAMETER + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=[tool.model_dump(exclude_none=True) for tool in sample_tools], + tool_choice="required", + ) + parser = TestParser(MagicMock(), tools=sample_tools) + parser._reasoning_parser = MagicMock(adjust_request=lambda request: request) + + out = parser.adjust_request(request) + + assert out.structured_outputs is not None + content = json.loads(out.structured_outputs.structural_tag)["format"]["tags"][0][ + "content" + ] + assert content["json_schema"] == sample_tools[0].function.parameters + + +def test_glm47_required_with_non_strict_tools_attaches_shallow_tag( + monkeypatch: pytest.MonkeyPatch, + sample_tools: list[ChatCompletionToolsParam], +): + monkeypatch.setattr(envs, "VLLM_ENFORCE_STRICT_TOOL_CALLING", True) + + class TestParser(DelegatingParser): + tool_parser_cls = Glm47MoeModelToolParser + + request = ChatCompletionRequest( + messages=[], + model="m", + tools=[tool.model_dump(exclude_none=True) for tool in sample_tools], + tool_choice="required", + ) + parser = TestParser(MagicMock(), tools=sample_tools) + parser._reasoning_parser = MagicMock(adjust_request=lambda request: request) + + out = parser.adjust_request(request) + + assert out.structured_outputs is not None + assert out.structured_outputs.json is None + grammar = Grammar.from_structural_tag(out.structured_outputs.structural_tag) + call = ( + "get_weathercity" + "Paris" + ) + assert _is_grammar_accept_string(grammar, call) + assert not _is_grammar_accept_string(grammar, "It is sunny in Paris.") + assert not _is_grammar_accept_string(grammar, call.replace("city", "ctiy")) + + @pytest.mark.parametrize("policy", ["auto", "required", "named"]) def test_mimo_strict_compact_xml(policy): tools = [ diff --git a/tests/tools/test_check_agent_files.py b/tests/tools/test_check_agent_files.py new file mode 100644 index 0000000000..ec9d4c8e09 --- /dev/null +++ b/tests/tools/test_check_agent_files.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the agent-file pre-commit check. + +The hook passes silently on a healthy tree, so CI would not notice a rule that +stopped firing. Each case breaks one convention in a minimal repo and asserts +the check reports it. +""" + +import os +from pathlib import Path + +import pytest + +from tools.pre_commit.check_agent_files import check + +SKILL_MD = ".agents/skills/demo/SKILL.md" +SKILL = "---\nname: demo\ndescription: Demo things. Use when testing.\n---\n" + + +def write(path: str, text: str) -> None: + Path(path).parent.mkdir(parents=True, exist_ok=True) + Path(path).write_text(text) + + +def tracked() -> dict[str, str]: + """Mimic `git ls-files -s` for every file and symlink under the cwd.""" + files = {} + for root, dirs, names in os.walk("."): + for name in dirs + names: + path = os.path.normpath(os.path.join(root, name)) + if os.path.islink(path): + files[path] = "120000" + elif os.path.isfile(path): + files[path] = "100644" + return files + + +@pytest.fixture(autouse=True) +def repo(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + write("AGENTS.md", "See the [guide](docs/guide.md#some-heading).\n") + write("docs/guide.md", "## Some Heading\n") + write(SKILL_MD, SKILL) + os.makedirs(".claude/skills") + os.symlink("../../.agents/skills/demo", ".claude/skills/demo") + + +def test_valid_repo_passes(): + assert check(tracked()) == ([], []) + + +@pytest.mark.parametrize( + ("break_repo", "expected"), + [ + pytest.param( + lambda: os.unlink(".claude/skills/demo"), + "must be a symlink", + id="skill-not-linked", + ), + pytest.param( + lambda: write(".claude/skills/other/SKILL.md", SKILL), + "only symlinks", + id="real-file-in-claude-skills", + ), + pytest.param( + lambda: os.symlink("../../.agents/skills/gone", ".claude/skills/gone"), + "only symlinks", + id="dangling-link", + ), + pytest.param( + lambda: write(SKILL_MD, SKILL.replace("demo", "other")), + "must be 'demo'", + id="name-differs-from-dir", + ), + pytest.param( + lambda: write(SKILL_MD, SKILL.replace("Demo things", "Demo: things")), + "not valid YAML", + id="unquoted-colon-in-description", + ), + pytest.param( + lambda: write(SKILL_MD, "# Demo\n"), + "missing YAML frontmatter", + id="no-frontmatter", + ), + pytest.param( + lambda: write("vllm/CLAUDE.md", "@AGENTS.md\n"), + "remove it", + id="claude-md", + ), + pytest.param( + lambda: write("AGENTS.md", "line\n" * 201), + "200-line budget", + id="root-guide-over-budget", + ), + pytest.param( + lambda: write("vllm/AGENTS.md", "line\n" * 301), + "300-line budget", + id="domain-guide-over-budget", + ), + pytest.param( + lambda: write(SKILL_MD, SKILL + "line\n" * 500), + "500-line budget", + id="skill-over-budget", + ), + pytest.param( + lambda: write("AGENTS.md", "[gone](docs/gone.md)\n"), + "does not exist", + id="dead-link", + ), + pytest.param( + lambda: write(SKILL_MD, SKILL + "[guide](../../../docs/guide.md#gone)\n"), + "no heading", + id="dead-anchor", + ), + ], +) +def test_broken_convention_is_reported(break_repo, expected): + break_repo() + errors, _ = check(tracked()) + assert len(errors) == 1 + assert expected in errors[0] + + +def test_description_without_trigger_is_advisory(): + write(SKILL_MD, SKILL.replace(" Use when testing.", "")) + errors, notes = check(tracked()) + assert not errors + assert "when to use" in notes[0] diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 5bb6eb4c32..8acfc0a334 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -23,6 +23,8 @@ from vllm.v1.attention.backends.mla.compressor_utils import ( from vllm.v1.attention.backends.mla.indexer import ( BuildPrefillChunkMetadataKernel, DeepseekV4IndexerBackend, + DeepSeekV32IndexerDecodeMetadata, + DeepseekV32IndexerMetadata, DeepseekV32IndexerMetadataBuilder, DeepseekV41IndexerBackend, get_max_prefill_buffer_size, @@ -99,6 +101,72 @@ def test_fused_indexer_decode_metadata(query_lens, padding): torch.testing.assert_close(per_req[:reqs], lengths) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_indexer_draft_decode_metadata_update_is_capture_safe(): + if not indexer_module.has_deep_gemm(): + pytest.skip("requires DeepGEMM") + device = torch.device("cuda") + num_reqs = 3 + seq_lens = torch.tensor([17, 33, 65], dtype=torch.int32, device=device) + decode_seq_lens = torch.zeros((num_reqs, 1), dtype=torch.int32, device=device) + block_table = torch.arange(num_reqs * 4, dtype=torch.int32, device=device).view( + num_reqs, 4 + ) + num_sms = torch.cuda.get_device_properties(device).multi_processor_count + schedule_metadata = indexer_module.get_paged_mqa_logits_metadata( + torch.div(seq_lens, 16, rounding_mode="floor").unsqueeze(1), + 64, + num_sms, + ).clone() + + builder = object.__new__(DeepseekV32IndexerMetadataBuilder) + builder.dcp_world_size = 1 + builder.compress_ratio = 16 + builder.num_sms = num_sms + builder.arange_buffer = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) + builder.kv_cache_spec = SimpleNamespace(num_states=64) + metadata = DeepseekV32IndexerMetadata( + seq_lens=seq_lens, + max_seq_len=128, + slot_mapping=torch.zeros(num_reqs, dtype=torch.int64, device=device), + num_decodes=num_reqs, + num_decode_tokens=num_reqs, + num_prefills=0, + num_prefill_tokens=0, + decode=DeepSeekV32IndexerDecodeMetadata( + block_table=block_table, + seq_lens=decode_seq_lens, + decode_lens=torch.zeros(num_reqs, dtype=torch.int32, device=device), + requires_padding=False, + schedule_metadata=schedule_metadata, + ), + ) + + builder.update_draft_decode_metadata(metadata) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + builder.update_draft_decode_metadata(metadata) + + seq_lens.copy_(torch.tensor([32, 48, 80], dtype=torch.int32, device=device)) + graph.replay() + torch.accelerator.synchronize() + + expected_decode_seq_lens = torch.div(seq_lens, 16, rounding_mode="floor") + torch.testing.assert_close(decode_seq_lens.flatten(), expected_decode_seq_lens) + torch.testing.assert_close( + schedule_metadata, + indexer_module.get_paged_mqa_logits_metadata( + expected_decode_seq_lens.unsqueeze(1), 64, num_sms + ), + ) + assert metadata.decode is not None + assert torch.all(metadata.decode.decode_lens == 1) + expected_slot_mapping = block_table[:, 0].to(torch.int64) * 64 + torch.tensor( + [1, 2, 4], dtype=torch.int64, device=device + ) + torch.testing.assert_close(metadata.slot_mapping, expected_slot_mapping) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize( "query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0], [0, 0, 0]] diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 157366e49c..da25eaefb2 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -300,6 +300,16 @@ def test_builder_build_falls_back_without_positions(): assert meta.slot_mapping is cam.slot_mapping +def test_builder_updates_draft_mapping(): + builder = make_tail_builder() + cam = make_common_metadata([[15], [16]], [5, 9]) + meta = KpoolTailMetadataBuilder.build(builder, 0, cam) + assert cam.positions is not None + cam.positions.add_(1) + builder.update_draft_decode_metadata(meta) + assert meta.slot_mapping[:2].tolist() == [5 * KPOOL, 9 * KPOOL + 1] + + def test_builder_reuses_slot_mapping_storage(): builder = make_tail_builder() first = make_common_metadata([list(range(10))], [5]) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index b41db5a0f1..468e36a787 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -65,6 +65,9 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import ( ) from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import current_stream +from vllm.v1.attention.backends.mla import ( + flashattn_mla_sparse as flashattn_sparse_module, +) from vllm.v1.attention.backends.mla import index_group as index_group_module from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( FlashAttnMLASparseImpl, @@ -3405,6 +3408,129 @@ def test_flashinfer_hisparse_decode_runs_batched_attention(): assert lse is None +@pytest.mark.skipif( + not current_platform.is_device_capability_family(90), + reason="FlashAttention MLA requires Hopper", +) +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.parametrize("block_stride_rows", [64, 128]) +@pytest.mark.parametrize( + "query_len,num_decode_tokens,max_decode_rows,num_layers", + [ + pytest.param(1, 2, 2, 4, id="decode"), + pytest.param(4, 8, 8, 4, id="mtp4"), + pytest.param(6, 12, 12, 4, id="mtp6"), + pytest.param(4, 8, 2, 4, id="scratch_overflow"), + pytest.param(4, 4, 8, 4, id="mixed"), + pytest.param(4, 0, 8, 4, id="prefill"), + pytest.param(4, 8, 8, 1, id="single_layer"), + ], +) +def test_flashattn_shared_indices_match_independent_conversion( + monkeypatch, + query_len, + num_decode_tokens, + max_decode_rows, + num_layers, + block_stride_rows, + use_cuda_graph, +): + """Reuse must preserve attention outputs as requests and pages change.""" + device = torch.device(DEVICE_TYPE) + num_tokens, block_size, topk, seq_len = 2 * query_len, 64, 2048, 8192 + blocks_per_req = seq_len // block_size + logical = torch.arange( + 0, seq_len, seq_len // topk, dtype=torch.int32, device=device + ).repeat(num_tokens, 1) + logical[:, ::7] = -1 + builder = SparseMLAIndexGroupBuilder(logical, max_decode_rows=max_decode_rows) + metadata = SimpleNamespace( + block_table=torch.randperm( + 2 * blocks_per_req, dtype=torch.int32, device=device + ).view(2, blocks_per_req), + block_size=block_size, + req_id_per_token=torch.arange( + 2, dtype=torch.int32, device=device + ).repeat_interleave(query_len), + num_decode_tokens=num_decode_tokens, + ) + q_nope = torch.randn(num_tokens, 8, 512, dtype=torch.bfloat16, device=device) + q_rope = torch.randn(num_tokens, 8, 64, dtype=torch.bfloat16, device=device) + caches = torch.randn( + num_layers, + 2 * blocks_per_req, + block_stride_rows, + 576, + dtype=torch.bfloat16, + device=device, + )[:, :, :block_size] + impls = [] + for i in range(num_layers): + impl = object.__new__(FlashAttnMLASparseImpl) + impl.topk_indices_buffer = logical + impl.index_group, impl.index_group_index = builder.register_layer(i == 0) + impl.dcp_world_size = 1 + impl.kv_lora_rank, impl.qk_rope_head_dim = 512, 64 + impl.scale = 256**-0.5 + impl.cu_seqlens_q_buffer = torch.arange( + num_tokens + 1, dtype=torch.int32, device=device + ) + impls.append(impl) + + convert = triton_convert_req_index_to_global_index + convert_spy = MagicMock(wraps=convert) + for module in (index_group_module, flashattn_sparse_module): + monkeypatch.setattr( + module, "triton_convert_req_index_to_global_index", convert_spy + ) + + def run(): + impls[0].record_logical_topk_ready() + return [ + impl.forward_mqa((q_nope, q_rope), cache, metadata, None)[0] + for impl, cache in zip(impls, caches) + ] + + stream = torch.cuda.Stream(device=device) + stream.wait_stream(current_stream()) + with torch.cuda.stream(stream): + run() # Warm up the conversion and FA3 kernels before capture. + stream.synchronize() + convert_spy.reset_mock() + if use_cuda_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = run() + else: + outputs = run() + reuse = num_decode_tokens == num_tokens and num_tokens <= max_decode_rows + 1 + assert convert_spy.call_count == (1 if reuse else num_layers) + + for _ in range(2): + metadata.req_id_per_token.copy_(metadata.req_id_per_token.flip(0)) + metadata.block_table.copy_(metadata.block_table.roll(1, dims=1)) + logical.copy_(torch.where(logical >= 0, (logical + 1) % seq_len, logical)) + if use_cuda_graph: + graph.replay() + else: + outputs = run() + physical, counts = convert( + metadata.req_id_per_token, + metadata.block_table, + logical, + BLOCK_SIZE=block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk, + return_valid_counts=True, + ) + for impl, cache, actual in zip(impls, caches, outputs): + expected = impl._run_mqa_kernel( + q_nope, q_rope, cache, physical, counts, block_size + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + stream.synchronize() + + def test_flashattn_hisparse_decode_uses_index_group(): num_tokens = 4 q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE) diff --git a/tests/v1/core/test_prefix_replay.py b/tests/v1/core/test_prefix_replay.py index b2e78e1f92..da57185808 100644 --- a/tests/v1/core/test_prefix_replay.py +++ b/tests/v1/core/test_prefix_replay.py @@ -4,6 +4,8 @@ the hit to rebuild the non-cacheable sliding-window group, keeps the hit's blocks, and hands the worker the replay start.""" +import dataclasses + import pytest import torch @@ -277,3 +279,53 @@ def test_hit_no_longer_than_window_is_ignored(via_connector): assert new_req.num_computed_tokens == 0 assert new_req.replay_start == 0 assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS + + +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize("loads_window", [True, False]) +def test_kv_load_zeroing_respects_loaded_groups(is_async, loads_window): + """Only groups restored by an async load skip zeroing. P/D loads include + the non-cacheable window; stores leave it to be initialized locally.""" + local = 5 * BLOCK_SIZE + scheduler = _replay_scheduler(use_kv_connector=MockKVConfig()) + if loads_window: + scheduler.connector.get_loaded_kv_cache_group_ids = lambda request: (FULL, SWA) + assert scheduler.needs_kv_cache_zeroing + first = create_requests( + num_requests=1, num_tokens=local, same_prompt=True, block_size=BLOCK_SIZE + )[0] + _prefill(scheduler, first) + scheduler.connector.config = dataclasses.replace( + scheduler.connector.config, + matched_tokens=NUM_PROMPT_TOKENS - 1 - local, + is_async=is_async, + ) + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + same_prompt=True, + block_size=BLOCK_SIZE, + req_ids=["load"], + )[0] + scheduler.add_request(request) + out = scheduler.schedule() + + manager = scheduler.kv_cache_manager + hit_ids = { + b.block_id for g in manager.get_blocks(first.request_id).blocks for b in g + } + swa_blocks = manager.get_blocks(request.request_id).blocks[SWA] + assert any( + not b.is_null and b.block_id not in hit_ids + for b in swa_blocks[: local // BLOCK_SIZE] + ) + zeroed = set(out.new_block_ids_to_zero or ()) + for group_id, blocks in enumerate(manager.get_blocks(request.request_id).blocks): + new_ids = { + b.block_id for b in blocks if not b.is_null and b.block_id not in hit_ids + } + assert new_ids + if is_async and (group_id == FULL or loads_window): + assert not new_ids & zeroed + else: + assert new_ids <= zeroed diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 2705b87f80..da7ab742b0 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -16,6 +16,7 @@ from vllm.v1.core.single_type_kv_cache_manager import ( ChunkedLocalAttentionManager, CircularBufferManager, FullAttentionManager, + KpoolTailManager, MambaManager, RSWAManager, SlidingWindowManager, @@ -24,6 +25,7 @@ from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, CircularBufferSpec, FullAttentionSpec, + KpoolTailSpec, MambaSpec, RSWASpec, SlidingWindowSpec, @@ -332,6 +334,37 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime(): assert manager.req_to_blocks[request_id] == blocks +@pytest.mark.parametrize("record_for_zeroing", [True, False]) +def test_external_kpool_tail_zeroing(record_for_zeroing): + """A transferred ring at index zero skips zeroing even with a local hit.""" + spec = KpoolTailSpec( + block_size=4, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + sliding_window=4, + ) + manager = KpoolTailManager( + spec, + block_pool=BlockPool(10, enable_caching=True, hash_block_size=4), + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=4, + needs_kv_cache_zeroing=True, + max_admission_blocks_per_request=1, + ) + manager.allocate_external_computed_blocks( + "request", 16, 4, record_for_zeroing=record_for_zeroing + ) + blocks = manager.req_to_blocks["request"] + assert len(blocks) == 1 + assert manager.take_new_block_ids() == ( + [blocks[0].block_id] if record_for_zeroing else [] + ) + assert manager.allocate_new_blocks("request", 20, 20) == [] + assert manager.take_new_block_ids() == [] + + def test_sliding_window_records_new_blocks_for_zeroing(): block_size = 2 spec = SlidingWindowSpec( diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index cae32e44dc..2d7b2f5cb8 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -192,6 +192,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable): full_sequence_must_fit: bool = False, reserved_blocks: int = 0, has_scheduled_reqs: bool = True, + skip_zeroing_group_ids: tuple[int, ...] = (), ): ret = original_allocate_slots_fn( self, @@ -206,6 +207,7 @@ def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable): full_sequence_must_fit, reserved_blocks, has_scheduled_reqs, + skip_zeroing_group_ids, ) if cur_step_action is not None: cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[ diff --git a/tests/v1/e2e/general/test_streaming_input.py b/tests/v1/e2e/general/test_streaming_input.py index c71e7ad0a6..5a226e75b1 100644 --- a/tests/v1/e2e/general/test_streaming_input.py +++ b/tests/v1/e2e/general/test_streaming_input.py @@ -491,7 +491,7 @@ async def test_streaming_input_per_chunk_sampling_params(engine: AsyncLLM): async def test_streaming_input_empty_generator(engine: AsyncLLM): """Test behavior when the input generator yields nothing. - An empty generator should still produce a finished output. + An empty generator should finish without producing output. """ request_id = "test_empty_generator" sampling_params = get_sampling_params(max_tokens=10) @@ -505,9 +505,7 @@ async def test_streaming_input_empty_generator(engine: AsyncLLM): async for output in engine.generate(empty_generator(), sampling_params, request_id): outputs.append(output) - # Should still get a finished marker - assert len(outputs) >= 1, "Should receive at least one output" - assert outputs[-1].finished, "Should have a finished output" + assert outputs == [] @pytest.mark.asyncio(loop_scope="module") diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py index d2c2ac87b7..10d449c86f 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -150,6 +150,21 @@ def test_draft_model_quantization( assert_draft_model_correctness(sd_case, vllm_runner) +@single_gpu_only +def test_draft_model_transformers_backend(vllm_runner): + """Ensure that Attention prefixes don't collide between the target and + draft models when both models use the Transformers modelling backend.""" + sd_case = ArgsTest( + target_model="allenai/OLMo-2-0425-1B-Instruct", + draft_model="allenai/OLMo-2-0425-1B-Instruct", + sampling_config=greedy_sampling(), + num_speculative_tokens=3, + expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1 + expected_acceptance_rate=0.98, # slight epsilon + ) + assert_draft_model_correctness(sd_case, vllm_runner) + + @multi_gpu_only(num_gpus=2) def test_draft_model_tensor_parallelism(vllm_runner): """Ensure spec decode works when running with TP > 1.""" diff --git a/tests/v1/e2e/spec_decode/mtp/_correctness.py b/tests/v1/e2e/spec_decode/mtp/_correctness.py index b88c1cce7a..352aa556e0 100644 --- a/tests/v1/e2e/spec_decode/mtp/_correctness.py +++ b/tests/v1/e2e/spec_decode/mtp/_correctness.py @@ -48,16 +48,6 @@ def check_mtp_correctness( elif "gemma-4" in model_name: extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0} - if draft_model is not None and "gemma-4" in draft_model: - import transformers - from packaging.version import Version - - if Version(transformers.__version__) < Version("5.8.0"): - pytest.skip( - "Gemma4 MTP assistant requires transformers>=5.8.0, " - f"got {transformers.__version__}" - ) - with vllm_runner( model_name, block_size=None, diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py index c9832f7d8e..b09f9405d7 100644 --- a/tests/v1/engine/test_llm_engine.py +++ b/tests/v1/engine/test_llm_engine.py @@ -2,11 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random from typing import TYPE_CHECKING +from unittest.mock import MagicMock, call import pytest from vllm import LLM from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.engine.llm_engine import LLMEngine from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector if TYPE_CHECKING: @@ -18,6 +20,38 @@ MODEL = "facebook/opt-125m" DTYPE = "half" +@pytest.mark.skip_global_cleanup +def test_step_invalidates_mm_processor_cache_misses(): + engine = object.__new__(LLMEngine) + engine.should_execute_dummy_batch = False + engine.log_stats = False + engine.logger_manager = None + + mm_processor_cache = MagicMock() + engine.renderer = MagicMock(mm_processor_cache=mm_processor_cache) + + engine_core_output = MagicMock( + mm_cache_miss_hashes=["missing-hash-1", "missing-hash-2"] + ) + outputs = MagicMock( + outputs=[engine_core_output], + timestamp=0.0, + scheduler_stats=None, + ) + engine.engine_core = MagicMock() + engine.engine_core.get_output.return_value = outputs + + processed_outputs = MagicMock(reqs_to_abort=[], request_outputs=[]) + engine.output_processor = MagicMock() + engine.output_processor.process_outputs.return_value = processed_outputs + + assert engine.step() == [] + assert mm_processor_cache.invalidate.call_args_list == [ + call("missing-hash-1"), + call("missing-hash-2"), + ] + + def _vllm_model( apc: bool, vllm_runner: type[VllmRunner], diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 40a6883aa9..44c3ee20b5 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, Mock import numpy as np import pytest +import torch from tests.v1.engine.utils import ( NUM_PROMPT_LOGPROBS_UNDER_TEST, @@ -18,7 +19,12 @@ from tests.v1.engine.utils import ( from vllm import PoolingParams from vllm.logprobs import FlatLogprobs, Logprob, PromptLogprobs, SampleLogprobs from vllm.lora.request import LoRARequest -from vllm.outputs import CompletionOutput, RequestOutput, SamplingMask +from vllm.outputs import ( + CompletionOutput, + PoolingRequestOutput, + RequestOutput, + SamplingMask, +) from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.v1.engine import ( @@ -1630,3 +1636,92 @@ def test_request_output_add_merges_delta_sampling_masks(): assert merged.outputs[0].token_ids == [10, 20] assert merged.outputs[0].sampling_mask.token_ids == [[10, 11], [20]] + + +@pytest.mark.skip_global_cleanup +def test_pooling_cache_miss_is_a_request_error_and_batch_continues(): + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + requests = [ + EngineCoreRequest( + request_id=f"request-{idx}-int", + external_req_id=f"request-{idx}", + prompt_token_ids=[1, 2], + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=None, + pooling_params=PoolingParams(task="embed"), + ) + for idx in range(2) + ] + for request in requests: + output_processor.add_request(request, prompt=None) + + processed = output_processor.process_outputs( + [ + EngineCoreOutput( + request_id=requests[0].request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + mm_cache_miss_hashes=["missing-hash"], + ), + EngineCoreOutput( + request_id=requests[1].request_id, + new_token_ids=[], + pooling_output=torch.tensor([1.0, 2.0]), + finish_reason=FinishReason.STOP, + ), + ] + ) + + assert len(processed.request_outputs) == 2 + failed, succeeded = processed.request_outputs + assert isinstance(failed, PoolingRequestOutput) + assert failed.error is not None + assert failed.error.code == "multimodal_cache_miss" + assert failed.error.retryable + assert isinstance(succeeded, PoolingRequestOutput) + assert succeeded.error is None + assert torch.equal(succeeded.outputs.data, torch.tensor([1.0, 2.0])) + + +@pytest.mark.skip_global_cleanup +def test_async_pooling_cache_miss_does_not_fail_output_processor(): + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + request = EngineCoreRequest( + request_id="request-int", + external_req_id="request", + prompt_token_ids=[1, 2], + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=None, + pooling_params=PoolingParams(task="embed"), + ) + queue = RequestOutputCollector( + output_kind=RequestOutputKind.FINAL_ONLY, + request_id=request.request_id, + ) + output_processor.add_request(request, prompt=None, queue=queue) + + processed = output_processor.process_outputs( + [ + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + mm_cache_miss_hashes=["missing-hash"], + ) + ] + ) + + assert processed.request_outputs == [] + output = queue.get_nowait() + assert isinstance(output, PoolingRequestOutput) + assert output.error is not None + assert output.error.retryable + assert not output_processor.has_unfinished_requests() diff --git a/tests/v1/kv_connector/unit/test_error_propagation.py b/tests/v1/kv_connector/unit/test_error_propagation.py index 91ec525f5c..35a3749e08 100644 --- a/tests/v1/kv_connector/unit/test_error_propagation.py +++ b/tests/v1/kv_connector/unit/test_error_propagation.py @@ -111,6 +111,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler): } fail_scheduler.connector = Mock() + fail_scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) fail_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) ) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 777bada58b..7ea2170584 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -327,6 +327,7 @@ def test_async_recompute_blocks_not_cached_when_invalid( recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True) ) + recompute_scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) recompute_scheduler.connector.request_finished.return_value = (False, None) recompute_scheduler.connector.take_events.return_value = () diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 108eb049d9..3dbbfdc6b2 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -42,6 +42,7 @@ def test_failed_receive_completion_honors_failure_policy(policy): request = create_request(num_tokens=3 * scheduler.block_size) scheduler.add_request(request) scheduler.connector = Mock() + scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = [ (2 * scheduler.block_size, True), (0, False), @@ -103,6 +104,7 @@ def test_async_load_failure( } scheduler.connector = Mock() + scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True) ) @@ -332,6 +334,7 @@ def test_async_progressive_load_failure( } scheduler.connector = Mock() + scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,) scheduler.connector.get_num_new_matched_tokens.side_effect = ( _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True) ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index c306103288..d326615aad 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -1549,6 +1549,40 @@ async def test_kv_producer(monkeypatch): prefill_worker.shutdown() +@pytest.mark.parametrize( + ("local_block_ids", "expected_finished"), + [ + ([], set()), + ([[]], set()), + ([[100]], {"d-req-1"}), + ], +) +def test_pull_completion_requires_local_blocks( + local_block_ids: list[list[int]], + expected_finished: set[str], +): + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.shutdown = MagicMock() + worker.finished_recving_reqs = set() + pull_meta = PullReqMeta( + d_req_id="d-req-1", + transfer_id="xfer-req-1", + local_block_ids=local_block_ids, + remote_engine_id="p-engine", + remote_bootstrap_addr="http://bootstrap:33333", + pull_tasks_count=1, + ) + response = MooncakeXferResponse( + status=MooncakeXferResponseStatus.FINISH, + ok_reqs=["d-req-1"], + ) + + worker.process_pulling_result(response, {"d-req-1": pull_meta}) + + assert pull_meta.pull_tasks_count == 0 + assert worker.finished_recving_reqs == expected_finished + + @pytest.mark.asyncio async def test_kv_consumuer(monkeypatch): """Simulates a Consumer Worker (Decoder) initiating a pull from a Producer. diff --git a/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py b/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py index 71ee0ef5d7..e3e04aedb8 100644 --- a/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py @@ -22,14 +22,20 @@ Role shapes below are exactly those in the RFC (#46107) deployments: """ import contextlib +import ctypes import importlib.util +import os +import subprocess import sys +import time import types from collections import Counter from pathlib import Path from typing import cast +import msgpack import pytest +import zmq _MISSING = object() @@ -203,3 +209,127 @@ def test_consecutive_requests_alternate_instances(route): # than filling one instance's ranks before moving on. insts = [inst for inst, _ in _route_n(route, _instances(3, 8), 3)] assert insts == [0, 1, 2] + + +@pytest.mark.skipif(sys.platform != "linux", reason="GIL probe uses libc.usleep") +def test_discovery_heartbeat_progresses_while_worker_holds_gil(): + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import ( + MoRIIOHeartbeat, + ) + + payload = { + "type": "P", + "http_address": "127.0.0.1:8000", + "zmq_address": "host:127.0.0.1,handshake:6301,notify:61005", + "dp_size": 1, + "tp_size": 8, + "transfer_mode": "READ", + } + with zmq.Context() as context, context.socket(zmq.ROUTER) as socket: + port = socket.bind_to_random_port("tcp://127.0.0.1") + heartbeat = MoRIIOHeartbeat(f"tcp://127.0.0.1:{port}", payload, 0.02) + try: + assert socket.poll(5000), "Discovery did not register" + assert msgpack.unpackb(socket.recv_multipart()[1]) == payload + # PyDLL retains the GIL throughout the native call. + ctypes.PyDLL(None).usleep(300_000) + received = 0 + while socket.poll(0): + assert msgpack.unpackb(socket.recv_multipart()[1]) == payload + received += 1 + assert received >= 3, "Discovery stalled while the worker held the GIL" + finally: + heartbeat.shutdown() + heartbeat.shutdown() + while socket.poll(0): + socket.recv_multipart() + assert not socket.poll(100), "Heartbeats continued after shutdown" + + +@pytest.mark.skipif(sys.platform != "linux", reason="Checks child exit via procfs") +def test_discovery_heartbeat_exits_when_worker_dies(): + from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_heartbeat + + owner_code = """ +import runpy +import sys +Heartbeat = runpy.run_path(sys.argv[1])["MoRIIOHeartbeat"] +heartbeat = Heartbeat(sys.argv[2], {"type": "D"}, 0.02) +print(heartbeat._process.pid, flush=True) +sys.stdin.read() +""" + with zmq.Context() as context, context.socket(zmq.ROUTER) as socket: + port = socket.bind_to_random_port("tcp://127.0.0.1") + owner = subprocess.Popen( + [ + sys.executable, + "-c", + owner_code, + moriio_heartbeat.__file__, + f"tcp://127.0.0.1:{port}", + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + ) + child_pid = None + child_exited = False + try: + assert socket.poll(5000), "Discovery did not register" + assert owner.stdout is not None + child_pid = int(owner.stdout.readline()) + owner.kill() + owner.wait(timeout=5) + stat = Path(f"/proc/{child_pid}/stat") + deadline = time.monotonic() + 5 + while stat.exists(): + # A terminated child may await reaping by the container's init. + try: + if stat.read_text().split(") ")[1].startswith("Z "): + break + except FileNotFoundError: + break + assert time.monotonic() < deadline, "Orphan heartbeat survived worker" + time.sleep(0.01) + child_exited = True + finally: + if owner.poll() is None: + owner.kill() + owner.communicate(timeout=5) + if child_pid is not None and not child_exited: + with contextlib.suppress(ProcessLookupError): + os.kill(child_pid, 9) + + +def test_discovery_heartbeat_recovers_after_router_outage(capfd): + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import ( + MoRIIOHeartbeat, + ) + + payload = {"type": "P", "http_address": "127.0.0.1:8000"} + with zmq.Context() as context: + with context.socket(zmq.ROUTER) as socket: + port = socket.bind_to_random_port("tcp://127.0.0.1") + endpoint = f"tcp://127.0.0.1:{port}" + heartbeat = MoRIIOHeartbeat(endpoint, payload, 0.001, max_retries=1) + try: + assert socket.poll(5000), "Discovery did not register" + assert msgpack.unpackb(socket.recv_multipart()[1]) == payload + except BaseException: + heartbeat.shutdown() + raise + try: + # Fill the real DEALER queue while discovery is unavailable. + deadline = time.monotonic() + 10 + errors = "" + while "heartbeat send failed" not in errors: + errors += capfd.readouterr().err + assert time.monotonic() < deadline, "No send timeout during outage" + time.sleep(0.01) + with context.socket(zmq.ROUTER) as socket: + socket.bind(endpoint) + assert socket.poll(5000), "Discovery did not recover after outage" + assert msgpack.unpackb(socket.recv_multipart()[1]) == payload + assert heartbeat._process.poll() is None + finally: + heartbeat.shutdown() diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py index a35f0a1169..34d193e272 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py @@ -269,6 +269,34 @@ def test_worker_get_finished_counts_structured_release_fan_in(): assert worker._completed_consumer_notifications == {"tx-fanin"} +@pytest.mark.parametrize("exit_kind", ["killed", "eof"]) +def test_worker_reports_unexpected_heartbeat_exit(exit_kind): + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import ( + MoRIIOHeartbeat, + ) + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.is_producer = True + worker.mode = MoRIIOMode.READ + worker.moriio_wrapper = SimpleNamespace( + pop_finished_req_ids=lambda: [], shutdown=lambda: None + ) + worker._pending_unmapped_acks = [] + worker._heartbeat = MoRIIOHeartbeat("tcp://127.0.0.1:1", {"type": "P"}, 0.02) + process = worker._heartbeat._process + try: + assert worker.get_finished() == (set(), set()) + if exit_kind == "killed": + process.kill() + else: + process.stdin.close() + returncode = process.wait(timeout=5) + with pytest.raises(RuntimeError, match=f"heartbeat.*{returncode}"): + worker.get_finished() + finally: + worker.shutdown() + + def test_read_completion_sends_structured_release_with_consumer_tp_size(): class DoneStatus: def Succeeded(self): diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index e636aabfd9..67cfdc1ab0 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -202,6 +202,20 @@ def mc() -> MultiConnector: return mc +@pytest.mark.parametrize("chosen", [0, 1]) +def test_loaded_groups_follow_the_connector_serving_the_request(mc, chosen): + request = MagicMock(request_id="request") + for i, connector in enumerate(mc._connectors): + connector.get_num_new_matched_tokens.return_value = ( + (16, True) if i == chosen else (0, False) + ) + connector.get_loaded_kv_cache_group_ids.return_value = (i,) + + assert mc.get_num_new_matched_tokens(request, 0) == (16, True) + assert mc.get_loaded_kv_cache_group_ids(request) == (chosen,) + mc._connectors[1 - chosen].get_loaded_kv_cache_group_ids.assert_not_called() + + def test_multi_connector_mem_pool_context_none(mc: MultiConnector): assert mc.get_mem_pool_context() is None diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 5a9a10a57d..2a1ca5a877 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -1522,34 +1522,6 @@ def _make_fake_kv_cache_manager(): return manager -@pytest.mark.cpu_test -def test_zeroing_block_ids_cover_only_loaded_attention_blocks(): - """Only zero-recorded (attention) groups contribute, sliced to the - externally-loaded token range; Mamba state blocks are never zeroed.""" - manager = _make_fake_kv_cache_manager() - - # Tokens [0, 16) are locally cached; the load covers tokens [16, 56). - assert manager.get_zeroing_block_ids_in_range("req-1", 16, 56) == [11, 12, 13] - - -@pytest.mark.cpu_test -def test_scheduler_filters_connector_loaded_blocks_from_zeroing(): - """Blocks that will be loaded by the connector must not be zeroed.""" - from vllm.v1.core.sched.scheduler import Scheduler - - class FakeKVCacheManager: - def take_new_block_ids(self): - return [9, 10, 11, 12] - - scheduler = object.__new__(Scheduler) - scheduler.needs_kv_cache_zeroing = True - scheduler.kv_cache_manager = FakeKVCacheManager() - scheduler._skip_zero_block_ids = {10, 12} - - assert scheduler._get_new_block_ids_to_zero() == [9, 11] - assert not scheduler._skip_zero_block_ids - - @pytest.mark.cpu_test def test_failed_load_rezeroes_unwritten_skipped_blocks(): """A failed async load leaves zeroing-skipped blocks unwritten beyond @@ -1575,7 +1547,6 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks(): # Attention blocks covering tokens >= 48 are re-recorded for zeroing # and flow into the next step's zero list; Mamba blocks are not. - scheduler._skip_zero_block_ids = set() assert scheduler._get_new_block_ids_to_zero() == [13, 14, 15] diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 2413c824fc..317ac10d41 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -11,6 +11,7 @@ from unittest.mock import patch import numpy as np import pytest import torch +from pydantic import TypeAdapter, ValidationError from tests.utils import large_gpu_mark from tests.v1.sample.utils import ( @@ -20,7 +21,7 @@ from tests.v1.sample.utils import ( compute_correct_cumulative_logprob, get_test_batch, ) -from vllm import SamplingParams +from vllm import SamplingParams, TokensPrompt from vllm.config.model import LogprobsMode from vllm.distributed import cleanup_dist_env_and_memory from vllm.exceptions import VLLMValidationError @@ -88,6 +89,7 @@ def hf_model(hf_runner) -> Generator[HfRunner, None, None]: def _model_config(vocab_size: int = 10, max_logprobs: int = 20): return SimpleNamespace( max_logprobs=max_logprobs, + max_model_len=16, logits_processors=None, is_diffusion=False, get_vocab_size=lambda: vocab_size, @@ -443,24 +445,58 @@ def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]): ) -def test_prompt_logprob_token_ids_bounded_by_max_logprobs(): - """The candidate count is bounded by max_logprobs.""" +def test_prompt_logprob_token_ids_validation(): + """Rows may be ragged (-1 padded); the width is bounded by max_logprobs.""" model_config = _model_config(vocab_size=100, max_logprobs=4) def verify(**kwargs): - SamplingParams(**kwargs).verify( + params = SamplingParams(**kwargs) + params.verify( model_config, speculative_config=None, structured_outputs_config=None, tokenizer=None, ) + return params.prompt_logprob_token_ids - verify(prompt_logprob_token_ids=[1, 2, 3, 4]) + assert verify(prompt_logprob_token_ids=[[], [1], [4, 3, 2, 1]]).tolist() == [ + [-1] * 4, + [1, -1, -1, -1], + [4, 3, 2, 1], + ] + # Any integer layout is normalized to what the worker uploads: + # C-contiguous native int32. + odd = np.asfortranarray(np.array([[1, 2], [3, -1]], dtype=">i8")) + ids = verify(prompt_logprob_token_ids=odd) + assert ids.dtype == np.int32 and ids.flags.c_contiguous + assert ids.tolist() == [[1, 2], [3, -1]] + # The pydantic path (token-in/token-out API, /docs) takes strict integer lists. + adapter = TypeAdapter(SamplingParams) + for bad in ("[[1.5]]", '[["3"]]', "[[true]]", "[1, 2]"): + with pytest.raises(ValidationError): + adapter.validate_json(f'{{"prompt_logprob_token_ids": {bad}}}') + with pytest.raises(VLLMValidationError, match=r"shape \[num_rows, num_ids\]"): + verify(prompt_logprob_token_ids=np.array([1, 2])) + with pytest.raises(VLLMValidationError, match=r"shape \[num_rows, num_ids\]"): + verify(prompt_logprob_token_ids=[1, 2]) + for bad_ids in ([[None]], [[2**63]], [[[1]]], [[[1], 2]]): + with pytest.raises(VLLMValidationError, match="integer token ids"): + verify(prompt_logprob_token_ids=bad_ids) + with pytest.raises(VLLMValidationError, match="out-of-vocab"): + verify(prompt_logprob_token_ids=np.array([[2**64 - 1]], dtype=np.uint64)) # The error names the value to raise max_logprobs to. with pytest.raises(VLLMValidationError, match=r"max_logprobs.*at least 5"): - verify(prompt_logprob_token_ids=[1, 2, 3, 4, 5]) - + verify(prompt_logprob_token_ids=[[1], [1, 2, 3, 4, 5]]) + with pytest.raises(VLLMValidationError, match="out-of-vocab"): + verify(prompt_logprob_token_ids=[[1], [100]]) + with pytest.raises(VLLMValidationError, match="out-of-vocab"): + verify(prompt_logprob_token_ids=[[1], [-2]]) + with pytest.raises(VLLMValidationError, match="non-empty"): + verify(prompt_logprob_token_ids=[[], []]) + # Rows are bounded by max_model_len before the table is padded. + with pytest.raises(VLLMValidationError, match="at most 15 scored rows"): + verify(prompt_logprob_token_ids=[[]] * 15 + [[1]]) # prompt_logprob_start alone is a caller mistake, not a silent no-op. with pytest.raises(VLLMValidationError, match="requires prompt_logprob_token_ids"): verify(prompt_logprob_start=3) @@ -483,7 +519,7 @@ def test_prompt_logprob_token_ids_require_v2_model_runner(): tokenizer=None, validate_logits_processors_params=lambda params: None, ) - params = SamplingParams(prompt_logprob_token_ids=[1, 2]) + params = SamplingParams(prompt_logprob_token_ids=[[1, 2]]) with patch.object(SamplingParams, "verify"): with pytest.raises(VLLMValidationError, match="V2 model runner"): InputProcessor._validate_params(processor, params, ("generate",)) @@ -1369,7 +1405,7 @@ def test_prompt_logprobs_with_chunking_and_preemption(): def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): - """Fixed-ID scores stay row-aligned across chunked prefill and preemption.""" + """Per-row scores stay row-aligned across chunked prefill and preemption.""" monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") prompts = [ @@ -1380,14 +1416,9 @@ def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): ] + [f"Tell me about the number {i}: " for i in range(32)] start = 2 - candidate_ids = [10, 100, 1000, 10000] - sampling_params = SamplingParams( - temperature=0.0, - max_tokens=40, - min_tokens=20, - prompt_logprob_token_ids=candidate_ids, - prompt_logprob_start=start, - ) + + def candidate_ids(j): + return [10 + j, 100 + j, 1000 + j, 10000 + j] with VllmRunner( "Qwen/Qwen3-0.6B", @@ -1398,15 +1429,38 @@ def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): disable_log_stats=False, gpu_memory_utilization=0.25, ) as vllm_model: + tokenizer = vllm_model.llm.get_tokenizer() + token_prompts = [ + TokensPrompt(prompt_token_ids=tokenizer(p).input_ids) for p in prompts + ] + + def make_params(row_ids): + return [ + SamplingParams( + temperature=0.0, + max_tokens=40, + min_tokens=20, + prompt_logprob_token_ids=[ + row_ids(j) + for j in range(len(p["prompt_token_ids"]) - 1 - start) + ], + prompt_logprob_start=start, + ) + for p in token_prompts + ] + metrics_before = vllm_model.llm.get_metrics() - outputs = vllm_model.llm.generate(prompts, sampling_params) + full_params = make_params(candidate_ids) + for params in full_params: # the array form, as a top-k tensor would be + params.prompt_logprob_token_ids = np.array(params.prompt_logprob_token_ids) + outputs = vllm_model.llm.generate(token_prompts, full_params) for i, output in enumerate(outputs): scores = output.prompt_token_id_logprobs - assert scores is not None, f"Output {i} missing fixed-ID scores" + assert scores is not None, f"Output {i} missing per-row scores" expected_shape = ( len(output.prompt_token_ids) - 1 - start, - len(candidate_ids), + 4, ) assert scores.shape == expected_shape, ( f"Output {i} scored {scores.shape}, expected {expected_shape}" @@ -1414,6 +1468,23 @@ def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): assert math.isfinite(float(scores.min())) assert float(scores.max()) <= 1e-3, "logprobs must be <= 0" + # Ragged rows: each row keeps a prefix of the candidates and must + # reproduce that prefix of the full scores, with -inf padding. + ragged_params = make_params(lambda j: candidate_ids(j)[: 1 + j % 4]) + ragged_outputs = vllm_model.llm.generate(token_prompts, ragged_params) + for output, ragged_output in zip(outputs, ragged_outputs): + expected = np.full_like(output.prompt_token_id_logprobs, -np.inf) + for j in range(len(expected)): + n = 1 + j % 4 + expected[j, :n] = output.prompt_token_id_logprobs[j, :n] + np.testing.assert_allclose( + ragged_output.prompt_token_id_logprobs, expected, atol=1e-2 + ) + + # The row count is checked against the prompt at admission. + with pytest.raises(VLLMValidationError, match="scored rows"): + vllm_model.llm.generate(token_prompts[0], ragged_params[1]) + metrics_after = vllm_model.llm.get_metrics() preemptions_before = next( (m.value for m in metrics_before if m.name == "vllm:num_preemptions"), 0 @@ -1429,31 +1500,32 @@ def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): # unchunked, unpreempted run of the same requests. Batch composition moves # bf16 tail logprobs by a few percent; a misaligned row differs by nats. with VllmRunner( - "Qwen/Qwen3-0.6B", - max_model_len=512, - max_logprobs=64, - gpu_memory_utilization=0.25, + "Qwen/Qwen3-0.6B", max_model_len=512, gpu_memory_utilization=0.25 ) as reference: - reference_outputs = reference.llm.generate(prompts, sampling_params) - own_ids = [list(dict.fromkeys(o.prompt_token_ids)) for o in reference_outputs] + reference_outputs = reference.llm.generate(token_prompts, full_params) self_scored = reference.llm.generate( - prompts, + token_prompts, [ SamplingParams( - max_tokens=1, prompt_logprobs=0, prompt_logprob_token_ids=ids + max_tokens=1, + prompt_logprobs=0, + prompt_logprob_token_ids=[[t] for t in p["prompt_token_ids"][1:]], ) - for ids in own_ids + for p in token_prompts ], ) for output, ref in zip(outputs, reference_outputs): np.testing.assert_allclose( - output.prompt_token_id_logprobs, ref.prompt_token_id_logprobs, rtol=0.1 + output.prompt_token_id_logprobs, + ref.prompt_token_id_logprobs, + rtol=0.1, + atol=0.2, ) - for output, ids in zip(self_scored, own_ids): - scores = output.prompt_token_id_logprobs + for output in self_scored: + scores = output.prompt_token_id_logprobs[:, 0] for row, target in enumerate(output.prompt_token_ids[1:]): expected = output.prompt_logprobs[row + 1][target].logprob - assert scores[row, ids.index(target)] == pytest.approx(expected, abs=1e-3) + assert scores[row] == pytest.approx(expected, abs=1e-3) def test_prompt_logprob_token_ids_drop_partially_scored_prefills(monkeypatch): @@ -1465,13 +1537,7 @@ def test_prompt_logprob_token_ids_drop_partially_scored_prefills(monkeypatch): """ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - prompt = "The capital of France is Paris. " * 20 - sampling_params = SamplingParams( - temperature=0.0, - max_tokens=1, - prompt_logprob_token_ids=[10, 100, 1000], - skip_reading_prefix_cache=False, - ) + text = "The capital of France is Paris. " * 20 with VllmRunner( "Qwen/Qwen3-0.6B", @@ -1479,6 +1545,15 @@ def test_prompt_logprob_token_ids_drop_partially_scored_prefills(monkeypatch): enable_prefix_caching=True, gpu_memory_utilization=0.25, ) as vllm_model: + token_ids = vllm_model.llm.get_tokenizer()(text).input_ids + prompt = TokensPrompt(prompt_token_ids=token_ids) + sampling_params = SamplingParams( + temperature=0.0, + max_tokens=1, + prompt_logprob_token_ids=np.tile([10, 100, 1000], (len(token_ids) - 1, 1)), + skip_reading_prefix_cache=False, + ) + first = vllm_model.llm.generate([prompt], sampling_params)[0] assert first.prompt_token_id_logprobs is not None assert math.isfinite(float(first.prompt_token_id_logprobs.min())) diff --git a/tests/v1/sample/test_trace_replay_params.py b/tests/v1/sample/test_trace_replay_params.py index cf6e3a83f7..43f8266ded 100644 --- a/tests/v1/sample/test_trace_replay_params.py +++ b/tests/v1/sample/test_trace_replay_params.py @@ -42,7 +42,9 @@ def test_sampling_params_trace_field_rejects_empty_list(): def test_sampling_params_trace_field_rejects_prompt_scoring(): - params = SamplingParams(trace_decode_token_ids=[1], prompt_logprob_token_ids=[2, 3]) + params = SamplingParams( + trace_decode_token_ids=[1], prompt_logprob_token_ids=[[2, 3]] + ) with pytest.raises(ValueError, match="prompt_logprob_token_ids"): params._validate_trace_replay( _make_model_config(vocab_size=100), speculative_config=None diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index aa6a65661d..57831fe4d3 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -39,6 +39,7 @@ def _make_fake_speculator( and a draft model. """ fake_input_buffers = SimpleNamespace( + positions=torch.arange(max_num_tokens, dtype=torch.int64), query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32), seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32), @@ -113,6 +114,7 @@ def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound(): captured = _run_build(fake, num_reqs=3, num_reqs_padded=4, base=base, step=2) + assert torch.equal(captured["positions"], fake.input_buffers.positions[:4]) bound = captured["seq_lens_cpu_upper_bound"] assert isinstance(bound, torch.Tensor), ( "seq_lens_cpu_upper_bound must be a tensor, not None" diff --git a/tests/v1/streaming_input/test_async_llm_streaming.py b/tests/v1/streaming_input/test_async_llm_streaming.py index b532eed15f..21a49024ce 100644 --- a/tests/v1/streaming_input/test_async_llm_streaming.py +++ b/tests/v1/streaming_input/test_async_llm_streaming.py @@ -8,7 +8,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from vllm.engine.protocol import StreamingInput -from vllm.outputs import RequestOutput +from vllm.outputs import STREAM_FINISHED, RequestOutput from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import RequestOutputCollector @@ -170,3 +170,31 @@ async def test_generate_with_async_generator(): assert outputs[2].finished is True # Both inputs were processed assert inputs_received == ["Hello", " world"] + + +@pytest.mark.asyncio +async def test_empty_input_stream_finishes_without_engine_request(): + llm = AsyncLLM.__new__(AsyncLLM) + llm._validate_streaming_input_sampling_params = MagicMock() + llm.get_supported_tasks = AsyncMock(return_value=("generate",)) + llm._run_output_handler = MagicMock() + llm._add_request = AsyncMock() + + final_req = MagicMock() + final_req.request_id = "request-int" + llm.input_processor = MagicMock() + llm.input_processor.process_inputs.return_value = final_req + + async def empty_stream(): + return + yield # pragma: no cover + + queue = await llm._add_streaming_input_request( + "request", empty_stream(), SamplingParams(max_tokens=4) + ) + task = queue._input_stream_task + assert task is not None + await task + + assert queue.get_nowait() is STREAM_FINISHED + llm._add_request.assert_not_awaited() diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index fde3605a20..a87c647256 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -85,6 +85,7 @@ def test_pcp_draft_metadata_keeps_graph_padding_in_decode(cg_mode): speculator.max_model_len = speculator.draft_max_seq_len = 32 speculator.draft_is_prefilling = torch.zeros(4, dtype=torch.bool) speculator.input_buffers = SimpleNamespace( + positions=torch.arange(4, dtype=torch.int64), query_start_loc=torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32), seq_lens=torch.tensor([11, 21, 0, 0], dtype=torch.int32), ) diff --git a/tests/v1/worker/test_late_interaction_runner.py b/tests/v1/worker/test_late_interaction_runner.py index 4befbf72cc..35ce3670a7 100644 --- a/tests/v1/worker/test_late_interaction_runner.py +++ b/tests/v1/worker/test_late_interaction_runner.py @@ -137,6 +137,52 @@ def test_finished_request_releases_unscored_doc_use(): ) +@pytest.mark.skip_global_cleanup +def test_release_queries_is_targeted_and_idempotent(): + runner = LateInteractionRunner() + query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) + + for query_key in ("query-release", "query-keep"): + query_params = _make_pooling_params( + build_late_interaction_query_params(query_key=query_key, query_uses=1) + ) + runner.postprocess_pooler_output( + raw_pooler_output=[query_emb], + pooling_params=[query_params], + req_ids=[f"{query_key}-req"], + finished_mask=[True], + ) + + released_doc_params = _make_pooling_params( + build_late_interaction_doc_params(query_key="query-release") + ) + kept_doc_params = _make_pooling_params( + build_late_interaction_doc_params(query_key="query-keep") + ) + runner.register_request("released-doc", released_doc_params) + runner.register_request("kept-doc", kept_doc_params) + + runner.release_queries(["query-release"]) + runner.release_queries(["query-release"]) + + with pytest.raises(ValueError, match="query cache miss"): + runner.postprocess_pooler_output( + raw_pooler_output=[query_emb], + pooling_params=[released_doc_params], + req_ids=["released-doc"], + finished_mask=[True], + ) + + kept_output = runner.postprocess_pooler_output( + raw_pooler_output=[query_emb], + pooling_params=[kept_doc_params], + req_ids=["kept-doc"], + finished_mask=[True], + ) + assert isinstance(kept_output, list) + assert kept_output[0] is not None + + def test_invalid_query_uses_raises(): runner = LateInteractionRunner() bad_meta = LateInteractionParams( diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 03c01b2485..4de6022e34 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -26,6 +26,7 @@ from vllm.v1.watermarking.spec_decode import ( speculative_target_watermark_key, ) from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.watermark import ( draft_watermarking_mask, @@ -1499,6 +1500,53 @@ def test_draft_watermarker_deduplicates_across_a_block_on_cuda( assert not draft_logits[:, :, -1].any() +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_draft_watermarker_cuda_graph_reads_current_prompt_lens(): + draft_watermarker = DraftWatermarker( + GumbelWatermarker(key=42, context_width=2), + max_num_reqs=1, + device=torch.device("cuda"), + num_speculative_steps=1, + deduplicate_contexts="single_turn", + deduplicate_contexts_max_history=None, + ) + prompt_lens = UvaBackedTensor(1, dtype=torch.int32) + contexts = torch.tensor([[5, 6]]).cuda() + enabled = torch.ones(1, dtype=torch.bool).cuda() + # [5, 6] repeats only while the prompt counts as history. + all_token_ids = torch.tensor([[5, 6, 9, 5, 6]]).cuda() + total_lens = torch.tensor([5]).cuda() + logits = torch.zeros(1, 8).cuda() + idx_mapping = torch.zeros(1, dtype=torch.int32).cuda() + temperature = torch.ones(1).cuda() + draft_step = torch.tensor(0).cuda() + + def prepare(): + draft_watermarker.prepare( + contexts, enabled, all_token_ids, prompt_lens.gpu, total_lens + ) + + def watermark_mask(): + _, mask = draft_watermarker._sampling_state( + logits, idx_mapping, temperature, draft_step, contexts + ) + return mask + + prepare() + assert not watermark_mask().item() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + mask = watermark_mask() + # A step that adds requests moves prompt_lens.gpu to the next pool buffer. + prompt_lens.np[:] = 3 + prompt_lens.copy_to_uva() + prepare() + graph.replay() + assert mask.item() + + def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch): speculator = object.__new__(DSparkSpeculator) speculator.draft_logits = torch.zeros(2, 1, 8) diff --git a/tools/install_torchcodec_rocm.sh b/tools/install_torchcodec_rocm.sh index e5ce3571f2..d511159ab2 100755 --- a/tools/install_torchcodec_rocm.sh +++ b/tools/install_torchcodec_rocm.sh @@ -117,8 +117,7 @@ echo "Building TorchCodec (MAX_JOBS=$MAX_JOBS)..." pip wheel . --no-build-isolation --no-deps -w "$BUILD_DIR/dist" # Install the built wheel -# shellcheck disable=SC2012 # `ls` sorting is relied on here; TODO: switch to `find` in a follow-up cleanup PR -BUILT_WHEEL=$(ls "$BUILD_DIR/dist"/torchcodec-*.whl 2>/dev/null | head -1) +BUILT_WHEEL=$(find "$BUILD_DIR/dist" -maxdepth 1 -name 'torchcodec-*.whl' -print -quit) if [ -z "$BUILT_WHEEL" ]; then echo "Error: No wheel produced" exit 1 diff --git a/tools/pre_commit/check_agent_files.py b/tools/pre_commit/check_agent_files.py new file mode 100644 index 0000000000..90833f29d4 --- /dev/null +++ b/tools/pre_commit/check_agent_files.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check agent instruction files and skills against vLLM's conventions. + +Agents fail silently on a broken setup: a skill with no `.claude/skills` symlink +or with unparsable frontmatter is simply never offered, and a dead link reads as +missing guidance. See `docs/contributing/editing-agent-instructions.md` for the +rules enforced here. + +Usage: + python tools/pre_commit/check_agent_files.py +""" + +import os +import posixpath +import subprocess +import sys +from pathlib import Path + +import regex as re +import yaml + +SKILLS = ".agents/skills" +CLAUDE_SKILLS = ".claude/skills" +SYMLINK = "120000" +ROOT_GUIDE_MAX_LINES = 200 +DOMAIN_GUIDE_MAX_LINES = 300 +SKILL_MAX_LINES = 500 + +FRONTMATTER = re.compile(r"\A---\n(.*?)\n---\n", re.DOTALL) +SKILL_NAME = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*") +WHEN_TO_USE = re.compile(r"(?i)\buse\b") +CODE_BLOCK = re.compile(r"```.*?```", re.DOTALL) +LINK = re.compile(r"(? dict[str, str]: + """Map each tracked path to its git mode.""" + out = subprocess.run( + ["git", "ls-files", "-s", "-z"], capture_output=True, text=True, check=True + ).stdout + files = {} + for entry in out.split("\0")[:-1]: + meta, path = entry.split("\t", 1) + files[path] = meta.split()[0] + return files + + +def link_target(path: str) -> str: + # Without core.symlinks, git checks a symlink out as a file holding its target. + return os.readlink(path) if os.path.islink(path) else Path(path).read_text() + + +def skill_names(files: dict[str, str]) -> list[str]: + return sorted({p.split("/")[2] for p in files if p.startswith(f"{SKILLS}/")}) + + +def check_skill_links(files: dict[str, str]) -> list[str]: + errors = [] + skills = skill_names(files) + for name in skills: + link = f"{CLAUDE_SKILLS}/{name}" + target = f"../../{SKILLS}/{name}" + if files.get(link) != SYMLINK or link_target(link) != target: + errors.append(f"{link}: must be a symlink to {target}") + for path, mode in files.items(): + if not path.startswith(f"{CLAUDE_SKILLS}/"): + continue + if mode != SYMLINK or path.split("/")[2] not in skills: + errors.append(f"{path}: only symlinks to {SKILLS}/ belong here") + return errors + + +def check_skill_frontmatter(files: dict[str, str]) -> tuple[list[str], list[str]]: + errors, notes = [], [] + for name in skill_names(files): + path = f"{SKILLS}/{name}/SKILL.md" + if path not in files: + errors.append(f"{SKILLS}/{name}: missing SKILL.md") + continue + match = FRONTMATTER.match(Path(path).read_text()) + try: + meta = yaml.safe_load(match.group(1)) if match else None + except yaml.YAMLError as exc: + reason = getattr(exc, "problem", None) or exc + errors.append(f"{path}: frontmatter is not valid YAML ({reason})") + continue + if not isinstance(meta, dict): + errors.append(f"{path}: missing YAML frontmatter") + continue + if meta.get("name") != name: + errors.append(f"{path}: name {meta.get('name')!r} must be {name!r}") + if not SKILL_NAME.fullmatch(name) or len(name) > 64: + errors.append(f"{path}: name must be lowercase-hyphenated, <= 64 chars") + description = meta.get("description") + if not isinstance(description, str) or not description.strip(): + errors.append(f"{path}: missing description") + elif len(description) > 1024: + errors.append(f"{path}: description exceeds 1024 chars") + elif not WHEN_TO_USE.search(description): + notes.append(f"{path}: description should say when to use the skill") + return errors, notes + + +def check_no_claude_md(files: dict[str, str]) -> list[str]: + return [ + f"{path}: remove it; Claude Code reads AGENTS.md, but not where a " + "CLAUDE.md exists" + for path in files + if posixpath.basename(path) == "CLAUDE.md" + ] + + +def check_line_budgets(files: dict[str, str]) -> list[str]: + errors = [] + for path in files: + if path == "AGENTS.md": + limit = ROOT_GUIDE_MAX_LINES + elif posixpath.basename(path) == "AGENTS.md": + limit = DOMAIN_GUIDE_MAX_LINES + elif path.startswith(f"{SKILLS}/") and path.endswith("/SKILL.md"): + limit = SKILL_MAX_LINES + else: + continue + lines = len(Path(path).read_text().splitlines()) + if lines > limit: + errors.append(f"{path}: {lines} lines exceeds the {limit}-line budget") + return errors + + +def anchor_key(text: str) -> str: + # Anchor slugs differ between renderers, so compare alphanumerics only. + return re.sub(r"[^a-z0-9]", "", text.lower()) + + +def has_heading(path: str, anchor: str) -> bool: + headings = HEADING.findall(CODE_BLOCK.sub("", Path(path).read_text())) + return anchor_key(anchor) in {anchor_key(heading) for heading in headings} + + +def check_links(files: dict[str, str]) -> list[str]: + errors = [] + dirs = {str(parent) for p in files for parent in Path(p).parents} + for path in files: + in_skill = path.startswith(f"{SKILLS}/") and path.endswith(".md") + if not in_skill and posixpath.basename(path) != "AGENTS.md": + continue + text = CODE_BLOCK.sub("", Path(path).read_text()) + for link in LINK.findall(text): + if URL_SCHEME.match(link): + continue + ref, _, anchor = link.partition("#") + dest = posixpath.normpath(posixpath.join(posixpath.dirname(path), ref)) + if dest not in files and dest not in dirs: + errors.append(f"{path}: link {link} does not exist") + elif anchor and dest.endswith(".md") and not has_heading(dest, anchor): + errors.append(f"{path}: link {link} has no heading #{anchor}") + return errors + + +def check(files: dict[str, str]) -> tuple[list[str], list[str]]: + """Return (errors, advisory notes) for the tracked files under the cwd.""" + errors, notes = check_skill_frontmatter(files) + errors += check_skill_links(files) + errors += check_no_claude_md(files) + errors += check_line_budgets(files) + errors += check_links(files) + return errors, notes + + +def main() -> int: + errors, notes = check(tracked_files()) + for note in notes: + print(f"note: {note}") + if not errors: + return 0 + print(f"{len(errors)} agent file problem(s):\n", file=sys.stderr) + for error in errors: + print(f" {error}", file=sys.stderr) + print( + "\nSee docs/contributing/editing-agent-instructions.md.", + file=sys.stderr, + ) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index f2d313bdd0..ca98e63b39 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1913,6 +1913,11 @@ def _sync_aiter_situv2_moe_env() -> None: os.environ.pop(name, None) +_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 +_GFX950_C4A_NATIVE_MAX_ROWS = 256 +_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 + + class rocm_aiter_ops: """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM. @@ -2546,6 +2551,100 @@ class rocm_aiter_ops: """Whether AITER has DSV4 native-I384 configs through the given M.""" return cls._probe_dsv4_i384_fhmoe_capability(num_tokens) + @classmethod + @if_aiter_supported + def is_indexer_top_k_enabled(cls) -> bool: + """gfx950 is the only arch with tuned AITER indexer top-k kernels.""" + from vllm.platforms.rocm import on_gfx950 + + return cls._AITER_ENABLED and on_gfx950() + + @classmethod + def is_indexer_top_k_supported( + cls, + *, + is_prefill: bool, + compress_ratio: int, + num_rows: int, + max_valid_seq_len: int | None = None, + ) -> bool: + """Whether AITER's sparse indexer top-k beats the in-tree kernel for + this shape.""" + if compress_ratio <= 1 or not cls.is_indexer_top_k_enabled(): + return False + + if not is_prefill: + assert max_valid_seq_len is not None + # AITER v0.1.19 decode is one-block only. This measured gfx950 + # FP32/k=1024 compressed-row boundary is independent of the native + # split-count boundary in sampler.cu. + if ( + num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS + and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN + ): + return False + + return True + + @staticmethod + def dsv4_indexer_prefers_native_top_k( + *, + num_rows: int, + num_columns: int, + topk_tokens: int, + ) -> bool: + """The in-tree decode kernel's measured advantage window over AITER. + Tuned on DSV4's compressed-KV logits, so it applies to that indexer + only.""" + return ( + topk_tokens == 512 + and 0 < num_rows <= 384 + and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS + ) + + @staticmethod + def indexer_top_k_decode( + logits: torch.Tensor, + next_n: int, + seq_lens: torch.Tensor, + topk_indices: torch.Tensor, + topk_tokens: int, + ) -> None: + from aiter.ops.topk import top_k_per_row_decode + + top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + k=topk_tokens, + ) + + @staticmethod + def indexer_top_k_prefill( + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + indices: torch.Tensor, + topk_tokens: int, + ) -> None: + from aiter.ops.topk import top_k_per_row_prefill + + top_k_per_row_prefill( + logits, + row_starts, + row_ends, + indices, + None, + logits.shape[0], + logits.stride(0), + logits.stride(1), + k=topk_tokens, + ) + @staticmethod def register_ops_once() -> None: global _OPS_REGISTERED diff --git a/vllm/config/ec_transfer.py b/vllm/config/ec_transfer.py index 5c7f99b4aa..5c2b8e589c 100644 --- a/vllm/config/ec_transfer.py +++ b/vllm/config/ec_transfer.py @@ -42,15 +42,6 @@ class ECTransferConfig: """Whether this vLLM instance produces, consumes EC cache, or both. Choices are 'ec_producer', 'ec_consumer', 'ec_both'.""" - ec_rank: int | None = None - """The rank of this vLLM instance in the EC cache transfer. Typical value: - 0 for encoder, 1 for pd instance. - Currently only 1P1D is supported.""" - - ec_parallel_size: int = 1 - """The number of parallel instances for EC cache transfer. For - PyNcclConnector, this should be 2.""" - ec_ip: str = "127.0.0.1" """The EC connector ip, used to build distributed connection.""" diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index a42ad62047..df69f17149 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -3,7 +3,7 @@ import contextlib from collections.abc import Callable, Iterable from dataclasses import asdict, fields -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Final, Literal from pydantic import Field, field_validator @@ -12,6 +12,7 @@ from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.config.parallel import ParallelConfig logger = init_logger(__name__) @@ -128,10 +129,10 @@ MoEBackend = Literal[ "flashinfer_trtllm", "flashinfer_cutlass", "flashinfer_cutedsl", + "flashinfer_moe_ep_cutedsl", "flashinfer_b12x", "b12x", "flashinfer_moe_ep_mega_deep_gemm", - "flashinfer_moe_ep_mega_cutedsl", "marlin", "humming", "triton_unfused", @@ -143,21 +144,28 @@ MoEBackend = Literal[ "rdna3", ] -# Backends that run the mega-MoE model path through the flashinfer moe_ep -# runtime. Only architectures in FLASHINFER_MOE_EP_ARCHITECTURES wire up -# these experts. +FLASHINFER_MOE_EP_CUTEDSL = "flashinfer_moe_ep_cutedsl" +FLASHINFER_MOE_EP_DEEP_GEMM = "flashinfer_moe_ep_mega_deep_gemm" + +# Backends implemented by the shared FlashInfer MoE-EP adapter. FLASHINFER_MOE_EP_BACKENDS = frozenset( { - "flashinfer_moe_ep_mega_deep_gemm", - "flashinfer_moe_ep_mega_cutedsl", + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, } ) +# Those backends dispatch and combine themselves, so the framework must not +# run a separate all2all around them. +PASSTHROUGH_ALL2ALL_BACKEND: Final = "passthrough" -# Backends that run a mega-MoE model path (fused expert module plus -# prepare_megamoe routing): vLLM's native deep_gemm path, which any model -# with a mega-MoE module may use (DeepSeek-V4, Kimi K3), plus the flashinfer -# moe_ep variants. -MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS +# "Native" here meaning that the backend does not use the fused MoE path. +# Currently only relevant for the vLLM provided DeepGEMM MegaMoE., +# TODO: migrate to the fused MoE path ultimately. +NATIVE_MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) + +# Every megakernel backend. The FlashInfer variants run through the standard +# fused-MoE model path; the quantization method binds the kernel. +MEGA_MOE_BACKENDS = NATIVE_MEGA_MOE_BACKENDS | FLASHINFER_MOE_EP_BACKENDS SparseIndexerTopkBackend = Literal[ "auto", @@ -167,10 +175,11 @@ SparseIndexerTopkBackend = Literal[ "per_row", "flashinfer", "torch", + "aiter", ] -# Architectures whose model code wires up the flashinfer moe_ep experts. MTP -# and DSpark draft variants inherit the setting from these target models. +# DeepGEMM consumes the DeepSeek-V4 MXFP4 checkpoint recipe. CuTeDSL is +# selected by quantization methods and is not architecture-specific. FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( { "DeepseekV4ForCausalLM", @@ -180,11 +189,48 @@ FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( ) +def bind_passthrough_all2all_backend( + moe_backend: str, parallel_config: "ParallelConfig" +) -> None: + """Bind ``all2all_backend="passthrough"`` to MoE backends that own their + expert-parallel communication (currently the FlashInfer MoE-EP megakernels). + + Such a deployment gets ``passthrough`` when the all2all backend is left at + its default and rejects any other explicit choice; the value is meaningless + with a MoE backend that relies on the framework to communicate. The CLI does + not offer the value; the early return keeps repeated + ``set_platform_defaults`` calls idempotent. + """ + a2a = parallel_config.all2all_backend + if moe_backend in FLASHINFER_MOE_EP_BACKENDS: + if a2a == PASSTHROUGH_ALL2ALL_BACKEND: + return + if a2a != "allgather_reducescatter": + raise ValueError( + f"moe_backend={moe_backend!r} dispatches and combines itself and " + f"cannot be paired with all2all_backend={a2a!r}; leave " + f"--all2all-backend unset." + ) + logger.info_once( + "moe_backend=%r: using all2all_backend=%r (dispatch and combine run " + "inside the MoE backend).", + moe_backend, + PASSTHROUGH_ALL2ALL_BACKEND, + ) + parallel_config.all2all_backend = PASSTHROUGH_ALL2ALL_BACKEND + elif a2a == PASSTHROUGH_ALL2ALL_BACKEND: + raise ValueError( + f"all2all_backend={a2a!r} is only valid with a moe_backend that " + f"dispatches and combines itself ({sorted(FLASHINFER_MOE_EP_BACKENDS)})," + f" got moe_backend={moe_backend!r}." + ) + + def validate_flashinfer_moe_ep_model( moe_backend: str, architectures: Iterable[str] ) -> None: - """Reject flashinfer moe_ep backends for models that lack the FI path.""" - if moe_backend not in FLASHINFER_MOE_EP_BACKENDS: + """Reject model-specific FlashInfer MoE-EP backends.""" + if moe_backend != FLASHINFER_MOE_EP_DEEP_GEMM: return if not any(arch in FLASHINFER_MOE_EP_ARCHITECTURES for arch in architectures): raise ValueError( @@ -254,6 +300,9 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) + - "flashinfer_moe_ep_cutedsl": Use FlashInfer's CuTeDSL MoE-EP + mega-kernel with NVFP4 weights (MXFP4 checkpoints are requantized at + load); requires Blackwell, expert parallelism, and NVSHMEM - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) - "b12x": Use b12x FP4 MoE kernels on SM12x @@ -261,10 +310,6 @@ class KernelConfig: expert-parallel mega-MoE with the DeepGEMM megakernel, which consumes an MXFP4 checkpoint verbatim (Blackwell, requires expert parallel; DeepSeek-V4 only) - - "flashinfer_moe_ep_mega_cutedsl": Same, with the CuteDSL megakernel - (additionally requires NVSHMEM). The checkpoint selects the weight path: - an NVFP4 checkpoint is consumed prequantized, MXFP4 weights are - requantized at load - "marlin": Use Marlin kernels (weight-only quantization) - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels @@ -289,6 +334,8 @@ class KernelConfig: - "per_row": Use vLLM's top_k_per_row_decode kernel - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel - "torch": Use a plain torch.topk implementation (debug reference) + - "aiter": Use AITER's top_k_per_row_decode kernel (ROCm gfx950 only, + requires VLLM_ROCM_USE_AITER) Explicit values raise RuntimeError when their constraints are not met. """ @@ -388,6 +435,7 @@ class KernelConfig: validate_flashinfer_moe_ep_model( self.moe_backend, vllm_config.model_config.architectures ) + bind_passthrough_all2all_backend(self.moe_backend, vllm_config.parallel_config) platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config) logger.debug( diff --git a/vllm/config/kv_transfer.py b/vllm/config/kv_transfer.py index dfc9e02f3d..59eb4c80be 100644 --- a/vllm/config/kv_transfer.py +++ b/vllm/config/kv_transfer.py @@ -68,25 +68,10 @@ class KVTransferConfig: """The device used by kv connector to buffer the KV cache. Choices are 'cuda', 'cpu' and 'xpu'.""" - kv_buffer_size: float = 1e9 - """The buffer size for TorchDistributedConnector. Measured in number of - bytes. Recommended value: 1e9 (about 1GB).""" - kv_role: KVRole | None = None """Whether this vLLM instance produces, consumes KV cache, or both. Choices are 'kv_producer', 'kv_consumer', and 'kv_both'.""" - kv_rank: int | None = None - """The rank of this vLLM instance in the KV cache transfer. Typical value: - 0 for prefill instance, 1 for decode instance. - Currently only 1P1D is supported.""" - - kv_parallel_size: int = 1 - """The number of parallel instances for KV cache transfer.""" - - kv_ip: str = "127.0.0.1" - """The KV connector ip, used to build distributed connection.""" - kv_port: int = 14579 """The KV connector port, used to build distributed connection.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 4de030e8f3..75ca774fa1 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -124,6 +124,22 @@ _RUNNER_CONVERTS: dict[RunnerType, list[ConvertType]] = { "draft": [], } + +def _modelopt_mixed_has_nvfp4(quant_config: dict[str, Any] | None) -> bool: + """Whether a ModelOpt MIXED_PRECISION checkpoint quantizes any layer to NVFP4. + + ``W4A16_NVFP4`` is excluded on purpose: it quantizes weights only, so there + is no activation quantization for the fusion passes to act on. + """ + layers = (quant_config or {}).get("quantized_layers") + if not isinstance(layers, dict): + return False + return any( + isinstance(info, dict) and info.get("quant_algo") == "NVFP4" + for info in layers.values() + ) + + AttnTypeStr = Literal[ "decoder", "encoder", "encoder_only", "encoder_decoder", "attention_free", "hybrid" ] @@ -2249,13 +2265,20 @@ class ModelConfig: return getattr(self.hf_config, "quantization_config", None) is not None def is_nvfp4_quantized(self) -> bool: + quant_config = self.model_arch_config.quantization_config + # ModelOpt NVFP4 checkpoints resolve to modelopt_fp4 quantization method if self.quantization in ("modelopt_fp4",): return True + # A checkpoint mixing NVFP4 with another algorithm declares + # quant_algo MIXED_PRECISION and resolves to modelopt_mixed, so the + # per-layer algorithms decide. + if self.quantization == "modelopt_mixed": + return _modelopt_mixed_has_nvfp4(quant_config) + # For Compressed Tensors we look for `"format": "nvfp4-pack-quantized"` # in the quantization config - quant_config = self.model_arch_config.quantization_config return ( self.quantization == "compressed-tensors" and quant_config is not None diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 53a28fa691..7df8f4e0a4 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -56,6 +56,7 @@ All2AllBackend = Literal[ "flashinfer_all2allv", # temporary alias for flashinfer_nvlink_two_sided "flashinfer_nvlink_two_sided", "flashinfer_nvlink_one_sided", + "passthrough", ] @@ -105,6 +106,17 @@ class EPLBConfig: on CUDA, falls back to "torch_gloo") """ + enable_migration_batching: bool = False + """Schedule expert migrations in batches where each rank communicates with + at most one peer. This reduces per-rank network contention at the cost of + additional sequential communication steps. This option only applies to + async EPLB and is disabled by default.""" + + @property + def migration_batching_enabled(self) -> bool: + """Whether contention-aware batching is active for this configuration.""" + return self.use_async and self.enable_migration_batching + @model_validator(mode="after") def _validate_eplb_config(self) -> Self: if self.use_async and self.policy != "default": @@ -211,7 +223,10 @@ class ParallelConfig: - "moonep": MoonEP balanced EP with dynamic redundant experts (NVLink) - "nixl_ep": Use nixl-ep kernels - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels - - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl""" + - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl + - "passthrough": No all2all at all: the MoE backend dispatches and combines + itself (e.g. the FlashInfer MoE-EP megakernels). Bound automatically for + such MoE backends; not selectable from the CLI.""" max_parallel_loading_workers: int | None = Field(default=None, ge=1) """Maximum number of parallel loading workers when loading model @@ -277,9 +292,6 @@ class ParallelConfig: worker_cls: str = "auto" """The full name of the worker class to use. If "auto", the worker class will be determined based on the platform.""" - sd_worker_cls: str = "auto" - """The full name of the worker class to use for speculative decoding. - If "auto", the worker class will be determined based on the platform.""" worker_extension_cls: str = "" """The full name of the worker extension class to use. The worker extension class is dynamically inherited by the worker class. This is used to inject @@ -745,6 +757,7 @@ class ParallelConfig: "deepep_low_latency", "deepep_v2", "flashinfer_nvlink_one_sided", + "passthrough", "mori_high_throughput", "mori_low_latency", "nixl_ep", @@ -894,7 +907,6 @@ class ParallelConfig: "placement_group", "distributed_executor_backend", "worker_cls", - "sd_worker_cls", "worker_extension_cls", "_api_process_count", "_api_process_rank", diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 7d32a17308..d950084726 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -22,6 +22,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static32, kInt8StaticChannelSym, kMxfp4Dynamic, kMxfp4Static, @@ -40,6 +41,8 @@ QUANT_KEY_NAMES: dict[str, QuantKey] = { "mxfp8": kMxfp8Dynamic, "mxfp4": kMxfp4Dynamic, "int8_per_channel_static": kInt8StaticChannelSym, + # Load-time MXFP4-to-int4 for Kimi-K3 on gfx942. + "int4_per_group_32": kInt4Static32, } diff --git a/vllm/device_allocator/sleep_mode_backend.py b/vllm/device_allocator/sleep_mode_backend.py index 023af04275..8ec9986ffc 100644 --- a/vllm/device_allocator/sleep_mode_backend.py +++ b/vllm/device_allocator/sleep_mode_backend.py @@ -129,7 +129,9 @@ class CuMemBackend(SleepModeBackend): self._state = "SUSPENDED" allocator = get_mem_allocator_instance() - allocator.sleep(offload_tags=("weights",) if level == 1 else tuple()) + # Runtime state is kept at every level; weights only at level 1. + offload_tags = ("weights", "runtime") if level == 1 else ("runtime",) + allocator.sleep(offload_tags=offload_tags) def resume(self, tags: list[str] | None = None) -> None: from vllm.device_allocator import get_mem_allocator_instance diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 5f339540b4..d3c48756ca 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -49,6 +49,20 @@ if has_flashinfer_nvlink_one_sided(): logger = init_logger(__name__) +class PassThroughAll2AllManager(All2AllManagerBase): + """Placeholder for ``all2all_backend="passthrough"``. + + The MoE backend dispatches and combines itself, so there is no all2all to + manage. + """ + + def get_handle(self, kwargs): + raise RuntimeError( + "passthrough has no all2all handle: dispatch and combine run " + "inside the MoE backend." + ) + + class AgRsAll2AllManager(All2AllManagerBase): """An implementation of all2all communication based on all-gather (dispatch) and reduce-scatter (combine). @@ -721,6 +735,7 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): rank: int world_size: int + low_precision_combine: bool = False def __init__(self, cpu_group): assert has_flashinfer_nvlink_one_sided(), ( @@ -741,6 +756,27 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): self.top_k = 0 self.num_experts = 0 self._combine_supports_output = False + self.low_precision_combine = self._resolve_low_precision_combine() + + def _resolve_low_precision_combine(self) -> bool: + """Whether to use the low-precision combine: requested via the env var + and supported by the installed FlashInfer. + """ + if not envs.VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE: + return False + try: + supported = supports_kw( + MoeAlltoAll.combine, "use_low_precision", allow_var_kwargs=False + ) + except (TypeError, ValueError): + supported = False + if not supported: + logger.warning_once( + "VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE is set, but the " + "installed FlashInfer MoeAlltoAll.combine() does not accept " + "`use_low_precision`. Falling back to a BF16 combine." + ) + return supported def initialize( self, @@ -758,7 +794,10 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): + top_k * 4 # int32 topks ids + top_k * 4 # float32 topk weights ) - combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states + # Sized from the bf16 payload passed to combine(), which is what the + # kernel checks this region against. Low-precision transport quantizes + # on write, so it does not shrink the requirement. + combine_payload_size_per_token = hidden_size * 2 needed_workspace_size = moe_a2a_get_workspace_size_per_rank( ep_size=self.world_size, max_num_tokens=max_num_tokens, @@ -874,8 +913,10 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase): payload=payload, runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, output=output, + use_low_precision=self.low_precision_combine, ) else: + # FlashInfer < 0.6.16 has neither `output` nor `use_low_precision`. combined_output = self.moe_alltoall.combine( payload=payload, runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index ef790f7732..0a203956f0 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -255,6 +255,12 @@ class CudaCommunicator(DeviceCommunicatorBase): from .all2all import FlashInferNVLinkOneSidedManager self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group) + elif self.all2all_backend == "passthrough": + from .all2all import PassThroughAll2AllManager + + self.all2all_manager = PassThroughAll2AllManager( + self.cpu_group, tcp_store_group + ) else: raise ValueError(f"Unknown all2all backend: {self.all2all_backend}") diff --git a/vllm/distributed/device_communicators/shm_object_storage.py b/vllm/distributed/device_communicators/shm_object_storage.py index 3da8c5dcf6..081fbea197 100644 --- a/vllm/distributed/device_communicators/shm_object_storage.py +++ b/vllm/distributed/device_communicators/shm_object_storage.py @@ -527,6 +527,9 @@ class SingleWriterShmObjectStorage: self.id_index: dict[int, str] = {} # Writer flag to track in-use status: monotonic_id -> count self.writer_flag: dict[int, int] = {} + # Items touch() keeps from being freed until they are looked up + # or release_touches() is called + self._touched: set[int] = set() else: if reader_lock is None: raise ValueError("Lock must be provided for readers.") @@ -540,6 +543,7 @@ class SingleWriterShmObjectStorage: self.key_index.clear() self.id_index.clear() self.writer_flag.clear() + self._touched.clear() logger.debug("Object storage cleared and reinitialized.") def copy_to_buffer( @@ -664,8 +668,13 @@ class SingleWriterShmObjectStorage: """Get the cached object by key if it exists.""" address, monotonic_id = self.key_index[key] self.increment_writer_flag(monotonic_id) + self._touched.discard(monotonic_id) return address, monotonic_id + def release_touches(self) -> None: + """Stop protecting touched items that were not looked up.""" + self._touched.clear() + def get_signature(self, key: str) -> list[int]: """Sign the handle of a cached object so readers can verify it.""" address, monotonic_id = self.key_index[key] @@ -756,10 +765,11 @@ class SingleWriterShmObjectStorage: monotonic_id: int = 0, signature: list[int] | None = None, ) -> None: - """Touch an existing cached item to update its eviction status. + """Touch an existing cached item. - For writers (ShmObjectStoreSenderCache): Increment writer_flag - For readers (ShmObjectStoreReceiverCache): Increment reader_count + For writers (ShmObjectStoreSenderCache): Protect the item from + eviction until get_cached or release_touches + For readers (ShmObjectStoreReceiverCache): Validate the handle Args: key: String key of the object to touch @@ -771,25 +781,12 @@ class SingleWriterShmObjectStorage: if self._reader_lock is None: if key not in self.key_index: return None - address, monotonic_id = self.key_index[key] - # Writer side: increment writer_flag to raise eviction threshold - self.increment_writer_flag(monotonic_id) + self._touched.add(self.key_index[key][1]) else: # Reject forged handles before dereferencing the supplied address. self.verify_signature(key, address, monotonic_id, signature) - with ( - self._reader_lock, - self.ring_buffer.access_buf(address) as (data_view, buf_metadata), - ): + with self.ring_buffer.access_buf(address) as (_, buf_metadata): self._validate_monotonic_id(address, monotonic_id, buf_metadata) - reader_count = self.ring_buffer.byte2int(data_view[: self.flag_bytes]) - - # NOTE(Long): - # Avoid increasing flag on newly added item (sync with sender) - # Since when a new item is added - # pre-touch has no effect on writer side - if reader_count >= self.n_readers: - self.increment_reader_flag(data_view[: self.flag_bytes]) def close(self) -> None: """Close the shared memory.""" @@ -823,6 +820,8 @@ class SingleWriterShmObjectStorage: """Default is_free function that checks if the first 4 bytes are zero. This indicates that the buffer is free. """ + if id in self._touched: + return False reader_count = int.from_bytes(buf[0:4], "little", signed=True) writer_count = self.writer_flag[id] return reader_count >= writer_count * self.n_readers diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py index 3b7bc908ec..0f09bb99c4 100644 --- a/vllm/distributed/eplb/async_worker.py +++ b/vllm/distributed/eplb/async_worker.py @@ -11,6 +11,7 @@ from vllm.distributed.parallel_state import get_eplb_group from vllm.logger import init_logger from .eplb_utils import CpuGpuEvent, device_stream +from .migration_scheduler import schedule_migration_batches_for_layers from .rebalance_execute import AsyncEplbLayerResult, transfer_layer if TYPE_CHECKING: @@ -99,6 +100,21 @@ def transfer_run_periodically( model_state, state, physical_to_logical_map_cpu, stream ) + migration_batches_by_layer = None + if state.parallel_config.eplb_config.migration_batching_enabled: + profile_name = "eplb: schedule migration batches" + num_local_experts = model_state.model.expert_weights[0][0].shape[0] + # The placement snapshots stay fixed for this rebalance cycle. + # Precompute all layer schedules once, but execute transfers one + # layer at a time below to preserve the migration lifecycle. This + # profiling range therefore covers scheduling for every MoE layer. + with torch.profiler.record_function(profile_name): + migration_batches_by_layer = schedule_migration_batches_for_layers( + num_local_experts, + physical_to_logical_map_cpu.numpy(), + new_physical_to_logical_map.numpy(), + ) + # Execute one EPLB layer transfer per model forward pass. Each iteration # of this loop will copy the new set of expert weights into # model_state.expert_buffer, which will be consumed by the main thread in @@ -134,6 +150,16 @@ def transfer_run_periodically( is_profile=is_profile, stream=stream, layer_idx=layer_idx, + enable_migration_batching=( + state.parallel_config.eplb_config.migration_batching_enabled + ), + migration_batches=( + # This is only the current layer's schedule; it does not + # combine transfers from different layers. + migration_batches_by_layer[layer_idx] + if migration_batches_by_layer is not None + else None + ), ) # Wait until all writes to expert_buffer have finished before making the diff --git a/vllm/distributed/eplb/migration_scheduler.py b/vllm/distributed/eplb/migration_scheduler.py new file mode 100644 index 0000000000..1844d79b0b --- /dev/null +++ b/vllm/distributed/eplb/migration_scheduler.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Schedule EPLB expert migrations into conflict-free batches.""" + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True, slots=True) +class MigrationFlow: + """Expert transfers between one directed rank pair.""" + + src_rank: int + dst_rank: int + # Experts being transferred. + expert_ids: tuple[int, ...] + + +def schedule_migration_batches( + num_local_experts: int, + old_indices: np.ndarray, + new_indices: np.ndarray, +) -> list[list[MigrationFlow]]: + """Schedule expert migrations without per-rank communication contention. + + Experts assigned to the same source and destination ranks form one flow. + Each flow is placed in the first batch that does not already use either + endpoint, so a rank communicates with at most one peer in each batch while + independent rank pairs can transfer concurrently. + """ + assert old_indices.shape == new_indices.shape + recv_ranks_by_expert: dict[int, list[int]] = {} + old_experts_by_rank: list[set[int]] = [] + old_by_rank = old_indices.reshape(-1, num_local_experts) + new_by_rank = new_indices.reshape(-1, num_local_experts) + + # First find the ranks that gain each expert. Existing local copies need no + # transfer, even if their physical row changes. + for rank, (old_local, new_local) in enumerate(zip(old_by_rank, new_by_rank)): + old_experts = set(old_local.tolist()) + old_experts.discard(-1) + old_experts_by_rank.append(old_experts) + + new_experts = set(new_local.tolist()) + new_experts.discard(-1) + for expert_id in new_experts - old_experts: + recv_ranks_by_expert.setdefault(expert_id, []).append(rank) + + old_ranks_by_expert: dict[int, list[int]] = {} + needed_experts = recv_ranks_by_expert.keys() + + # Only index source copies for experts that need a remote transfer. + for rank, old_experts in enumerate(old_experts_by_rank): + for expert_id in old_experts & needed_experts: + old_ranks_by_expert.setdefault(expert_id, []).append(rank) + + expert_ids_by_pair: dict[tuple[int, int], list[int]] = {} + for expert_id in sorted(recv_ranks_by_expert): + send_ranks = old_ranks_by_expert.get(expert_id, []) + if not send_ranks: + continue + recv_ranks = recv_ranks_by_expert[expert_id] + + num_per_sender, remainder = divmod(len(recv_ranks), len(send_ranks)) + remainder_start = len(send_ranks) * num_per_sender + + # Preserve the unbatched path's balanced source assignment. Batching + # changes transfer order, not which existing replica supplies a copy. + for sender_idx, src_rank in enumerate(send_ranks): + start = sender_idx * num_per_sender + for dst_rank in recv_ranks[start : start + num_per_sender]: + expert_ids_by_pair.setdefault((src_rank, dst_rank), []).append( + expert_id + ) + if sender_idx < remainder: + dst_rank = recv_ranks[remainder_start + sender_idx] + expert_ids_by_pair.setdefault((src_rank, dst_rank), []).append( + expert_id + ) + + # Transfers with the same directed rank pair have identical contention + # constraints, so coalesce them before grouping. The loop below is a + # deterministic first-fit edge coloring of the rank communication graph. + batches: list[list[MigrationFlow]] = [] + endpoints_used: list[int] = [] + for (src_rank, dst_rank), expert_ids in expert_ids_by_pair.items(): + flow = MigrationFlow(src_rank, dst_rank, tuple(expert_ids)) + + # Direction is irrelevant for NIC contention. This endpoint mask keeps + # each rank connected to at most one peer in a batch. The conservative + # limit leaves bandwidth headroom for inference during async migration. + endpoints = (1 << src_rank) | (1 << dst_rank) + for batch_idx, (batch, used) in enumerate(zip(batches, endpoints_used)): + if not endpoints & used: + batch.append(flow) + endpoints_used[batch_idx] |= endpoints + break + else: + batches.append([flow]) + endpoints_used.append(endpoints) + + return batches + + +def schedule_migration_batches_for_layers( + num_local_experts: int, + old_indices: np.ndarray, + new_indices: np.ndarray, +) -> list[list[list[MigrationFlow]]]: + """Precompute the independent migration schedule for every MoE layer. + + This reduces scheduling entry points from once per layer to once per + rebalance cycle. Only CPU scheduling is grouped; transfers remain ordered + by layer and preserve each layer's batch boundaries. + """ + assert old_indices.shape == new_indices.shape + assert old_indices.ndim == 2 + return [ + schedule_migration_batches(num_local_experts, old_layer, new_layer) + for old_layer, new_layer in zip(old_indices, new_indices) + ] + + +__all__ = [ + "MigrationFlow", + "schedule_migration_batches", + "schedule_migration_batches_for_layers", +] diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 7a52dd5f56..5479ef01df 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -14,6 +14,10 @@ from torch.distributed import ProcessGroup, all_gather from vllm.distributed.eplb.eplb_communicator import EplbCommunicator from vllm.distributed.eplb.eplb_utils import CpuGpuEvent, device_stream +from vllm.distributed.eplb.migration_scheduler import ( + MigrationFlow, + schedule_migration_batches, +) from vllm.logger import init_logger from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -167,6 +171,107 @@ def get_ep_ranks_with_experts_batch( return ranks_to_send_map, ranks_to_recv_map +def _execute_migration_batch( + batch: list[MigrationFlow], + old_indices: np.ndarray, + old_rows: dict[int, int], + new_rows: dict[int, int], + ep_rank: int, + expert_weights: Sequence[torch.Tensor], + expert_weights_buffers: Sequence[torch.Tensor], + communicator: EplbCommunicator, + layer_idx: int, +) -> None: + """Execute one batch of rank-pair-disjoint expert transfers.""" + # NIXL drains its transfer context in execute(), so each batch restores it. + communicator.set_transfer_context(old_indices, layer_idx) + for flow in batch: + if flow.src_rank == ep_rank: + for expert_id in flow.expert_ids: + src_row = _get_migration_row( + old_rows, expert_id, flow, ep_rank, mapping_name="old" + ) + communicator.add_send( + [weight[src_row] for weight in expert_weights], + flow.dst_rank, + expert_id=expert_id, + ) + elif flow.dst_rank == ep_rank: + for expert_id in flow.expert_ids: + dst_row = _get_migration_row( + new_rows, expert_id, flow, ep_rank, mapping_name="new" + ) + communicator.add_recv( + [buffer[dst_row] for buffer in expert_weights_buffers], + flow.src_rank, + expert_id=expert_id, + ) + # Combining execute calls would erase the scheduler's batch boundaries. + communicator.execute() + + +def _execute_migration_batches( + num_local_experts: int, + old_indices: np.ndarray, + new_indices: np.ndarray, + ep_rank: int, + expert_weights: Sequence[torch.Tensor], + expert_weights_buffers: Sequence[torch.Tensor], + communicator: EplbCommunicator, + layer_idx: int, + batches: list[list[MigrationFlow]] | None = None, +) -> None: + """Execute all contention-aware migration batches for one layer.""" + if batches is None: + with torch.profiler.record_function("eplb: schedule migration batches"): + batches = schedule_migration_batches( + num_local_experts, old_indices, new_indices + ) + base = ep_rank * num_local_experts + old_local = old_indices[base : base + num_local_experts] + new_local = new_indices[base : base + num_local_experts] + old_rows: dict[int, int] = {} + new_rows: dict[int, int] = {} + for row, expert_id in enumerate(old_local): + if expert_id != -1: + old_rows.setdefault(int(expert_id), row) + for row, expert_id in enumerate(new_local): + if expert_id != -1: + new_rows.setdefault(int(expert_id), row) + for batch in batches: + _execute_migration_batch( + batch=batch, + old_indices=old_indices, + old_rows=old_rows, + new_rows=new_rows, + ep_rank=ep_rank, + expert_weights=expert_weights, + expert_weights_buffers=expert_weights_buffers, + communicator=communicator, + layer_idx=layer_idx, + ) + + +def _get_migration_row( + rows: dict[int, int], + expert_id: int, + flow: MigrationFlow, + ep_rank: int, + mapping_name: str, +) -> int: + row = rows.get(expert_id) + if row is not None: + return row + + message = ( + f"Expert {expert_id} is missing from the {mapping_name} " + f"mapping on rank {ep_rank} for migration " + f"{flow.src_rank} -> {flow.dst_rank}" + ) + logger.error(message) + raise RuntimeError(message) + + def move_to_buffer( num_local_experts: int, old_indices: np.ndarray, @@ -177,6 +282,8 @@ def move_to_buffer( ep_rank: int, communicator: EplbCommunicator, layer_idx: int = 0, + enable_migration_batching: bool = False, + migration_batches: list[list[MigrationFlow]] | None = None, ) -> TransferMetadata: """Rearranges expert weights during EPLB rebalancing. @@ -192,6 +299,9 @@ def move_to_buffer( ep_rank: Rank of this process in expert parallel group. communicator: EplbCommunicator instance for P2P communication. layer_idx: Index of the MoE layer being transferred. + enable_migration_batching: Schedule remote transfers in batches where + each rank communicates with at most one peer. + migration_batches: Optional batches precomputed for this layer. Returns: TransferMetadata: Metadata needed for completing remote weight transfers. @@ -264,6 +374,30 @@ def move_to_buffer( with device_stream(stream): for w, b in zip(expert_weights, expert_weights_buffers): b[dst].copy_(w[src_local], non_blocking=True) + transfer_metadata = TransferMetadata( + is_unchanged=is_unchanged, + is_received_locally=is_received_locally, + recv_primary_mask=recv_primary_mask, + recv_count=recv_count, + recv_expert_ids=recv_expert_ids, + recv_dst_rows=recv_dst_rows, + ) + + # Migration batching is disabled by default and must be explicitly enabled. + if enable_migration_batching: + _execute_migration_batches( + num_local_experts=num_local_experts, + old_indices=old_indices, + new_indices=new_indices, + ep_rank=ep_rank, + expert_weights=expert_weights, + expert_weights_buffers=expert_weights_buffers, + communicator=communicator, + layer_idx=layer_idx, + batches=migration_batches, + ) + return transfer_metadata + communicator.set_transfer_context(old_indices, layer_idx) # 2. Post sends @@ -334,14 +468,7 @@ def move_to_buffer( # 4. Execute transfers and wait for completion. communicator.execute() - return TransferMetadata( - is_unchanged=is_unchanged, - is_received_locally=is_received_locally, - recv_primary_mask=recv_primary_mask, - recv_count=recv_count, - recv_expert_ids=recv_expert_ids, - recv_dst_rows=recv_dst_rows, - ) + return transfer_metadata def move_from_buffer( @@ -432,6 +559,8 @@ def transfer_layer( stream: torch.Stream | None = None, rank_mapping: dict[int, int] | None = None, layer_idx: int = 0, + enable_migration_batching: bool = False, + migration_batches: list[list[MigrationFlow]] | None = None, ) -> TransferMetadata: """Rearranges the expert weights in place according to the new expert indices. @@ -453,6 +582,9 @@ def transfer_layer( stream: CUDA stream for async copies (can be None for sync mode). rank_mapping: Optional rank mapping for elastic expert parallelism. layer_idx: Index of the MoE layer being transferred. + enable_migration_batching: Schedule remote transfers in batches where + each rank communicates with at most one peer. + migration_batches: Optional batches precomputed for this layer. Returns: TransferMetadata: Metadata needed for completing remote weight transfers, @@ -502,6 +634,8 @@ def transfer_layer( ep_rank=ep_group.rank(), communicator=communicator, layer_idx=layer_idx, + enable_migration_batching=enable_migration_batching, + migration_batches=migration_batches, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 6df9d1cfbe..61de1802d5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -482,6 +482,16 @@ class KVConnectorBase_V1(ABC): """ return + def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]: + """KV cache groups restored by this connector's load for ``request``. + + Called after ``get_num_new_matched_tokens`` returned a positive count. + Defaults to the prefix-cacheable groups, all a hash-addressed store + holds; a connector transferring the request's own blocks restores + every transfer group. + """ + return self._kv_cache_config.prefix_cacheable_group_ids + @abstractmethod def get_num_new_matched_tokens( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 46118f4056..172d41efdb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -748,6 +748,9 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): # Scheduler Side Methods ############################################################ + def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]: + return self._kv_cache_config.transfer_group_ids + def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int, bool]: @@ -2404,7 +2407,9 @@ class MooncakeConnectorWorker: continue # No race because we are in async loop. pull_meta.pull_tasks_count -= 1 - if pull_meta.pull_tasks_count == 0: + # Empty pulls only release the producer's blocks; the consumer + # did not enter WAITING_FOR_REMOTE_KVS. + if pull_meta.pull_tasks_count == 0 and any(pull_meta.local_block_ids): self.finished_recving_reqs.add(pull_meta.d_req_id) if ok_reqs: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index ec24f0d199..1dc6f8459a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -59,6 +59,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( MoRIIOWrapper, MoRIIOWriter, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import ( + MoRIIOHeartbeat, +) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import ( LayerTransferGeometry, MambaOffsetTemplate, @@ -1600,7 +1603,7 @@ class MoRIIOConnectorWorker: self.moriio_engine = None self._handle_request_thread = None - self._ping_thread = None + self._heartbeat: MoRIIOHeartbeat | None = None self._writer = MoRIIOWriter(self) # Completions that arrived before transfer_id_to_request_id was populated. # Retried each step until the mapping is established. @@ -1630,10 +1633,7 @@ class MoRIIOConnectorWorker: ) if self._rank == 0 and self.moriio_config.proxy_ip: - self._ping_thread = threading.Thread( - target=self._ping, args=(self.zmq_context,), daemon=True - ) - self._ping_thread.start() + self._heartbeat = self._start_heartbeat() logger.info( "Initializing MoRIIO Engine, engine = %s, role = %s", @@ -1881,75 +1881,31 @@ class MoRIIOConnectorWorker: remote_engine_id ] - def _ping(self, zmq_context): - # Use host:port format for http_address (compatible with official router) - http_address = f"{self.request_address}" - # Include host so the router embeds it in the request_id; the connector - # on the other side parses host/ports from there. - zmq_address = ( - f"host:{self.local_ip}," - f"handshake:{self.handshake_port}," - f"notify:{self.notify_port}" + def _start_heartbeat(self) -> MoRIIOHeartbeat: + payload = { + "type": "P" if self.is_producer else "D", + "http_address": self.request_address, + "zmq_address": ( + f"host:{self.local_ip},handshake:{self.handshake_port}," + f"notify:{self.notify_port}" + ), + "dp_size": self.moriio_config.dp_size, + "tp_size": self.moriio_config.tp_size, + "transfer_mode": self.mode.name, + } + return MoRIIOHeartbeat( + f"tcp://{self.proxy_ip}:{self.proxy_ping_port}", + payload, + MoRIIOConstants.PING_INTERVAL, + MoRIIOConstants.MAX_PING_RETRIES, ) - role = "P" if self.is_producer else "D" - - retry_count = 0 - index = 1 - with zmq_context.socket(zmq.DEALER) as sock: - sock.connect(f"tcp://{self.proxy_ip}:{self.proxy_ping_port}") - - while True: - try: - data = { - "type": role, # "P" or "D" - "http_address": http_address, - "zmq_address": zmq_address, - # dp_size/tp_size are not used by the official vLLM router - # (routing operates at the http_address level); they are - # consumed only by the toy proxy server. - "dp_size": self.moriio_config.dp_size, - "tp_size": self.moriio_config.tp_size, - # transfer_mode is included so the router can distinguish - # READ (prefill-then-decode, sequential) from WRITE (concurrent) - # scheduling. - "transfer_mode": self.mode.name, - } - - sock.send(msgpack.dumps(data)) - # logger.debug(f"Successfully sent ping message #{index}") - retry_count = 0 - - except ConnectionRefusedError: - logger.info( - "Connection refused: %s:%s -> %s:%s", - self.local_ip, - self.local_ping_port, - self.proxy_ip, - self.proxy_ping_port, - ) - retry_count += 1 - - except OSError as e: - logger.info("OS error when sending ping: %s", e) - retry_count += 1 - - except Exception as e: - logger.info("Unexpected error when sending ping: %s", e) - retry_count += 1 - if retry_count >= MoRIIOConstants.MAX_PING_RETRIES: - logger.error( - "Max retries (%s) exceeded. Stopping ping loop.", - MoRIIOConstants.MAX_PING_RETRIES, - ) - raise RuntimeError( - f"Ping failed after {retry_count} retries" - ) from e - - finally: - time.sleep(MoRIIOConstants.PING_INTERVAL) - index += 1 def shutdown(self): + heartbeat = getattr(self, "_heartbeat", None) + if heartbeat is not None: + heartbeat.shutdown() + self._heartbeat = None + if hasattr(self, "moriio_wrapper") and self.moriio_wrapper: self.moriio_wrapper.shutdown() @@ -2580,6 +2536,9 @@ class MoRIIOConnectorWorker: The scheduler process (via the MultiprocExecutor) will use this output to track which workers are done. """ + heartbeat = getattr(self, "_heartbeat", None) + if heartbeat is not None: + heartbeat.check_health() done_sending, done_recving = set(), set() if self.is_producer: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py new file mode 100644 index 0000000000..87eefe8f42 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Keep discovery heartbeats independent of the model worker's GIL.""" + +import json +import logging +import select +import subprocess +import sys +from typing import Any + +import msgpack +import zmq + +logger = logging.getLogger(__name__) + + +class MoRIIOHeartbeat: + def __init__( + self, + endpoint: str, + payload: dict[str, Any], + interval: float, + max_retries: int = 100, + ): + self._process = subprocess.Popen( + [sys.executable, __file__, endpoint, str(interval), str(max_retries)], + stdin=subprocess.PIPE, + text=True, + ) + assert self._process.stdin is not None + try: + self._process.stdin.write(json.dumps(payload) + "\n") + self._process.stdin.flush() + except BaseException: + self.shutdown() + raise + + def check_health(self) -> None: + returncode = self._process.poll() + if returncode is not None: + raise RuntimeError( + f"MoRIIO discovery heartbeat exited with code {returncode}" + ) + + def shutdown(self) -> None: + assert self._process.stdin is not None + try: + if not self._process.stdin.closed: + self._process.stdin.close() + except BrokenPipeError: + pass + finally: + try: + self._process.wait(timeout=5) + except subprocess.TimeoutExpired: + self._process.terminate() + try: + self._process.wait(timeout=1) + except subprocess.TimeoutExpired: + self._process.kill() + self._process.wait(timeout=1) + + +def main() -> None: + endpoint, interval_arg, max_retries_arg = sys.argv[1:] + initial = sys.stdin.readline() + if not initial: + return + payload = msgpack.dumps(json.loads(initial)) + interval = float(interval_arg) + max_retries = int(max_retries_arg) + failures = 0 + with zmq.Context() as context, context.socket(zmq.DEALER) as sock: + sock.setsockopt(zmq.LINGER, 0) + sock.setsockopt(zmq.SNDTIMEO, 1000) + sock.connect(endpoint) + while not select.select([sys.stdin], [], [], 0)[0]: + try: + sock.send(payload) + failures = 0 + except zmq.Again: + failures = 0 + logger.warning("MoRIIO discovery heartbeat send failed; retrying") + except zmq.ZMQError: + failures += 1 + logger.exception("MoRIIO discovery heartbeat send failed") + if failures >= max_retries: + raise + # EOF also stops the child when the worker exits without cleanup. + if select.select([sys.stdin], [], [], interval)[0]: + break + + +if __name__ == "__main__": + main() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 73b6b0ccf6..752acb01ec 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -426,6 +426,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): # ============================== # Scheduler-side methods # ============================== + def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]: + connector = self._connectors[self._requests_to_connector[request.request_id]] + return connector.get_loaded_kv_cache_group_ids(request) + def get_num_new_matched_tokens( self, request: "Request", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index ac929f7fcd..27b4d410f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -164,6 +164,9 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): def get_finished_count(self) -> int | None: return None + def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]: + return self._kv_cache_config.transfer_group_ids + def get_num_new_matched_tokens( self, request: "Request", num_computed_tokens: int ) -> tuple[int | None, bool]: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ce24ac28f2..f4cbcf9012 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -80,6 +80,7 @@ from vllm.config.cache import ( ) from vllm.config.device import Device from vllm.config.kernel import ( + PASSTHROUGH_ALL2ALL_BACKEND, IrOpPriorityConfig, LinearBackend, MoEBackend, @@ -1104,6 +1105,10 @@ class EngineArgs: # Parallel arguments parallel_kwargs = get_kwargs(ParallelConfig) + # Bound from --moe-backend in KernelConfig.set_platform_defaults(). + parallel_kwargs["all2all_backend"]["choices"].remove( + PASSTHROUGH_ALL2ALL_BACKEND + ) parallel_group = parser.add_argument_group( title="ParallelConfig", description=ParallelConfig.__doc__, diff --git a/vllm/entrypoints/launchers/cli_args.py b/vllm/entrypoints/launchers/cli_args.py index 4c1f121372..28fe79fd69 100644 --- a/vllm/entrypoints/launchers/cli_args.py +++ b/vllm/entrypoints/launchers/cli_args.py @@ -207,9 +207,9 @@ class BaseFrontendArgs: enable_scale_out: bool = False """ If set to True, register the scale-out endpoints (`/render`, `/derender`, - and `/inference/v1/generate`) on `vllm serve`. Has no effect on - `vllm launch render` or `vllm serve --tokens-only`, which always register - their required endpoints regardless of this flag. + `/inference/v1/generate` and `/inference/v1/abort_requests`) on `vllm serve`. + Has no effect on `vllm launch render` or `vllm serve --tokens-only` which + always register their required endpoints regardless of this flag. """ fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full" """Controls the ``system_fingerprint`` field on responses. diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index c0c0c15e9c..92cdfaec25 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -21,6 +21,7 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.engine.typing import AnyRequest from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.exceptions import GenerationError, RetryableRequestError from vllm.lora.request import LoRARequest from vllm.renderers.base import BaseRenderer from vllm.tracing import ( @@ -206,6 +207,10 @@ class PoolingBaseServing(ABC, BaseServing): final_res_batch = [None] * num_inputs async for i, res in ctx.result_generator: + if res.error is not None: + if res.error.retryable: + raise RetryableRequestError(res.error.message) + raise GenerationError(res.error.message) final_res_batch[i] = res if None in final_res_batch: diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 8ffa2e4ddd..91c5f762cf 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -191,6 +191,27 @@ class ScoringIOProcessor(PoolingIOProcessor): scoring_data = self.valid_inputs(data_1, data_2) return scoring_data + def _pair_error_output( + self, + output_1: PoolingRequestOutput, + output_2: PoolingRequestOutput, + ) -> PoolingRequestOutput | None: + failed = output_1 if output_1.error is not None else output_2 + if failed.error is None: + return None + + padding = [self.pad_token_id] if self.pad_token_id is not None else [] + return PoolingRequestOutput( + request_id=f"{output_1.request_id}_{output_2.request_id}", + outputs=failed.outputs, + prompt_token_ids=( + output_1.prompt_token_ids + padding + output_2.prompt_token_ids + ), + num_cached_tokens=(output_1.num_cached_tokens + output_2.num_cached_tokens), + finished=failed.finished, + error=failed.error, + ) + class BiEncoderIOProcessor(ScoringIOProcessor): name = "bi-encoder" @@ -310,6 +331,10 @@ class BiEncoderIOProcessor(ScoringIOProcessor): final_res_batch: list[PoolingRequestOutput] = [] for emb_1, emb_2 in zip(emb_data_1, emb_data_2): + if error_output := self._pair_error_output(emb_1, emb_2): + final_res_batch.append(error_output) + continue + pair_score = F.cosine_similarity( emb_1.outputs.data.float(), emb_2.outputs.data.float(), dim=0 ) @@ -352,6 +377,10 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor): # Compute MaxSim scores for emb_1, emb_2 in zip(emb_data_1, emb_data_2): + if error_output := self._pair_error_output(emb_1, emb_2): + final_res_batch.append(error_output) + continue + # emb_1.outputs.data: [query_len, dim] # emb_2.outputs.data: [doc_len, dim] q_emb = emb_1.outputs.data @@ -907,6 +936,10 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM final_res_batch: list[PoolingRequestOutput] = [] for i in range(len(outputs)): + if outputs[i].error is not None: + final_res_batch.append(outputs[i]) + continue + embeds = outputs[i].outputs.data.float() # The JinaForRanking model concatenates docs first, then query. diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index 674522b7eb..b84d3310d6 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -195,13 +195,36 @@ class ServingScores(PoolingServing): ctx = await self._init_ctx(self.io_processor, *args, **kwargs) await self._preprocessing(self.io_processor, ctx) - # stage 1: encode queries and cache token embeddings on workers. - await self._flash_late_interaction_encode_queries(ctx) - # stage 2: encode docs and return scalar scores from workers. - await self._flash_late_interaction_encode_docs(ctx) + try: + # stage 1: encode queries and cache token embeddings on workers. + await self._flash_late_interaction_encode_queries(ctx) + # stage 2: encode docs and return scalar scores from workers. + await self._flash_late_interaction_encode_docs(ctx) + except BaseException: + try: + await self._cleanup_flash_late_interaction(ctx) + except Exception: + logger.exception("Failed to clean up late-interaction query cache.") + raise return await self._postprocessing_async(self.io_processor, ctx) + async def _cleanup_flash_late_interaction(self, ctx: ScoringServeContext) -> None: + query_keys = ctx.late_interaction_query_keys + if not query_keys: + return + + doc_keys = ctx.late_interaction_doc_keys or [] + + try: + # Stop documents before removing the query tensors they reference. + await self.engine_client.abort([*query_keys, *doc_keys]) + finally: + await self.engine_client.collective_rpc( + "release_late_interaction_query_cache", + args=(query_keys,), + ) + async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext): assert ctx.n_queries is not None assert ctx.engine_inputs is not None @@ -259,7 +282,9 @@ class ServingScores(PoolingServing): query_keys = ctx.late_interaction_query_keys if query_keys is None: raise RuntimeError("Late-interaction query keys were not initialized.") - doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)] + doc_namespace = random_uuid() + doc_keys = [f"late-interaction-{doc_namespace}-doc-{i}" for i in range(n_docs)] + ctx.late_interaction_doc_keys = doc_keys for i in range(n_docs): query_idx = 0 if n_queries == 1 else i diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 7ce26dd6cf..8e45ce33f4 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -115,6 +115,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): ## for flash-late-interaction query_final_res_batch: list[PoolingRequestOutput] | None = None late_interaction_query_keys: list[str] | None = None + late_interaction_doc_keys: list[str] | None = None @dataclass diff --git a/vllm/entrypoints/scale_out/derender/serving.py b/vllm/entrypoints/scale_out/derender/serving.py index 0783d3a78f..72ef434770 100644 --- a/vllm/entrypoints/scale_out/derender/serving.py +++ b/vllm/entrypoints/scale_out/derender/serving.py @@ -29,7 +29,7 @@ from ..token_in_token_out.protocol import ( DerenderCompletionRequest, DerenderCompletionStreamRequest, DerenderStreamState, - GenerateResponse, + GenerateTokensResponse, MultiModalFeatures, ) @@ -54,7 +54,7 @@ class ServingDerender(BaseServing): def _validate_derender_bounds( self, - generate_responses: list[GenerateResponse], + generate_responses: list[GenerateTokensResponse], ) -> ErrorResponse | None: """Reject derender payloads that exceed resource bounds. diff --git a/vllm/entrypoints/scale_out/token_in_token_out/api_router.py b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py index 28677b6321..ee0e8c1623 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/api_router.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/api_router.py @@ -21,7 +21,7 @@ from vllm.logger import init_logger from ...serve.engine.protocol import ErrorResponse from .protocol import ( GenerateRequest, - GenerateResponse, + GenerateResponseBase, ) from .serving import ServingTokens @@ -67,35 +67,51 @@ async def generate(request: GenerateRequest, raw_request: Request): content=generator.model_dump(), status_code=generator.error.code ) - elif isinstance(generator, GenerateResponse): + elif isinstance(generator, GenerateResponseBase): return JSONResponse(content=generator.model_dump()) return StreamingResponse(content=generator, media_type="text/event-stream") +abort_router = APIRouter() + + +async def abort_requests(raw_request: Request): + """Abort one or more requests. To be used in a + Disaggregated Everything setup. + """ + try: + body = await raw_request.json() + except json.JSONDecodeError as e: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"JSON decode error: {e}", + ) from e + request_ids = body.get("request_ids") + if request_ids is None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Missing 'request_ids' in request body", + ) + # Abort requests in background + asyncio.create_task(engine_client(raw_request).abort(request_ids)) + return Response(status_code=200) + + +# Under the `/inference` prefix, so `--api-key` guards it. +router.add_api_route("/inference/v1/abort_requests", abort_requests, methods=["POST"]) +# Unauthenticated legacy path, kept for existing `--tokens-only` deployments. +abort_router.add_api_route("/abort_requests", abort_requests, methods=["POST"]) + + def attach_router(app: FastAPI): - if getattr(app.state.args, "tokens_only", False): - - @router.post("/abort_requests") - async def abort_requests(raw_request: Request): - """Abort one or more requests. To be used in a - Disaggregated Everything setup. - """ - try: - body = await raw_request.json() - except json.JSONDecodeError as e: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"JSON decode error: {e}", - ) from e - request_ids = body.get("request_ids") - if request_ids is None: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="Missing 'request_ids' in request body", - ) - # Abort requests in background - asyncio.create_task(engine_client(raw_request).abort(request_ids)) - return Response(status_code=200) - app.include_router(router) + if not getattr(app.state.args, "tokens_only", False): + return + # The RLHF dev router registers its own /abort_requests first. Registering + # this one too would only add a shadowed route with a duplicate operation ID. + has_abort_route = any( + getattr(route, "path", None) == "/abort_requests" for route in app.routes + ) + if not has_abort_route: + app.include_router(abort_router) 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 60cffd07c2..31bd7e2eae 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -4,9 +4,11 @@ from typing import Annotated, Any, Literal, TypeAlias from pydantic import ( BaseModel, + Discriminator, Field, NonNegativeInt, PrivateAttr, + Tag, field_validator, model_validator, ) @@ -34,6 +36,8 @@ from vllm.utils import random_uuid ####### Tokens IN <> Tokens OUT ####### +OutputMode: TypeAlias = Literal["tokens", "text"] + class PlaceholderRangeInfo(BaseModel): """Serializable placeholder location for a single multi-modal item.""" @@ -44,7 +48,7 @@ class PlaceholderRangeInfo(BaseModel): length: int = Field(gt=0) """Number of placeholder tokens.""" - # TODO: add ``is_embed: list[bool] | None`` once the /generate side + # TODO: add `is_embed: list[bool] | None` once the /generate side # consumes features — some models (e.g. Qwen-VL) use sparse # placeholder masks that cannot be recomputed from offset+length alone. @@ -74,22 +78,22 @@ class MultiModalFeatures(BaseModel): kwargs_data: dict[str, list[str | None]] | None = None """Per-modality serialized tensor data. - Each value is a list parallel to ``mm_hashes[modality]``. A ``str`` - entry is a base64-encoded ``MultiModalKwargsItem``; ``None`` means + Each value is a list parallel to `mm_hashes[modality]`. A `str` + entry is a base64-encoded `MultiModalKwargsItem`; `None` means the item should be resolved from cache. The entire field is - ``None`` for metadata-only (cache-hit) responses. + `None` for metadata-only (cache-hit) responses. """ mm_metadata: dict[str, list[str | None]] | None = None """Per-modality serialized metadata for disaggregated prefill. - Each value is a list parallel to ``mm_hashes[modality]``. A ``str`` - entry is a base64-encoded ``MultiModalKwargsItem`` containing only - placeholder-metadata and ``keep_on_cpu`` fields. ``None`` means that + Each value is a list parallel to `mm_hashes[modality]`. A `str` + entry is a base64-encoded `MultiModalKwargsItem` containing only + placeholder-metadata and `keep_on_cpu` fields. `None` means that the metadata is unavailable for that item. Prefill can use this - instead of ``kwargs_data`` only when ``ec_transfer_params`` is also + instead of `kwargs_data` only when `ec_transfer_params` is also set, so embeddings arrive through the EC connector rather than from - ``pixel_values``. + `pixel_values`. """ @model_validator(mode="after") @@ -174,7 +178,7 @@ class GenerateRequest(BaseModel): content_parts: list[dict[str, Any]] | None = None """Raw multimodal input; server resolves media. Mutually exclusive - with ``features``.""" + with `features`.""" @model_validator(mode="after") def _check_mm_fields_exclusive(self) -> "GenerateRequest": @@ -217,6 +221,16 @@ class GenerateRequest(BaseModel): ), ) + output_mode: OutputMode = Field( + default="tokens", + description=( + "Response level. 'tokens' returns token IDs only. 'text' also " + "returns the detokenized text and logprobs with decoded token " + "strings. 'text' needs a server that loads a tokenizer and " + "sampling_params.detokenize to be true." + ), + ) + stream: bool | None = False stream_options: StreamOptions | None = None cache_salt: str | None = Field( @@ -253,12 +267,12 @@ class GenerateRequest(BaseModel): ), ) - # Tracks which keys the caller explicitly set inside ``sampling_params`` + # Tracks which keys the caller explicitly set inside `sampling_params` # when the request was parsed from a JSON body. Lets the server tell # "client said max_tokens=16" from "client said nothing → dataclass # default 16" so it can apply server-side defaulting only in the latter - # case. ``None`` means the request was constructed with a pre-built - # ``SamplingParams`` instance (e.g. from internal callers that have + # case. `None` means the request was constructed with a pre-built + # `SamplingParams` instance (e.g. from internal callers that have # already resolved values), in which case all fields are considered set. _sampling_params_provided_keys: set[str] | None = PrivateAttr(default=None) _response_mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = ( @@ -284,6 +298,15 @@ class GenerateRequest(BaseModel): validate_cache_salt(data.get("cache_salt")) return data + @model_validator(mode="after") + def _check_output_mode_detokenizes(self) -> "GenerateRequest": + if self.output_mode != "tokens" and not self.sampling_params.detokenize: + raise ValueError( + f"output_mode={self.output_mode!r} requires " + "sampling_params.detokenize to be true" + ) + return self + @model_validator(mode="after") def _validate_multimodal_feature_bounds(self) -> "GenerateRequest": if self.features is None: @@ -299,10 +322,10 @@ class GenerateRequest(BaseModel): return self def is_sampling_param_provided(self, name: str) -> bool: - """Whether the caller explicitly set ``sampling_params.``. + """Whether the caller explicitly set `sampling_params.`. For requests parsed from a JSON body, this reflects the raw input - dict. For requests constructed with a pre-built ``SamplingParams`` + dict. For requests constructed with a pre-built `SamplingParams` instance, all fields are considered provided so server-side defaults do not clobber values already resolved upstream. """ @@ -317,21 +340,22 @@ class GenerateRequest(BaseModel): ) -class GenerateResponseChoice(BaseModel): +class GenerateChoiceBase(BaseModel): + """Fields shared by every `output_mode` of a non-streaming choice.""" + index: int - logprobs: ChatCompletionLogProbs | None = None # per OpenAI spec this is the default finish_reason: str | None = "stop" token_ids: list[int] | None = None - # Per-token expert routing decisions, base64-encoded ``.npy`` bytes + # 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/int32 - # ``num_tokens - 1`` because the last sampled token has not been + # `num_tokens - 1` because the last sampled token has not been # forwarded yet and therefore has no routing data. # Decode: # np.load(io.BytesIO(base64.b64decode(s))) - # ``None`` if (a) the request was aborted before any forward pass, - # or (b) ``enable_return_routed_experts`` is off server-side. + # `None` if (a) the request was aborted before any forward pass, + # or (b) `enable_return_routed_experts` is off server-side. routed_experts: str | None = None sampling_mask: list[list[int]] | None = None @@ -343,16 +367,43 @@ class GenerateResponseChoice(BaseModel): return v -class GenerateResponseStreamChoice(BaseModel): - index: int +class GenerateTokensChoice(GenerateChoiceBase): logprobs: ChatCompletionLogProbs | None = None + """Logprobs whose tokens are `token_id:N` placeholders.""" + + +class GenerateTextChoice(GenerateChoiceBase): + text: str + """Detokenized output. Excludes a matched stop string unless + `include_stop_str_in_output` is set, while `token_ids` keeps every + generated token.""" + + logprobs: ChatCompletionLogProbs | None = None + """Logprobs with decoded token strings and `bytes`.""" + + +class GenerateStreamChoiceBase(BaseModel): + """Fields shared by every `output_mode` of a streaming choice.""" + + index: int finish_reason: str | None = None token_ids: list[int] | None = None routed_experts: str | None = None sampling_mask: list[list[int]] | None = None -class GenerateStreamResponse(BaseModel): +class GenerateTokensStreamChoice(GenerateStreamChoiceBase): + logprobs: ChatCompletionLogProbs | None = None + + +class GenerateTextStreamChoice(GenerateStreamChoiceBase): + text: str + """Text delta since the previous chunk of this choice.""" + + logprobs: ChatCompletionLogProbs | None = None + + +class GenerateStreamResponseBase(BaseModel): request_id: str = Field( default_factory=lambda: f"{random_uuid()}", description=( @@ -361,14 +412,23 @@ class GenerateStreamResponse(BaseModel): "through out the inference process and return in response." ), ) - choices: list[GenerateResponseStreamChoice] usage: UsageInfo | None = Field(default=None) prompt_token_ids: list[NonNegativeInt] | None = None mm_placeholders: dict[str, list[PlaceholderRangeInfo]] | None = None metrics: PerRequestMetrics | None = None -class GenerateResponse(BaseModel): +class GenerateTokensStreamResponse(GenerateStreamResponseBase): + output_mode: Literal["tokens"] = "tokens" + choices: list[GenerateTokensStreamChoice] + + +class GenerateTextStreamResponse(GenerateStreamResponseBase): + output_mode: Literal["text"] = "text" + choices: list[GenerateTextStreamChoice] + + +class GenerateResponseBase(BaseModel): request_id: str = Field( default_factory=lambda: f"{random_uuid()}", description=( @@ -379,7 +439,6 @@ class GenerateResponse(BaseModel): ) model: str | None = None created: int | None = None - choices: list[GenerateResponseChoice] usage: UsageInfo | None = Field(default=None) prompt_logprobs: list[dict[int, Logprob] | None] | None = None prompt_token_id_logprobs: str | None = None @@ -399,6 +458,42 @@ class GenerateResponse(BaseModel): ) +class GenerateTokensResponse(GenerateResponseBase): + output_mode: Literal["tokens"] = "tokens" + choices: list[GenerateTokensChoice] + + +class GenerateTextResponse(GenerateResponseBase): + output_mode: Literal["text"] = "text" + choices: list[GenerateTextChoice] + + +def output_mode_or_tokens(value: Any) -> Any: + """Discriminator for `GenerateResponse` and `GenerateStreamResponse`. + + A missing `output_mode` means `tokens`. Existing clients and servers + that predate the field never send it. + """ + if isinstance(value, dict): + return value.get("output_mode", "tokens") + return getattr(value, "output_mode", "tokens") + + +# Type aliases, not classes: use `GenerateResponseBase` or +# `GenerateStreamResponseBase` for `isinstance` checks and the +# `Generate...` classes to construct. +GenerateResponse: TypeAlias = Annotated[ + Annotated[GenerateTokensResponse, Tag("tokens")] + | Annotated[GenerateTextResponse, Tag("text")], + Discriminator(output_mode_or_tokens), +] +GenerateStreamResponse: TypeAlias = Annotated[ + Annotated[GenerateTokensStreamResponse, Tag("tokens")] + | Annotated[GenerateTextStreamResponse, Tag("text")], + Discriminator(output_mode_or_tokens), +] + + class DerenderChatRequest(BaseModel): """Request for the /v1/chat/completions/derender endpoint (non-streaming). @@ -412,8 +507,13 @@ class DerenderChatRequest(BaseModel): model: str | None = None """Served model name. Defaults to the server's served model name.""" - generate_response: GenerateResponse - """The complete token-in / token-out engine response to derender.""" + generate_response: GenerateTokensResponse + """The complete token-in / token-out engine response to derender. + + Only `output_mode="tokens"` responses are accepted. Other modes are + already detokenized and re-decoding their `token_ids` would undo the + engine's stop string truncation. + """ prompt_tokens: int | None = None """Prompt token count for usage; defaults to 0 if omitted. @@ -455,9 +555,9 @@ class DerenderCompletionRequest(BaseModel): model: str | None = None """Served model name. Defaults to the server's served model name.""" - generate_responses: list[GenerateResponse] - """One response per prompt, parallel to the list[GenerateRequest] - returned by /v1/completions/render.""" + generate_responses: list[GenerateTokensResponse] + """One `output_mode="tokens"` response per prompt, parallel to the + list[GenerateRequest] returned by /v1/completions/render.""" prompt_tokens: list[int] | None = None """One prompt token count per response; each defaults to 0 if omitted. @@ -589,7 +689,7 @@ class DerenderStreamState(BaseModel): token in here through `parse_delta` (discarding the result) before processing the current chunk's tokens since parser internal state cannot be serialized into this stateless model. Unavoidably O(n) - bounded by ``max_model_len`` (enforced server side, not by a field + bounded by `max_model_len` (enforced server side, not by a field validator here since the bound is model dependent). """ @@ -630,24 +730,25 @@ class DerenderChatStreamRequest(BaseModel): """One chunk streaming derender request for /v1/chat/completions/derender. The client sends one request per SSE chunk received from - ``/inference/v1/generate``. Each request carries the generate chunk - plus the ``stream_state`` returned by the previous call (``None`` on the + `/inference/v1/generate`. Each request carries the generate chunk + plus the `stream_state` returned by the previous call (`None` on the first call). The response contains the derendered chunk and the updated state to be passed to the next call. This implements stateless no server side session. All mutable state lives in - the client carried ``stream_state``. + the client carried `stream_state`. """ # --8<-- [start:derender-chat-stream-request] stream: Literal[True] model: str | None = None - generate_chunk: GenerateStreamResponse - """One SSE chunk from ``/inference/v1/generate`` (``stream=True``).""" + generate_chunk: GenerateTokensStreamResponse + """One `output_mode="tokens"` SSE chunk from `/inference/v1/generate` + (`stream=True`).""" stream_state: DerenderStreamState | None = None - """Client carried detok state from the previous call. ``None`` on first.""" + """Client carried detok state from the previous call. `None` on first.""" prompt_tokens: int | None = None """Prompt token count for usage. Forwarded from the render step.""" @@ -655,7 +756,7 @@ class DerenderChatStreamRequest(BaseModel): prompt_token_ids: list[NonNegativeInt] | None = None """Prompt token IDs. Required by the parser path's `parse_delta` to settle its initial reasoning state (e.g. chat templates that pre-open - ````). `prompt_tokens` is a usage count and cannot serve this + ``). `prompt_tokens` is a usage count and cannot serve this purpose. Sourced from `GenerateRequest.token_ids` at the render step. Rejected with a 400 (by `ServingDerender`) when a tool or reasoning @@ -680,7 +781,7 @@ class DerenderChatStreamRequest(BaseModel): class DerenderCompletionStreamRequest(BaseModel): """One chunk streaming derender request for /v1/completions/derender. - Parallel to ``DerenderChatStreamRequest`` for the completions endpoint. + Parallel to `DerenderChatStreamRequest` for the completions endpoint. Each call processes one SSE chunk (one output sequence's delta) and returns the derendered chunk plus updated state. """ @@ -689,11 +790,11 @@ class DerenderCompletionStreamRequest(BaseModel): stream: Literal[True] model: str | None = None - generate_chunk: GenerateStreamResponse - """One SSE chunk from ``/inference/v1/generate``.""" + generate_chunk: GenerateTokensStreamResponse + """One `output_mode="tokens"` SSE chunk from `/inference/v1/generate`.""" stream_state: DerenderStreamState | None = None - """Client-carried detok state. ``None`` on the first call.""" + """Client-carried detok state. `None` on the first call.""" prompt_tokens: int | None = None """Prompt token count for usage.""" @@ -723,16 +824,16 @@ class DerenderChatStreamResponse(BaseModel): class DerenderCompletionStreamResponse(BaseModel): """Response for one streaming completions derender chunk. - Parallel to ``DerenderChatStreamResponse`` for the completions endpoint. + Parallel to `DerenderChatStreamResponse` for the completions endpoint. """ chunk: CompletionStreamResponse stream_state: DerenderStreamState -# Determines the type by checking the ``stream`` field's literal value. A body without -# ``stream`` validates as the non-streaming member -# (``stream`` defaults to ``False`` there), so FastAPI can validate and dispatch both +# Determines the type by checking the `stream` field's literal value. A body without +# `stream` validates as the non-streaming member +# (`stream` defaults to `False` there), so FastAPI can validate and dispatch both # shapes on a single path. DerenderChatRequestUnion: TypeAlias = DerenderChatRequest | DerenderChatStreamRequest DerenderCompletionRequestUnion: TypeAlias = ( 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 249ed48093..a30e519cb2 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -6,6 +6,7 @@ import asyncio import time from collections.abc import AsyncGenerator from collections.abc import Sequence as GenericSequence +from typing import Any import msgspec from fastapi import Request @@ -20,6 +21,7 @@ from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, build_spec_decoding_metrics, clamp_prompt_logprobs, + format_token_id_placeholder, ) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProb, @@ -46,6 +48,7 @@ from vllm.multimodal.inputs import ( from vllm.outputs import RequestOutput from vllm.renderers.online_renderer import OnlineRenderer from vllm.sampling_params import RequestOutputKind, SamplingParams +from vllm.tokenizers import TokenizerLike from vllm.utils.collection_utils import as_list from vllm.utils.serial_utils import numpy2base64 @@ -56,14 +59,45 @@ from .mm_features import ( from .protocol import ( GenerateRequest, GenerateResponse, - GenerateResponseChoice, - GenerateResponseStreamChoice, - GenerateStreamResponse, + GenerateTextChoice, + GenerateTextResponse, + GenerateTextStreamChoice, + GenerateTextStreamResponse, + GenerateTokensChoice, + GenerateTokensResponse, + GenerateTokensStreamChoice, + GenerateTokensStreamResponse, ) logger = init_logger(__name__) +def _logprob_token( + token_id: int, logprob: Logprob | None, tokenizer: TokenizerLike | None +) -> tuple[str, list[int] | None]: + """Token string and UTF-8 bytes for one logprob entry. + + Without a tokenizer the token is a ``token_id:N`` placeholder with no bytes. + """ + if tokenizer is None: + return format_token_id_placeholder(token_id), None + token = ( + tokenizer.decode([token_id]) + if logprob is None + else GenerateBaseServing._get_decoded_token(logprob, token_id, tokenizer) + ) + return token, list(token.encode("utf-8", errors="replace")) + + +def _top_logprob( + token_id: int, logprob: Logprob, tokenizer: TokenizerLike | None +) -> ChatCompletionLogProb: + token, token_bytes = _logprob_token(token_id, logprob, tokenizer) + return ChatCompletionLogProb( + token=token, logprob=max(logprob.logprob, -9999.0), bytes=token_bytes + ) + + class ServingTokens(GenerateBaseServing): """Provides Tokens IN <> Tokens OUT functionality to vLLM API.""" @@ -89,6 +123,9 @@ class ServingTokens(GenerateBaseServing): self.enable_prompt_tokens_details = enable_prompt_tokens_details self.enable_log_outputs = enable_log_outputs self.force_no_detokenize = force_no_detokenize + self.has_tokenizer = ( + not force_no_detokenize and self.renderer.tokenizer is not None + ) if force_no_detokenize: logger.info( "Tokens-only mode is enabled, skipping detokenization " @@ -171,6 +208,15 @@ class ServingTokens(GenerateBaseServing): "because detokenization is disabled. Check stop strings on " "the coordinator, or use stop_token_ids." ) + if request.output_mode != "tokens" and not self.has_tokenizer: + # The no-op detokenizer returns "", so without this check the + # request would succeed with empty text. + return self.create_error_response( + f"output_mode={request.output_mode!r} requires a tokenizer, but " + "this server does not load one (--tokens-only or " + "--skip-tokenizer-init). Send the request to a server that " + "loads a tokenizer, or use output_mode='tokens'." + ) try: msgspec.msgpack.encode( ( @@ -331,6 +377,8 @@ class ServingTokens(GenerateBaseServing): created_time = int(time.time()) final_res: RequestOutput | None = None sampling_params: SamplingParams = request.sampling_params + text_mode = request.output_mode == "text" + tokenizer = self._logprobs_tokenizer(text_mode) try: async for res in result_generator: @@ -340,7 +388,8 @@ class ServingTokens(GenerateBaseServing): assert final_res is not None - choices: list[GenerateResponseChoice] = [] + tokens_choices: list[GenerateTokensChoice] = [] + text_choices: list[GenerateTextChoice] = [] num_generated_tokens = 0 for output in final_res.outputs: self._raise_if_error(output.finish_reason, request_id) @@ -355,6 +404,7 @@ class ServingTokens(GenerateBaseServing): token_ids=token_ids, top_logprobs=out_logprobs, num_output_top_logprobs=sampling_params.logprobs, + tokenizer=tokenizer, ) else: logprobs = None @@ -369,7 +419,7 @@ class ServingTokens(GenerateBaseServing): if output.sampling_mask is not None: sampling_mask = output.sampling_mask.token_ids - choice_data = GenerateResponseChoice( + choice_fields: dict[str, Any] = dict( index=output.index, logprobs=logprobs, finish_reason=output.finish_reason if output.finish_reason else "stop", @@ -377,8 +427,12 @@ class ServingTokens(GenerateBaseServing): routed_experts=routed_experts_b64, sampling_mask=sampling_mask, ) - - choices.append(choice_data) + if text_mode: + text_choices.append( + GenerateTextChoice(text=output.text, **choice_fields) + ) + else: + tokens_choices.append(GenerateTokensChoice(**choice_fields)) num_generated_tokens += len(output.token_ids) assert final_res.prompt_token_ids is not None @@ -407,11 +461,10 @@ class ServingTokens(GenerateBaseServing): spec_stats = build_spec_decoding_metrics(final_res) if spec_stats is not None: per_request_metrics = PerRequestMetrics(speculative_decoding=spec_stats) - response = GenerateResponse( + response_fields: dict[str, Any] = dict( request_id=request_id, created=created_time, model=model_name, - choices=choices, usage=usage, prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs), prompt_token_id_logprobs=( @@ -427,10 +480,15 @@ class ServingTokens(GenerateBaseServing): kv_transfer_params=final_res.kv_transfer_params, ec_transfer_params=final_res.ec_transfer_params, ) + response: GenerateTokensResponse | GenerateTextResponse + if text_mode: + response = GenerateTextResponse(choices=text_choices, **response_fields) + else: + response = GenerateTokensResponse(choices=tokens_choices, **response_fields) # Log complete response if output logging is enabled if self.enable_log_outputs and self.request_logger: - for choice in choices: + for choice in response.choices: # Get the corresponding output token IDs output_token_ids = None if choice.index < len(final_res.outputs): @@ -464,6 +522,8 @@ class ServingTokens(GenerateBaseServing): num_cached_tokens = None sampling_params: SamplingParams = request.sampling_params last_res: RequestOutput | None = None + text_mode = request.output_mode == "text" + tokenizer = self._logprobs_tokenizer(text_mode) include_usage, include_continuous_usage = should_include_usage( request.stream_options, False @@ -491,11 +551,21 @@ class ServingTokens(GenerateBaseServing): finish_reason = output.finish_reason self._raise_if_error(finish_reason, request_id) - # Still emit a terminal empty chunk while prompt metadata - # is pending, so zero-token completions deliver it. - if not delta_token_ids and ( - finish_reason is None or prompt_token_ids is None - ): + if text_mode: + # Text held back for stop string matching is flushed + # with the final output, and the abort output has no + # new token IDs, so text or a finish reason is enough. + emit = bool( + delta_token_ids or output.text or finish_reason is not None + ) + else: + # Still emit a terminal empty chunk while prompt + # metadata is pending, so zero-token completions + # deliver it. + emit = bool(delta_token_ids) or ( + finish_reason is not None and prompt_token_ids is not None + ) + if not emit: continue if sampling_params.logprobs is not None: @@ -505,6 +575,7 @@ class ServingTokens(GenerateBaseServing): token_ids=delta_token_ids, top_logprobs=out_logprobs, num_output_top_logprobs=sampling_params.logprobs, + tokenizer=tokenizer, ) else: logprobs = None @@ -518,20 +589,30 @@ class ServingTokens(GenerateBaseServing): sampling_mask = None if output.sampling_mask is not None: sampling_mask = output.sampling_mask.token_ids - - chunk = GenerateStreamResponse( - request_id=request_id, - choices=[ - GenerateResponseStreamChoice( - index=i, - logprobs=logprobs, - finish_reason=finish_reason, - token_ids=as_list(delta_token_ids), - routed_experts=routed_experts_b64, - sampling_mask=sampling_mask, - ) - ], + choice_fields: dict[str, Any] = dict( + index=i, + logprobs=logprobs, + finish_reason=finish_reason, + token_ids=as_list(delta_token_ids), + routed_experts=routed_experts_b64, + sampling_mask=sampling_mask, ) + chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse + if text_mode: + chunk = GenerateTextStreamResponse( + request_id=request_id, + choices=[ + GenerateTextStreamChoice( + text=output.text, **choice_fields + ) + ], + ) + else: + chunk = GenerateTokensStreamResponse( + request_id=request_id, + choices=[GenerateTokensStreamChoice(**choice_fields)], + ) + if prompt_token_ids is not None: chunk.prompt_token_ids = prompt_token_ids chunk.mm_placeholders = request._response_mm_placeholders @@ -571,12 +652,21 @@ class ServingTokens(GenerateBaseServing): per_request_metrics = PerRequestMetrics( speculative_decoding=spec_stats ) - final_chunk = GenerateStreamResponse( - request_id=request_id, - choices=[], - usage=final_usage_info, - metrics=per_request_metrics, - ) + final_chunk: GenerateTokensStreamResponse | GenerateTextStreamResponse + if text_mode: + final_chunk = GenerateTextStreamResponse( + request_id=request_id, + choices=[], + usage=final_usage_info, + metrics=per_request_metrics, + ) + else: + final_chunk = GenerateTokensStreamResponse( + request_id=request_id, + choices=[], + usage=final_usage_info, + metrics=per_request_metrics, + ) yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" request_metadata.final_usage_info = final_usage_info @@ -591,37 +681,49 @@ class ServingTokens(GenerateBaseServing): yield f"data: {data}\n\n" yield "data: [DONE]\n\n" + def _logprobs_tokenizer(self, text_mode: bool) -> TokenizerLike | None: + """Tokenizer that resolves logprob tokens, or None for placeholders. + + ``--return-tokens-as-token-ids`` keeps placeholders at every level, as + it does on ``/v1/completions``. + """ + if not text_mode or self.return_tokens_as_token_ids: + return None + return self.renderer.tokenizer + def _create_tokens_logprobs( self, token_ids: GenericSequence[int], top_logprobs: GenericSequence[dict[int, Logprob] | None], num_output_top_logprobs: int | None = None, + tokenizer: TokenizerLike | None = None, ) -> ChatCompletionLogProbs: - """Create OpenAI-style logprobs.""" + """Create OpenAI-style logprobs. + + Tokens are ``token_id:N`` placeholders. With a ``tokenizer`` they are + decoded strings and carry ``bytes``. + """ logprobs_content: list[ChatCompletionLogProbsContent] = [] for i, token_id in enumerate(token_ids): - token = f"token_id:{token_id}" step_top_logprobs = top_logprobs[i] if step_top_logprobs is None or step_top_logprobs.get(token_id) is None: + token, token_bytes = _logprob_token(token_id, None, tokenizer) logprobs_content.append( - ChatCompletionLogProbsContent( - token=token, - ) + ChatCompletionLogProbsContent(token=token, bytes=token_bytes) ) else: step_token = step_top_logprobs[token_id] + token, token_bytes = _logprob_token(token_id, step_token, tokenizer) logprobs_content.append( ChatCompletionLogProbsContent( token=token, logprob=max(step_token.logprob, -9999.0), + bytes=token_bytes, top_logprobs=[ - ChatCompletionLogProb( - token=f"token_id:{token_id}", - logprob=max(logprob.logprob, -9999.0), - ) - for i, (token_id, logprob) in enumerate( + _top_logprob(top_id, logprob, tokenizer) + for i, (top_id, logprob) in enumerate( step_top_logprobs.items() ) if num_output_top_logprobs is not None diff --git a/vllm/entrypoints/serve/exception_handling/error_response.py b/vllm/entrypoints/serve/exception_handling/error_response.py index 505c5ac35c..46e2504663 100644 --- a/vllm/entrypoints/serve/exception_handling/error_response.py +++ b/vllm/entrypoints/serve/exception_handling/error_response.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from http import HTTPStatus -from vllm.exceptions import GenerationError +from vllm.exceptions import GenerationError, RetryableRequestError from vllm.logger import init_logger from ..engine.protocol import ErrorInfo, ErrorResponse @@ -56,7 +56,11 @@ def create_error_response( status_code = HTTPStatus.BAD_REQUEST param = None elif isinstance(exc, GenerationError): - err_type = "InternalServerError" + err_type = ( + "ServiceUnavailableError" + if isinstance(exc, RetryableRequestError) + else "InternalServerError" + ) status_code = exc.status_code param = None elif isinstance(exc, VLLMServerError): diff --git a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py index c023413937..7b6954b8f7 100644 --- a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py +++ b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py @@ -70,8 +70,8 @@ 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. + that should be returned to the client 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/envs.py b/vllm/envs.py index e63ed5cc02..f51100fe8b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -23,7 +23,6 @@ if TYPE_CHECKING: VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE: int = 256 LOCAL_RANK: int = 0 CUDA_VISIBLE_DEVICES: str | None = None - VLLM_ENGINE_ITERATION_TIMEOUT_S: int = 60 VLLM_ENGINE_READY_TIMEOUT_S: int = 600 VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: float = 30.0 VLLM_API_KEY: str | None = None @@ -166,9 +165,6 @@ if TYPE_CHECKING: VLLM_DISABLE_COMPILE_CACHE: bool = False VLLM_REPLICATE_EMBED: bool = False VLLM_USE_LAYERNAME: bool = True - Q_SCALE_CONSTANT: int = 200 - K_SCALE_CONSTANT: int = 200 - V_SCALE_CONSTANT: int = 100 VLLM_USE_RUST_FRONTEND: bool = False VLLM_USE_RUST_BENCH: bool = False VLLM_RUST_FRONTEND_PATH: str | None = "auto" @@ -226,6 +222,7 @@ if TYPE_CHECKING: VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS: list[str] | None = None VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto" VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024 + VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE: bool = False VLLM_XGRAMMAR_CACHE_MB: int = 0 VLLM_REGEX_COMPILATION_TIMEOUT_S: int = 5 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256 @@ -804,10 +801,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), # used to control the visible devices in the distributed setting "CUDA_VISIBLE_DEVICES": lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None), - # timeout for each iteration in the engine - "VLLM_ENGINE_ITERATION_TIMEOUT_S": lambda: int( - os.environ.get("VLLM_ENGINE_ITERATION_TIMEOUT_S", "60") - ), # Timeout in seconds for waiting for engine cores to become ready # during startup. Default is 600 seconds (10 minutes). "VLLM_ENGINE_READY_TIMEOUT_S": lambda: int( @@ -1805,6 +1798,12 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int( os.getenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024)) ), + # Transmit MoE all-to-all combine (expert-output) payloads in FP8 instead + # of BF16, halving NVLink traffic on the combine leg. Only takes effect + # when the installed FlashInfer MoeAlltoAll kernel supports it. + "VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE": lambda: bool( + int(os.getenv("VLLM_FLASHINFER_MOE_A2A_LOW_PRECISION_COMBINE", "0")) + ), # Control the maximum number of tokens per expert supported by the # NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for # the blockscale tensor of activations NVFP4 Quantization. @@ -2368,7 +2367,6 @@ def compile_factors() -> dict[str, object]: "VLLM_TUNED_CONFIG_FOLDER", "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS", - "VLLM_ENGINE_ITERATION_TIMEOUT_S", "VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", "VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS", diff --git a/vllm/exceptions.py b/vllm/exceptions.py index c7900a0347..9bfd4634bc 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -130,6 +130,14 @@ class GenerationError(VLLMServerError): self.status_code = HTTPStatus.INTERNAL_SERVER_ERROR +class RetryableRequestError(GenerationError): + """A transient request-level failure that is safe for clients to retry.""" + + def __init__(self, message: str = "Request failed; please retry"): + super().__init__(message) + self.status_code = HTTPStatus.SERVICE_UNAVAILABLE + + class GracefulHTTPError(VLLMError): """Exception that should be translated into an HTTP error response. diff --git a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py index 703f766889..18bef4406b 100644 --- a/vllm/lora/ops/triton_ops/fused_moe_lora_op.py +++ b/vllm/lora/ops/triton_ops/fused_moe_lora_op.py @@ -106,14 +106,18 @@ def _get_c_ptrs( _LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {} +# The one-shot kernel keeps the rank-dimension intermediate in registers. +# Larger ranks require rank tiling, so route them through the two-kernel path. +_FUSED_MOE_LORA_ONE_SHOT_MAX_RANK = 128 + # --------------------------------------------------------------------------- # Fully-fused MoE-LoRA kernel (one-shot): shrink + expand combined into a single # launch with the rank-dim intermediate kept in registers. Used by the fast -# path of `_fused_moe_lora` for `fully_sharded=False`. The legacy two-kernel -# path (`_fused_moe_lora_kernel` above) is retained for `fully_sharded=True` -# because that path needs to materialise the intermediate cache for an -# all_reduce / all_gather between shrink and expand. +# path of `_fused_moe_lora` for supported ranks with `fully_sharded=False`. +# The legacy two-kernel path (`_fused_moe_lora_kernel` above) is retained for +# larger ranks and for `fully_sharded=True`, where the intermediate cache must +# be materialised for an all_reduce / all_gather between shrink and expand. # --------------------------------------------------------------------------- @@ -349,8 +353,9 @@ def _run_fused_moe_lora_one_shot( # regressed across all M (+8 to +40%): the (64,32) fp32 accumulator + # widened B tile pushed register count past spill threshold, lowering # occupancy by more than the MMA gain saved. - assert rank <= 128, ( - f"fused_moe_lora_one_shot supports max_lora_rank<=128; got rank={rank}" + assert rank <= _FUSED_MOE_LORA_ONE_SHOT_MAX_RANK, ( + "fused_moe_lora_one_shot supports " + f"max_lora_rank<={_FUSED_MOE_LORA_ONE_SHOT_MAX_RANK}; got rank={rank}" ) BLOCK_R = max(triton.next_power_of_2(rank), 16) @@ -1318,6 +1323,7 @@ def _fused_moe_lora_expand( offset: int = 0, use_gdc: bool = False, use_tma: bool = False, + add_inputs: bool = True, ) -> None: b_ptr = _get_ptr(lora_b_stacked, device) K = max_lora_rank @@ -1408,7 +1414,7 @@ def _fused_moe_lora_expand( token_mapping_factor=1, naive_block_assignment=sorted_token_ids is None, MUL_ROUTED_WEIGHT=mul_routed_weight, - ADD_INPUTS=True, + ADD_INPUTS=add_inputs, USE_B_L2_CACHE=True, sort_c=False, IS_PRIMARY=False, @@ -1478,7 +1484,7 @@ def _fused_moe_lora( assert top_k_num == topk_weights.shape[1] # Fast path: single fused kernel - if not fully_sharded: + if not fully_sharded and max_lora_rank <= _FUSED_MOE_LORA_ONE_SHOT_MAX_RANK: M_pairs = topk_weights.numel() if ( sorted_token_ids is None @@ -1524,9 +1530,8 @@ def _fused_moe_lora( ) return - assert add_inputs, ( - "fused_moe_lora(add_inputs=False) is only supported on the " - "fully_sharded=False fast path" + assert add_inputs or not fully_sharded, ( + "fused_moe_lora(add_inputs=False) is not supported with fully_sharded=True" ) device = qcurr_hidden_states.device @@ -1660,6 +1665,7 @@ def _fused_moe_lora( offset, use_gdc=use_gdc, use_tma=use_tma, + add_inputs=add_inputs, ) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 3186b5d3c4..6a0a2b2358 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -31,6 +31,9 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501 FlashInferNVLinkTwoSidedPrepareAndFinalize, ) +from vllm.model_executor.layers.fused_moe.prepare_finalize.passthrough import ( + PassThroughPrepareAndFinalize, +) from vllm.platforms import current_platform from vllm.utils.import_utils import ( has_deep_ep, @@ -185,6 +188,10 @@ def maybe_make_prepare_finalize( all2all_manager: Any | None = None, input_dtype: torch.dtype | None = None, ) -> FusedMoEPrepareAndFinalize | None: + if moe.moe_parallel_config.use_passthrough_all2all: + # The experts dispatch and combine themselves, in every EP topology. + return PassThroughPrepareAndFinalize() + if not moe.moe_parallel_config.use_all2all_kernels: if not allow_new_interface: return None diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 302ed121fe..158bc90769 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1068,6 +1068,12 @@ class FusedMoEParallelConfig: and self.all2all_backend == "allgather_reducescatter" ) + @property + def use_passthrough_all2all(self): + # Not gated on use_all2all_kernels: the experts dispatch and combine + # internally in every EP topology, including TP-only. + return self.all2all_backend == "passthrough" + @property def use_mori_kernels(self): return self.use_all2all_kernels and self.all2all_backend in ( @@ -1530,6 +1536,10 @@ class FusedMoEConfig: def use_deepep_ll_kernels(self): return self.moe_parallel_config.use_deepep_ll_kernels + @property + def use_passthrough_all2all(self): + return self.moe_parallel_config.use_passthrough_all2all + @property def use_mori_kernels(self): return self.moe_parallel_config.use_mori_kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py index 4776a908fb..5423a93b5f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a16_moe.py @@ -374,6 +374,7 @@ class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: routing_method = self.moe_config.routing_method if routing_method == RoutingMethodType.DeepseekV4: diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 407dc097b1..bc1a4af6c0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -289,6 +289,7 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: assert self.moe_config.intermediate_size_per_partition_unpadded is not None assert self.moe_config.hidden_dim_unpadded is not None diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py index 5ee98a8c84..3f594dc0b3 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py @@ -198,6 +198,7 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: """Apply the monolithic 4-bit INT MoE forward pass. @@ -215,6 +216,7 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: Bias for expert scores routed_scaling_factor: Scaling factor for routing topk_group: Group size for topk + routing_replay_out: Unused; this kernel does not capture routing Returns: Output tensor after MoE computation diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py new file mode 100644 index 0000000000..8fbfcd9178 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_moe_ep.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer MoE-EP megakernels as a modular-kernel experts implementation. + +The megakernel routes on vLLM's top-k output and then dispatches, computes and +combines in one launch, so it is paired with a pass-through prepare/finalize. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + FlashInferMoeEp, + FlashInferMoeEpEpilogue, + FlashInferMoeEpWeights, + apply_topk_in_fc1, + flashinfer_moe_ep_unsupported_reasons, + supports_current_device, + validate_flashinfer_moe_ep_layer, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, + kNvfp4Static, + require_finite_positive_scales, +) +from vllm.model_executor.utils import replace_parameter + +if TYPE_CHECKING: + pass + +# Layer parameters the megakernel reads. After ``process_weights_after_loading`` +# they alias the megakernel's own tensors, so EPLB permutes exactly what it uses. +KERNEL_WEIGHT_NAMES = ("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale") + + +_supported_weight_quant_schemes: frozenset[QuantKey] = frozenset( + (kNvfp4Static, kMxfp4Static) +) + +_supported_activations: frozenset[MoEActivation] = frozenset((MoEActivation.SILU,)) + + +class FlashInferMoeEpExperts(mk.FusedMoEExpertsModular): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ) -> None: + super().__init__(moe_config, quant_config) + self._adapter: FlashInferMoeEp | None = None + + @staticmethod + def is_supported_config( + cls: type[mk.FusedMoEExperts], + moe_config: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + activation_format: mk.FusedMoEActivationFormat, + ) -> tuple[bool, str | None]: + supported, reason = mk.FusedMoEExperts.is_supported_config( + cls, moe_config, weight_key, activation_key, activation_format + ) + if not supported: + return False, reason + reasons = flashinfer_moe_ep_unsupported_reasons( + moe_config, weight_key, activation_key + ) + if reasons: + return False, f"kernel does not support {', '.join(reasons)}" + return True, None + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + return supports_current_device() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return weight_key in _supported_weight_quant_schemes + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in _supported_activations + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return ( + moe_parallel_config.use_ep + and not moe_parallel_config.use_batched_activation_format + ) + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + return (0,), (0,), (M, K) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Build the megakernel from the layer's canonical weights. + + The oracle has already put the weights in kernel format; the per-expert + weight global scales arrive through the quant config and become the + epilogue alphas. + """ + validate_flashinfer_moe_ep_layer(layer) # type: ignore[arg-type] + adapter = FlashInferMoeEp( + self.moe_config, + FlashInferMoeEpWeights( + w13=layer.w13_weight, + w2=layer.w2_weight, + w13_scale=getattr(layer, "w13_weight_scale", None), + w2_scale=getattr(layer, "w2_weight_scale", None), + ), + epilogue_from_quant_config(self.quant_config), + apply_topk_in_fc1=apply_topk_in_fc1( + self.moe_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + ), + ) + adapter.warmup() + for name, tensor in zip(KERNEL_WEIGHT_NAMES, adapter.kernel_weights()): + replace_parameter(layer, name, tensor) + self._adapter = adapter + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ) -> None: + if self._adapter is None: + raise RuntimeError("FlashInfer MoE-EP weights have not been processed") + output.copy_( + self._adapter(hidden_states, topk_ids.to(torch.int32), topk_weights) + ) + + +def epilogue_from_quant_config( + quant_config: FusedMoEQuantConfig, +) -> FlashInferMoeEpEpilogue: + """Per-expert epilogue constants from the canonical quant config. + + The weight global scales (``g1_alphas``/``g2_alphas``) become the fc1/fc2 + alphas. Activations are quantized dynamically inside the megakernel, so the + checkpoint's input scales drop out. MXFP4 checkpoints carry no global scales + and keep the default epilogue. The alphas alias the layer's parameters, so + EPLB permutations reach the kernel. + """ + fc1_alpha, fc2_alpha = quant_config.g1_alphas, quant_config.g2_alphas + if fc1_alpha is None and fc2_alpha is None: + return FlashInferMoeEpEpilogue() + assert fc1_alpha is not None and fc2_alpha is not None + require_finite_positive_scales("g1_alphas", fc1_alpha) + require_finite_positive_scales("g2_alphas", fc2_alpha) + if fc1_alpha.dim() != 1 or fc2_alpha.dim() != 1: + raise ValueError("the megakernel needs one weight global scale per expert") + return FlashInferMoeEpEpilogue( + fc1_alpha=fc1_alpha.float().contiguous(), + fc2_alpha=fc2_alpha.float().contiguous(), + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 5bf5065858..d23edaa01a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -1580,6 +1580,7 @@ class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: return triton_kernel_moe_forward( hidden_states=hidden_states, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 54227d3cd1..782d6bddd5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -272,6 +272,7 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: import flashinfer @@ -282,10 +283,6 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol # dummy 0-token forward has to keep the finalized (empty) form. defer = self.moe_config.should_defer_moe_finalize(num_tokens) - routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=num_tokens, - device=hidden_states.device, - ) flashinfer_output = flashinfer.fused_moe.trtllm_bf16_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -312,5 +309,4 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol num_tokens=num_tokens, top_k=self.topk, ) - self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) return routed_output diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index fb52445d84..532e25b745 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -450,6 +450,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: import flashinfer from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout @@ -487,11 +488,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit n_group = num_expert_group or 0 selected_topk_group = topk_group or 0 - routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=hidden_states.shape[0], - device=hidden_states.device, - ) - kwargs = dict( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -522,9 +518,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL: kwargs["activation_type"] = activation_type result = flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs) - self._maybe_dispatch_routing_replay( - routing_replay_out, num_tokens=hidden_states.shape[0] - ) return result def _apply_per_tensor( @@ -543,6 +536,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: # Delay import for non-CUDA. import flashinfer @@ -558,11 +552,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit else: assert not apply_router_weight_on_input - routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=hidden_states.shape[0], - device=hidden_states.device, - ) - out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -586,9 +575,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), routing_replay_out=routing_replay_out, ) - self._maybe_dispatch_routing_replay( - routing_replay_out, num_tokens=hidden_states.shape[0] - ) return out def apply( @@ -607,6 +593,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: if self.quant_config.block_shape is not None: return self._apply_block_scale( @@ -623,6 +610,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, topk_group=topk_group, + routing_replay_out=routing_replay_out, ) elif self.quant_config.is_per_tensor: return self._apply_per_tensor( @@ -638,6 +626,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit num_expert_group=num_expert_group, e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, + routing_replay_out=routing_replay_out, ) else: raise NotImplementedError( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index e1557a8f0d..6b62e8ba9c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -215,6 +215,7 @@ class TrtLlmMxfp4ExpertsMonolithic( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: from flashinfer import trtllm_fp4_block_scale_moe @@ -236,10 +237,6 @@ class TrtLlmMxfp4ExpertsMonolithic( device=hidden_states.device, ) - routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=num_tokens, - device=hidden_states.device, - ) flashinfer_output = trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, @@ -279,7 +276,6 @@ class TrtLlmMxfp4ExpertsMonolithic( top_k=self.topk, finalized_output=finalized_output, ) - self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) return routed_output diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py index e3e9bc6f2e..a2f06b1be9 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -139,6 +139,7 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501 flashinfer_trtllm_mxint4_moe, @@ -147,10 +148,6 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): assert self.w1_scale is not None assert self.w2_scale is not None - routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=hidden_states.shape[0], - device=hidden_states.device, - ) result = flashinfer_trtllm_mxint4_moe( x=hidden_states, router_logits=router_logits, @@ -169,7 +166,4 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): routing_method_type=self.routing_method, routing_replay_out=routing_replay_out, ) - self._maybe_dispatch_routing_replay( - routing_replay_out, num_tokens=hidden_states.shape[0] - ) return result diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 9dcfc6af3f..c9af86b43a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -515,6 +515,7 @@ class TrtLlmNvFp4ExpertsMonolithic( e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: import flashinfer @@ -545,10 +546,6 @@ class TrtLlmNvFp4ExpertsMonolithic( # dummy 0-token forward has to keep the finalized (empty) form. defer = self.moe_config.should_defer_moe_finalize(num_tokens) - routing_replay_out = self._maybe_make_routing_replay_buffer( - num_tokens=num_tokens, - device=hidden_states.device, - ) # Invoke kernel. # NOTE: Activation padding and output # truncation are handled by the MoE runner's @@ -592,5 +589,4 @@ class TrtLlmNvFp4ExpertsMonolithic( num_tokens=num_tokens, top_k=self.topk, ) - self._maybe_dispatch_routing_replay(routing_replay_out, num_tokens=num_tokens) return routed_output diff --git a/vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py b/vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py index 0076b65966..b98fb226bb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py @@ -151,6 +151,7 @@ class ZentorchExpertsInt4(mk.FusedMoEExpertsMonolithic): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor: if apply_router_weight_on_input: raise NotImplementedError( diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py new file mode 100644 index 0000000000..bac5345255 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import contextlib +import sys +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, cast + +import torch + +from vllm.config import get_current_vllm_config +from vllm.config.kernel import ( + FLASHINFER_MOE_EP_BACKENDS, + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, +) +from vllm.distributed import get_ep_group +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Static, + kNvfp4Static, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import _import_deep_gemm + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + +FlashInferMoeEpKernel = Literal["cutedsl", "deep_gemm"] +FlashInferMoeEpWeightFormat = Literal["nvfp4", "mxfp4"] +_WEIGHT_FORMATS: dict[QuantKey | None, FlashInferMoeEpWeightFormat] = { + kNvfp4Static: "nvfp4", + kMxfp4Static: "mxfp4", +} + +_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)}) + + +def supports_current_device() -> bool: + capability = current_platform.get_device_capability() + return capability is not None and ( + (capability.major, capability.minor) in _SUPPORTED_CAPABILITIES + ) + + +_E2M1_LUT = ( + 0.0, + 0.5, + 1.0, + 1.5, + 2.0, + 3.0, + 4.0, + 6.0, + -0.0, + -0.5, + -1.0, + -1.5, + -2.0, + -3.0, + -4.0, + -6.0, +) + + +@dataclass(frozen=True) +class FlashInferMoeEpBackendSpec: + kernel: FlashInferMoeEpKernel + weight_formats: frozenset[FlashInferMoeEpWeightFormat] + + +@dataclass(frozen=True) +class FlashInferMoeEpWeights: + w13: torch.Tensor + w2: torch.Tensor + w13_scale: torch.Tensor | None = None + w2_scale: torch.Tensor | None = None + + +@dataclass(frozen=True) +class FlashInferMoeEpEpilogue: + input_norm_const: float = 1.0 + fc1_alpha: torch.Tensor | None = None + fc2_alpha: torch.Tensor | None = None + fc1_norm_const: torch.Tensor | None = None + + +_BACKEND_SPECS = { + FLASHINFER_MOE_EP_CUTEDSL: FlashInferMoeEpBackendSpec( + kernel="cutedsl", + weight_formats=frozenset({"nvfp4", "mxfp4"}), + ), + FLASHINFER_MOE_EP_DEEP_GEMM: FlashInferMoeEpBackendSpec( + kernel="deep_gemm", + weight_formats=frozenset({"mxfp4"}), + ), +} + +assert frozenset(_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS + + +def is_flashinfer_moe_ep_backend(moe_backend: str) -> bool: + return moe_backend in _BACKEND_SPECS + + +def flashinfer_moe_ep_backend_spec( + moe_backend: str, +) -> FlashInferMoeEpBackendSpec: + try: + return _BACKEND_SPECS[moe_backend] + except KeyError: + raise ValueError( + f"{moe_backend!r} is not a FlashInfer MoE-EP backend; expected " + f"one of {sorted(_BACKEND_SPECS)}" + ) from None + + +def flashinfer_moe_ep_unsupported_reasons( + moe: FusedMoEConfig, + weight_key: QuantKey | None, + activation_key: QuantKey | None, +) -> tuple[str, ...]: + """Deployment constraints of the megakernel beyond the framework's checks. + + ``activation_key`` is ``None`` for MXFP4 checkpoints (the kernel quantizes + bf16 activations itself) but marks a W4A16 recipe for NVFP4 weights. + """ + spec = flashinfer_moe_ep_backend_spec(moe.moe_backend) + unsupported: list[str] = [] + weight_format = _WEIGHT_FORMATS.get(weight_key) + if weight_format is None: + unsupported.append(f"weight scheme {weight_key}") + elif weight_format not in spec.weight_formats: + unsupported.append(f"{weight_format.upper()} weights") + if weight_key == kNvfp4Static and activation_key is None: + unsupported.append("A16 activations with NVFP4 weights") + if moe.skip_final_all_reduce: + unsupported.append("skip_final_all_reduce") + if moe.in_dtype != torch.bfloat16: + unsupported.append(f"activation dtype {moe.in_dtype}") + if moe.has_bias: + unsupported.append("expert bias") + if moe.swiglu_alpha is not None or moe.swiglu_beta is not None: + unsupported.append("custom SwiGLU alpha or beta") + if ( + spec.kernel == "deep_gemm" + and moe.routing_method is not RoutingMethodType.DeepseekV4 + ): + unsupported.append(f"routing method {moe.routing_method.name}") + + vllm_config = get_current_vllm_config() + if vllm_config.weight_transfer_config is not None: + unsupported.append("runtime weight transfer") + if vllm_config.parallel_config.enable_dbo: + unsupported.append("dual batch overlap") + if spec.kernel == "deep_gemm" and vllm_config.parallel_config.enable_eplb: + unsupported.append("EPLB") + return tuple(unsupported) + + +def validate_flashinfer_moe_ep_layer(layer: RoutedExperts) -> None: + if layer.expert_map_manager.placement_strategy != "linear": + raise ValueError( + f"{layer.moe_config.moe_backend} requires contiguous linear " + "expert placement" + ) + + +def apply_topk_in_fc1( + moe: FusedMoEConfig, + *, + apply_router_weight_on_input: bool = False, +) -> bool: + return ( + apply_router_weight_on_input + or moe.routing_method is RoutingMethodType.DeepseekV4 + ) + + +def _dequant_mxfp4_ue8m0_gran32( + packed: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + raw = packed.view(torch.uint8) + lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device) + values = torch.empty( + raw.shape[0], + raw.shape[1] * 2, + dtype=torch.float32, + device=raw.device, + ) + values[:, ::2] = lut[(raw & 0x0F).to(torch.int64)] + values[:, 1::2] = lut[(raw >> 4).to(torch.int64)] + encoded_scale = ( + scale + if scale.dtype == torch.float8_e8m0fnu + else scale.view(torch.float8_e8m0fnu) + ) + decoded_scale = encoded_scale.float() + return (values * decoded_scale.repeat_interleave(32, dim=-1)).to(torch.bfloat16) + + +def _dequant_mxfp4_experts( + weight: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + num_experts, output_size, packed_input_size = weight.shape + result = torch.empty( + num_experts, + output_size, + packed_input_size * 2, + dtype=torch.bfloat16, + device=weight.device, + ) + for expert_id in range(num_experts): + result[expert_id] = _dequant_mxfp4_ue8m0_gran32( + weight[expert_id], scale[expert_id] + ) + return result + + +def mxfp4_moe_ep_weights( + moe_backend: str, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, +) -> FlashInferMoeEpWeights: + spec = flashinfer_moe_ep_backend_spec(moe_backend) + if spec.kernel == "deep_gemm": + return FlashInferMoeEpWeights(w13, w2, w13_scale, w2_scale) + return FlashInferMoeEpWeights( + w13=_dequant_mxfp4_experts(w13, w13_scale), + w2=_dequant_mxfp4_experts(w2, w2_scale), + ) + + +class _FlashInferMoeEpApi: + def __init__(self) -> None: + try: + from flashinfer.moe_ep import ( + BootstrapConfig, + DeepGemmMegaMoeConfig, + FleetParams, + MegaConfig, + MoEEpMegaLayer, + MoEEpTensors, + MoEWeightPack, + Nvfp4CutedslMegaMoeConfig, + ) + except ImportError as err: + raise RuntimeError( + "FlashInfer MoE-EP requires the public API shipped with " + "FlashInfer 0.7.0 or newer" + ) from err + + self.BootstrapConfig = BootstrapConfig + self.DeepGemmMegaMoeConfig = DeepGemmMegaMoeConfig + self.FleetParams = FleetParams + self.MegaConfig = MegaConfig + self.MoEEpMegaLayer = MoEEpMegaLayer + self.MoEEpTensors = MoEEpTensors + self.MoEWeightPack = MoEWeightPack + self.Nvfp4CutedslMegaMoeConfig = Nvfp4CutedslMegaMoeConfig + + +def _load_flashinfer_moe_ep_api() -> _FlashInferMoeEpApi: + return _FlashInferMoeEpApi() + + +def _expose_deep_gemm_to_flashinfer() -> None: + deep_gemm = _import_deep_gemm() + if deep_gemm is None: + raise RuntimeError( + "FlashInfer MoE-EP DeepGEMM requires a usable DeepGEMM installation" + ) + sys.modules.setdefault("deep_gemm", deep_gemm) + + +class FlashInferMoeEp: + def __init__( + self, + moe: FusedMoEConfig, + weights: FlashInferMoeEpWeights, + epilogue: FlashInferMoeEpEpilogue | None = None, + *, + apply_topk_in_fc1: bool, + ) -> None: + spec = flashinfer_moe_ep_backend_spec(moe.moe_backend) + if spec.kernel == "deep_gemm": + _expose_deep_gemm_to_flashinfer() + api = _load_flashinfer_moe_ep_api() + if epilogue is None: + epilogue = FlashInferMoeEpEpilogue() + ep_group = get_ep_group() + bootstrap = api.BootstrapConfig( + world_size=ep_group.world_size, + rank=ep_group.rank_in_group, + process_group=ep_group.device_group, + device=torch.accelerator.current_device_index(), + ) + fleet_params = api.FleetParams( + num_experts=moe.num_experts, + max_tokens_per_rank=moe.max_num_tokens, + token_hidden_size=moe.hidden_dim, + ) + weight_pack = api.MoEWeightPack( + w13=weights.w13, + w2=weights.w2, + w13_scale=weights.w13_scale, + w2_scale=weights.w2_scale, + ) + if spec.kernel == "cutedsl": + megakernel = api.Nvfp4CutedslMegaMoeConfig( + intermediate_size=moe.intermediate_size, + top_k=moe.experts_per_token, + gate_up_clamp=moe.swiglu_limit, + fast_math=True, + apply_topk_in_fc1=apply_topk_in_fc1, + enable_in_kernel_fc2_reduce=False, + combine_dtype="bf16", + input_norm_const=epilogue.input_norm_const, + fc1_alpha=None, + fc2_alpha=None, + fc1_norm_const=None, + knobs=None, + ) + else: + megakernel = api.DeepGemmMegaMoeConfig( + intermediate_size=moe.intermediate_size, + top_k=moe.experts_per_token, + activation_clamp=moe.swiglu_limit, + fast_math=True, + ) + backend = api.MegaConfig( + megakernel=megakernel, + quantize_input=True, + preprocess_weights=True, + ) + + self._moe_ep_tensors_cls = api.MoEEpTensors + self._mega_layer: Any | None = api.MoEEpMegaLayer( + bootstrap, + fleet_params, + weight_pack, + backend, + ) + self._device = weights.w13.device + self._hidden_size = moe.hidden_dim + self._max_num_tokens = moe.max_num_tokens + self._num_experts = moe.num_experts + self._top_k = moe.experts_per_token + self._epilogue: FlashInferMoeEpEpilogue | None = epilogue + self._return_workspace_view = moe.routing_method is RoutingMethodType.DeepseekV4 + + @property + def can_overlap_shared_experts(self) -> bool: + return False + + @property + def output_is_reduced(self) -> bool: + return True + + @property + def topk_indices_dtype(self) -> torch.dtype: + return torch.int32 + + @property + def is_monolithic(self) -> bool: + return False + + def _tensors( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> Any: + epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue) + return self._moe_ep_tensors_cls( + hidden_states=hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + fc1_alpha=epilogue.fc1_alpha, + fc2_alpha=epilogue.fc2_alpha, + fc1_norm_const=epilogue.fc1_norm_const, + ) + + def __call__( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + if self._mega_layer is None: + raise RuntimeError("FlashInfer MoE-EP layer was destroyed") + if hidden_states.shape[0] > self._max_num_tokens: + raise ValueError( + f"FlashInfer MoE-EP got {hidden_states.shape[0]} tokens, but " + f"the workspace supports at most {self._max_num_tokens}" + ) + tensors = self._tensors(hidden_states, topk_ids, topk_weights) + if self._return_workspace_view and getattr( + self._mega_layer, "supports_output_view", False + ): + return self._mega_layer.forward(tensors, return_workspace_view=True) + return self._mega_layer(tensors) + + @torch.inference_mode() + def warmup(self) -> None: + if self._mega_layer is None: + return + hidden_states = torch.zeros( + 1, + self._hidden_size, + dtype=torch.bfloat16, + device=self._device, + ) + topk_ids = ( + torch.arange(self._top_k, dtype=torch.int32, device=self._device) + .mul_(self._num_experts // self._top_k) + .view(1, self._top_k) + ) + topk_weights = torch.full( + (1, self._top_k), + 1.0 / self._top_k, + dtype=torch.float32, + device=self._device, + ) + self._mega_layer.warmup(self._tensors(hidden_states, topk_ids, topk_weights)) + + def kernel_weights( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Kernel-resident fc1 weight, fc1 scales, fc2 weight, fc2 scales. + + Contiguous per-expert byte tensors that alias the memory the megakernel + reads, so a caller may permute experts in place (EPLB). + """ + if self._mega_layer is None: + raise RuntimeError("FlashInfer MoE-EP layer was destroyed") + (fc1_weight, fc1_sf), (fc2_weight, fc2_sf) = self._mega_layer._transformed + return ( + _expert_major_bytes(fc1_weight), + _expert_major_bytes(fc1_sf), + _expert_major_bytes(fc2_weight), + _expert_major_bytes(fc2_sf), + ) + + def expert_weight_views(self) -> tuple[torch.Tensor, ...]: + """Every per-expert tensor the kernel reads: weights plus epilogue.""" + epilogue = cast(FlashInferMoeEpEpilogue, self._epilogue) + vectors = (epilogue.fc1_alpha, epilogue.fc2_alpha, epilogue.fc1_norm_const) + return (*self.kernel_weights(), *(v for v in vectors if v is not None)) + + def destroy(self) -> None: + mega_layer = self._mega_layer + if mega_layer is None: + return + mega_layer.destroy() + self._mega_layer = None + self._epilogue = None + + def __del__(self) -> None: + with contextlib.suppress(Exception): + self.destroy() + + +def _expert_major_bytes(tensor: torch.Tensor) -> torch.Tensor: + """Contiguous storage of a per-expert kernel tensor, 1-byte dtypes as uint8. + + The CuTeDSL kernel keeps (E, K, N) transpose views over (E, N, K) storage. + """ + if not tensor.is_contiguous(): + tensor = tensor.transpose(1, 2) + if not tensor.is_contiguous(): + raise ValueError(f"expert tensor is not expert-major: {tensor.stride()}") + return tensor.view(torch.uint8) if tensor.element_size() == 1 else tensor diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index f563c819e4..9bde03795c 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, ) from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsSink, +) from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, SharedExpertsOrder, @@ -979,9 +982,6 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts): def is_monolithic() -> bool: return True - routing_replay_capture_fn: Callable[[torch.Tensor], None] | None = None - _routing_replay_buffer: torch.Tensor | None = None - def supports_routing_replay_capture(self) -> bool: """Whether this expert supports routing replay capture. @@ -990,53 +990,6 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts): """ return False - def set_capture_fn( - self, - capture_fn: Callable[[torch.Tensor], None] | None, - ) -> None: - self.routing_replay_capture_fn = capture_fn - if capture_fn is None: - self._routing_replay_buffer = None - return - # Allocate for per-rank batches gathered across the DP or EP group. - dispatch_group_size = ( - self.moe_config.ep_size - if self.moe_config.use_ep - else self.moe_config.dp_size - ) - max_num_replay_tokens = self.moe_config.max_num_tokens * dispatch_group_size - self._routing_replay_buffer = torch.empty( - (max_num_replay_tokens, self.moe_config.experts_per_token), - dtype=torch.int16, - device=self.moe_config.device, - ) - - def _maybe_make_routing_replay_buffer( - self, - num_tokens: int, - device: torch.device, - ) -> torch.Tensor | None: - if self.routing_replay_capture_fn is None: - return None - buf = self._routing_replay_buffer - assert buf is not None - if buf.shape[0] < num_tokens or buf.device != device: - raise ValueError( - "Routing replay buffer was initialized for " - f"{buf.shape[0]} tokens on {buf.device}, but the kernel " - f"received {num_tokens} tokens on {device}." - ) - return buf - - def _maybe_dispatch_routing_replay( - self, - routing_replay_out: torch.Tensor | None, - num_tokens: int, - ) -> None: - if routing_replay_out is None or self.routing_replay_capture_fn is None: - return - self.routing_replay_capture_fn(routing_replay_out[:num_tokens]) - def apply( self, hidden_states: torch.Tensor, @@ -1053,10 +1006,14 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts): e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + routing_replay_out: torch.Tensor | None = None, ) -> torch.Tensor | UnfinalizedMoEOutput: """Same as ``FusedMoEExperts.apply``, except uses router_logits as opposed to the topk_ids and topk_weights. This is useful for kernels with fused router and fused_experts (e.g. FLASHINFER_TRTLLM). + + Kernels that ``supports_routing_replay_capture`` write each token's + routed expert ids into the leading rows of ``routing_replay_out``. """ raise NotImplementedError @@ -1563,6 +1520,8 @@ class FusedMoEKernelMonolithicImpl: e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + *, + routing_sink: RoutedExpertsSink | None, ) -> torch.Tensor | UnfinalizedMoEOutput: """Same as forward(), except uses router_logits as opposed to the topk_ids and topk_weights. This is used for kernels @@ -1575,6 +1534,15 @@ class FusedMoEKernelMonolithicImpl: defer_input_quant=self.fused_experts.expects_unquantized_inputs, ) + routing_replay_out = None + if routing_sink is not None: + # Checked per call: a weight reload may rebuild a different kernel. + if not self.fused_experts.supports_routing_replay_capture(): + raise ValueError( + "Routed-experts capture is not supported with monolithic MoE " + f"kernel {type(self.fused_experts).__name__}." + ) + routing_replay_out = routing_sink.buffer[: len(a1q)] fused_out = self.fused_experts.apply( hidden_states=a1q, w1=w1, @@ -1590,7 +1558,10 @@ class FusedMoEKernelMonolithicImpl: e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, topk_group=topk_group, + routing_replay_out=routing_replay_out, ) + if routing_sink is not None: + routing_sink.capture_fn(routing_replay_out) if isinstance(fused_out, UnfinalizedMoEOutput): if not self.prepare_finalize.supports_deferred_moe_finalize(): @@ -1698,6 +1669,8 @@ class FusedMoEKernel: e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float | None = None, topk_group: int | None = None, + *, + routing_sink: RoutedExpertsSink | None, ) -> torch.Tensor | UnfinalizedMoEOutput: assert isinstance(self.impl, FusedMoEKernelMonolithicImpl) return self.impl.apply( @@ -1713,6 +1686,7 @@ class FusedMoEKernel: e_score_correction_bias=e_score_correction_bias, routed_scaling_factor=routed_scaling_factor, topk_group=topk_group, + routing_sink=routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py index 8d16a41975..63d33eadd8 100644 --- a/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py +++ b/vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py @@ -3,6 +3,7 @@ import torch from torch._subclasses.fake_tensor import FakeTensor +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -75,11 +76,18 @@ def moe_fused_mul_sum_kernel( def _heuristic_config( hidden_size: int, element_size: int, + device_index: int = 0, ): is_fp32 = element_size > 2 max_block_k = 256 if is_fp32 else 512 - BLOCK_K = max(128, min(triton.next_power_of_2(hidden_size), max_block_k)) num_warps = 4 if is_fp32 else 2 + if not is_fp32 and hidden_size >= 1024 and current_platform.is_cuda(): + capability = current_platform.get_device_capability(device_index) + if capability is not None and capability >= (7, 5): + # Larger tiles reduce the serial hidden-dimension loop. + max_block_k = 4096 + num_warps = 8 if capability == (7, 5) else 16 + BLOCK_K = max(128, min(triton.next_power_of_2(hidden_size), max_block_k)) num_stages = 3 return BLOCK_K, num_warps, num_stages @@ -149,6 +157,7 @@ def moe_fused_mul_sum( BLOCK_K, num_warps, num_stages = _heuristic_config( hidden_size, inputs.element_size(), + inputs.device.index or 0, ) grid = (num_tokens,) moe_fused_mul_sum_kernel[grid]( diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 27d59ba2e0..a3c65dc4a5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -8,7 +8,11 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import envs from vllm.config import get_current_vllm_config -from vllm.config.kernel import MoEBackend +from vllm.config.kernel import ( + FLASHINFER_MOE_EP_CUTEDSL, + FLASHINFER_MOE_EP_DEEP_GEMM, + MoEBackend, +) from vllm.config.quantization import QuantizationConfigArgs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( @@ -145,6 +149,9 @@ class Mxfp4MoeBackend(Enum): EMULATION = "EMULATION" # Humming HUMMING = "HUMMING" + # FlashInfer MoE-EP megakernels: routing outside, dispatch/GEMMs/combine fused. + FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL" + FLASHINFER_MOE_EP_DEEP_GEMM = "FLASHINFER_MOE_EP_DEEP_GEMM" # Backends that share the same TRTLLM weight format @@ -165,6 +172,13 @@ B12X_BACKENDS = ( Mxfp4MoeBackend.B12X_MXFP4_BF16, ) +# Megakernels consume the checkpoint's MXFP4 tensors through the shared +# FlashInfer MoE-EP adapter; vLLM does not convert their weight layout. +FLASHINFER_MOE_EP_MXFP4_BACKENDS = ( + Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, + Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM, +) + def backend_to_kernel_cls( backend: Mxfp4MoeBackend, @@ -291,6 +305,13 @@ def backend_to_kernel_cls( return [OCP_MXQuantizationEmulationTritonExperts] + elif backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( # noqa: E501 + FlashInferMoeEpExperts, + ) + + return [FlashInferMoeEpExperts] + else: raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}") @@ -314,6 +335,10 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, ], "flashinfer_cutlass_afp8": [Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8], + "flashinfer_moe_ep_cutedsl": [Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL], + "flashinfer_moe_ep_mega_deep_gemm": [ + Mxfp4MoeBackend.FLASHINFER_MOE_EP_DEEP_GEMM + ], "triton": [Mxfp4MoeBackend.TRITON], "triton_unfused": [Mxfp4MoeBackend.TRITON_UNFUSED], "humming": [Mxfp4MoeBackend.HUMMING], @@ -735,6 +760,9 @@ 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 in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + # The megakernel takes the model's dimensions and validates them itself. + return hidden_size, intermediate_size if backend in B12X_BACKENDS: # b12x plans for the exact model dimensions. B12xExperts validates the # required MXFP4 block alignment before selecting the backend. @@ -1827,6 +1855,29 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias=w2_bias, _cache_permute_indices=_cache_permute_indices, ) + elif mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + mxfp4_moe_ep_weights, + ) + + # DeepGEMM consumes MXFP4 directly; CuTeDSL gets bf16 and requantizes. + weights = mxfp4_moe_ep_weights( + FLASHINFER_MOE_EP_CUTEDSL + if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL + else FLASHINFER_MOE_EP_DEEP_GEMM, + w13_weight, + w2_weight, + w13_weight_scale, + w2_weight_scale, + ) + return ( + weights.w13, + weights.w2, + weights.w13_scale, + weights.w2_scale, + w13_bias, + w2_bias, + ) elif mxfp4_backend == Mxfp4MoeBackend.CPU: from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( prepare_mxfp4_moe_layer_for_cpu, @@ -1874,6 +1925,9 @@ def make_mxfp4_moe_quant_config( layer: "RoutedExperts | None" = None, ) -> FusedMoEQuantConfig | None: """Create a FusedMoEQuantConfig for the given MXFP4 backend.""" + if mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + # The megakernel quantizes activations itself and owns its scales. + return FusedMoEQuantConfig.make("nvfp4", weight_dtype="mxfp4") if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8: return mxfp4_mxfp8_moe_quant_config( w1_bias=w1_bias, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index e1952c0032..68ba62782c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -50,6 +50,7 @@ class NvFp4MoeBackend(Enum): MARLIN = "MARLIN" HUMMING = "HUMMING" EMULATION = "EMULATION" + FLASHINFER_MOE_EP_CUTEDSL = "FLASHINFER_MOE_EP_CUTEDSL" FLASHINFER_NVFP4_MOE_BACKENDS = [ @@ -148,6 +149,12 @@ def backend_to_kernel_cls( ) return [Nvfp4QuantizationEmulationTritonExperts] + elif backend == NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import ( # noqa: E501 + FlashInferMoeEpExperts, + ) + + return [FlashInferMoeEpExperts] else: raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}") @@ -160,6 +167,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, + "flashinfer_moe_ep_cutedsl": NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X, "marlin": NvFp4MoeBackend.MARLIN, "humming": NvFp4MoeBackend.HUMMING, @@ -207,6 +215,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.EMULATION, @@ -281,8 +290,9 @@ def select_nvfp4_moe_backend( f"Model sets swiglu_limit={config.swiglu_limit}, but the " f"explicitly requested moe_backend={runner_backend!r} does " f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', " - f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', " - f"'b12x', 'marlin', or 'humming' instead." + f"'flashinfer_cutlass', 'flashinfer_cutedsl', " + f"'flashinfer_moe_ep_cutedsl', 'cutlass', 'b12x', 'marlin', " + f"or 'humming' instead." ) return _return_or_raise( requested_backend, config, weight_key, activation_key, activation_format @@ -491,6 +501,9 @@ def convert_to_nvfp4_moe_kernel_format( # for other experts - other selection strategies may be used. a13_scale = 1.0 / a13_scale.max().to(torch.float32) a2_scale = 1.0 / a2_scale.max().to(torch.float32) + elif nvfp4_backend == NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL: + # The megakernel transforms the checkpoint tensors itself at init. + pass else: raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}") diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py new file mode 100644 index 0000000000..a6c541d2b9 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/passthrough.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prepare/finalize stages that do no expert-parallel communication. + +Selected by ``all2all_backend="passthrough"``: the experts implementation +dispatches, combines and reduces itself, so nothing is left for these stages. +""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, +) + + +class PassThroughPrepareAndFinalize(MoEPrepareAndFinalizeNoDPEPModular): + """Hands activations and routing straight to experts that own their EP + communication.""" + + def output_is_reduced(self) -> bool: + return True + + def supports_deferred_moe_finalize(self) -> bool: + # The experts apply the top-k weights and combine internally; there + # is no finalize left for a consumer to take over. + return False + + def prepare( + self, + a1: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, + expert_map: torch.Tensor | None, + apply_router_weight_on_input: bool, + quant_config: FusedMoEQuantConfig, + defer_input_quant: bool = False, + ) -> mk.PrepareResultType: + return a1, None, None, None, None diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index de08d9aa54..f3c9522df1 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -20,6 +20,9 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( FusedMoEMethodBase, ) from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsSink, +) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) @@ -138,6 +141,8 @@ class RoutedExperts(PluggableLayer): self.swiglu_alpha = swiglu_alpha self.swiglu_beta = swiglu_beta self.e_score_correction_bias = e_score_correction_bias + # Set by bind_routed_experts_capturer for monolithic kernels. + self.routing_sink: RoutedExpertsSink | None = None self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters self._loaded_expert_biases: set[str] = set() diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index c49d614549..5d111fa635 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -15,6 +15,7 @@ import torch from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_tp_group from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.platforms import current_platform logger = logging.getLogger(__name__) @@ -186,6 +187,23 @@ class RoutedExpertsCapturer: return self.device_buffer[:num_tokens].to(self.output_dtype) +class RoutedExpertsSink: + """Layer-owned buffer and callback for the expert ids a monolithic kernel + routes to; it outlives kernel rebuilds on weight reload.""" + + def __init__( + self, moe_config: FusedMoEConfig, capture_fn: Callable[[torch.Tensor], None] + ) -> None: + # Sized for per-rank batches gathered across the DP or EP group. + group_size = moe_config.ep_size if moe_config.use_ep else moe_config.dp_size + self.buffer = torch.empty( + (moe_config.max_num_tokens * group_size, moe_config.experts_per_token), + dtype=torch.int16, + device=moe_config.device, + ) + self.capture_fn = capture_fn + + def bind_routed_experts_capturer( model: torch.nn.Module, capturer: RoutedExpertsCapturer, @@ -219,7 +237,9 @@ def bind_routed_experts_capturer( "Routed-experts capture is not supported with monolithic " f"MoE kernel {type(fused_experts).__name__}." ) - fused_experts.set_capture_fn(capture_fn) + module.routed_experts.routing_sink = RoutedExpertsSink( + fused_experts.moe_config, capture_fn + ) num_bound += 1 elif isinstance(module.router, BaseRouter): module.router.set_capture_fn(capture_fn) diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index 17481d4a91..e9a401cf45 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -362,4 +362,5 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) diff --git a/vllm/model_executor/layers/indexer_topk.py b/vllm/model_executor/layers/indexer_topk.py index c3422445b2..f0da0312e1 100644 --- a/vllm/model_executor/layers/indexer_topk.py +++ b/vllm/model_executor/layers/indexer_topk.py @@ -163,6 +163,14 @@ class SparseIndexerTopk(torch.nn.Module): current_platform.has_device_capability(90) and not current_platform.is_device_capability_family(120) ) + self._is_rocm = current_platform.is_rocm() + self._aiter_enabled = False + self._aiter_capable = False + if self._is_rocm: + from vllm._aiter_ops import rocm_aiter_ops + + self._aiter_enabled = bool(rocm_aiter_ops.is_enabled()) + self._aiter_capable = bool(rocm_aiter_ops.is_indexer_top_k_enabled()) def resolve_backend( self, logits: torch.Tensor, topk_tokens: int, num_rows: int @@ -193,6 +201,19 @@ class SparseIndexerTopk(torch.nn.Module): f"inputs violate DeepSelect's constraints: dtype={logits.dtype}," f" stride={logits.stride()}, topk={topk_tokens}" ) + elif self._backend == "aiter": + if not self._is_rocm: + failures.append("requires a ROCm platform") + elif not self._aiter_enabled: + failures.append("requires AITER (VLLM_ROCM_USE_AITER=1)") + elif not self._aiter_capable: + failures.append( + "AITER has no tuned indexer top-k kernels for this GPU arch" + ) + if topk_tokens not in (512, 1024, 2048, 4096): + failures.append( + f"topk_tokens must be in (512, 1024, 2048, 4096), got {topk_tokens}" + ) elif self._backend == "flashinfer": if not self._is_cuda: failures.append("requires a CUDA platform") @@ -216,11 +237,14 @@ class SparseIndexerTopk(torch.nn.Module): self, logits: torch.Tensor, topk_tokens: int, num_rows: int ) -> str: """The priority chain: cooperative -> persistent -> per_row. - deep_select/flashinfer/torch are opt-in only. + aiter/deep_select/flashinfer/torch are opt-in only. cooperative_topk is preferred whenever it is applicable, i.e. within its AUTO_COOPERATIVE_MAX_ROWS row limit; larger batches go to persistent_topk. + + aiter not auto-resolved as it does not outperform the per-row on every shape, + hence we delegate to the caller to decide when to use it. """ if not self._cooperative_constraints(logits, topk_tokens, num_rows): return "cooperative" @@ -285,7 +309,8 @@ class SparseIndexerTopk(torch.nn.Module): max_seq_len: int, ) -> None: """Run the resolved decode top-k implementation, writing into - topk_indices (int32, -1 fill for rows shorter than topk_tokens).""" + topk_indices (int32, -1 fill for rows shorter than topk_tokens). + """ backend = self.resolve_backend(logits, topk_tokens, logits.shape[0]) if backend == "deep_select": row_ends = self._row_ends(seq_lens, next_n, logits.shape[0]) @@ -344,15 +369,37 @@ class SparseIndexerTopk(torch.nn.Module): ) < row_ends.unsqueeze(1) indices = torch.where(in_range, indices, -1) topk_indices.copy_(indices) - else: - assert backend == "per_row", f"unknown topk backend: {backend}" - ops.top_k_per_row_decode( + elif backend == "aiter": + from vllm._aiter_ops import rocm_aiter_ops + + # AITER's decode kernel only implements the 1D/next_n == 1 form, + # which is equivalent once the ragged ends are expanded. + rocm_aiter_ops.indexer_top_k_decode( logits, - next_n, - seq_lens, + 1, + self._row_ends(seq_lens, next_n, logits.shape[0]), topk_indices, - logits.shape[0], - logits.stride(0), - logits.stride(1), topk_tokens, ) + else: + assert backend == "per_row", f"unknown topk backend: {backend}" + self._per_row(logits, seq_lens, next_n, topk_indices, topk_tokens) + + @staticmethod + def _per_row( + logits: torch.Tensor, + seq_lens: torch.Tensor, + next_n: int, + topk_indices: torch.Tensor, + topk_tokens: int, + ) -> None: + ops.top_k_per_row_decode( + logits, + next_n, + seq_lens, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + topk_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 0f22f0c911..3f13dc19d3 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 @@ -551,28 +551,33 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): 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 - ) + if self.speculative_config is not None: + logger.info_once( + "Falling back to the Triton GDN spec decode path: %s", + reason, + ) self.gdn_decode_kernel = "triton" elif current_platform.is_cpu(): self.gdn_decode_kernel = "CPU" elif current_platform.is_xpu(): self.gdn_decode_kernel = "XPU" - self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda" + self.enable_fused_gdn_spec_decode = ( + self.gdn_decode_kernel == "cuda" and self.speculative_config is not None + ) # RecoverSSM verify replaces the per-draft snapshot path (Triton decode # only) self.use_gdn_recoverssm = self._uses_gdn_recoverssm() if self.use_gdn_recoverssm: self.gdn_decode_kernel = "triton" - self.enable_fused_gdn_decode = False + self.enable_fused_gdn_spec_decode = False self.enable_packed_recurrent_decode = False logger.info_once( "GDN RecoverSSM speculative verify active (spec_query_len %d)", self.num_spec + 1, ) - logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel) + if self.speculative_config is not None: + logger.info_once("GDN spec decode kernel: %s", self.gdn_decode_kernel) compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: @@ -969,12 +974,12 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention): mixed_qkvz, _ = self.in_proj_qkvz(hidden_states) ba, _ = self.in_proj_ba(hidden_states) - use_fused_gdn_decode = ( - self.enable_fused_gdn_decode + use_fused_gdn_spec_decode = ( + self.enable_fused_gdn_spec_decode and hidden_states.dtype == torch.bfloat16 and self.norm.weight.dtype in (torch.bfloat16, torch.float32) ) - if use_fused_gdn_decode: + if use_fused_gdn_spec_decode: core_attn_out = torch.zeros( (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim), dtype=hidden_states.dtype, diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 0c3e980551..e0eccc708e 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -790,6 +790,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index ed430f7445..37260ac5c8 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -779,4 +779,5 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index 643fd75e82..4dd1e7aafe 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -24,10 +24,15 @@ from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsMxFp4, ) +from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import ( + is_flashinfer_moe_ep_backend, +) from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( B12X_BACKENDS, + FLASHINFER_MOE_EP_MXFP4_BACKENDS, Mxfp4MoeBackend, + convert_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, select_mxfp4_moe_backend, @@ -38,7 +43,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_moe_fp4_layer_for_marlin, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform logger = init_logger(__name__) @@ -49,12 +54,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): super().__init__(moe) self.group_size = 32 self.mxfp4_backend = Mxfp4MoeBackend.MARLIN + self.experts_cls: type[mk.FusedMoEExperts] | None = None + # Backend selection must match the weight preparation below: CUTLASS - # swizzles scales, b12x and XPU consume checkpoint packing, and Marlin - # repacks weights and scales. + # swizzles scales, b12x and XPU consume checkpoint packing, Marlin + # repacks weights and scales, and the megakernels take the checkpoint + # tensors as they are. self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() - self.experts_cls: type[mk.FusedMoEExperts] - if moe.moe_backend == "b12x": + if moe.moe_backend == "b12x" or is_flashinfer_moe_ep_backend(moe.moe_backend): self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe) assert experts_cls is not None self.experts_cls = experts_cls @@ -70,6 +77,10 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): logger.info_once("Using MarlinExperts for MXFP4 MoE") self.experts_cls = MarlinExperts + @property + def supports_eplb(self) -> bool: + return self.mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS + def create_weights( self, layer: torch.nn.Module, @@ -198,6 +209,21 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight_scale = torch.nn.Parameter( torch.stack(swizzled_w2), requires_grad=False ) + elif self.mxfp4_backend in FLASHINFER_MOE_EP_MXFP4_BACKENDS: + w13, w2, w13_scale, w2_scale, _, _ = ( + convert_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + ) + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu(): pass else: @@ -211,6 +237,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): self.moe_quant_config = self.get_fused_moe_quant_config(layer) if self.moe_quant_config is not None: + assert self.experts_cls is not None self.moe_kernel = make_mxfp4_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 51b7d0460a..dff0c62669 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -38,6 +38,7 @@ _EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset( NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_MOE_EP_CUTEDSL, } ) @@ -61,7 +62,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): weight_key=kNvfp4Static, activation_key=None if use_a16 else kNvfp4Dynamic, ) - self.use_global_sf = is_global_sf_supported_for_nvfp4_backend( self.nvfp4_backend ) @@ -155,7 +155,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): set_weight_attrs(w13_weight_scale_2, extra_weight_attrs) w2_weight_scale_2 = torch.nn.Parameter( - torch.empty(num_experts, dtype=torch.float32), requires_grad=False + torch.empty(num_experts, dtype=torch.float32), + requires_grad=False, ) layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2) extra_weight_attrs.update( @@ -175,7 +176,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): set_weight_attrs(w13_input_scale, extra_weight_attrs) w2_input_scale = torch.nn.Parameter( - torch.empty(num_experts, dtype=torch.float32), requires_grad=False + torch.empty(num_experts, dtype=torch.float32), + requires_grad=False, ) layer.register_parameter("w2_input_global_scale", w2_input_scale) extra_weight_attrs.update( @@ -293,6 +295,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py index 0abae681f0..f3952afe0d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_int8.py @@ -287,6 +287,7 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py index 82f3ada446..c20b217010 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -366,6 +366,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py index 1f6e87a672..75dfc5fa8f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py @@ -291,4 +291,5 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py index b5326e1875..90f7cb2cf3 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py @@ -179,6 +179,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 8bbaaecd64..fdd37d1203 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -528,6 +528,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index 37bb9c6840..1c1d18614e 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -818,6 +818,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py index 53e8ca2b01..9db2907a26 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -4,6 +4,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoeWeightScaleSupported, @@ -53,6 +54,11 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase): def __init__(self, moe) -> None: super().__init__(moe) + if moe.moe_backend in FLASHINFER_MOE_EP_BACKENDS: + raise ValueError( + f"moe_backend={moe.moe_backend!r} is not supported for AutoRound " + "(INC) MXFP4 MoE checkpoints" + ) self.group_size = 32 # Backend selection must stay consistent with the weight preparation in # process_weights_after_loading / get_fused_moe_quant_config, which use diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index 161baf446d..52cfb249dd 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -201,7 +201,7 @@ class QuantFP8(CustomOp): if scale is None: if self.group_shape == GroupShape.PER_TOKEN: - x_max, _ = x.abs().max(dim=-1) + x_max = x.abs().amax(dim=-1) x_max = x_max.unsqueeze(-1).to(torch.float32) if scale_ub is not None: x_max = x_max.clamp(max=scale_ub) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 90ac2b8c79..272eb457bb 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -680,6 +680,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( @@ -1092,6 +1093,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( @@ -1461,6 +1463,7 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( @@ -2314,9 +2317,9 @@ def select_linear_kernel( w = spec.weight assert isinstance(w, QuantKey), f"resolve() must supply a weight key, got {w!r}" if w.dtype == FP4_DTYPE: - # W4A16 (activation is None) → use_a16=True defaults to Marlin *and* - # honors --linear-backend (matches upstream ModelOptNvFp4W4A16LinearMethod - # after #50273); W4A4 → use_a16=False. + # W4A16 (activation is None) → use_a16=True, auto prefers FlashInfer + # CuTe-DSL backend on SM100/103 when available and Marlin otherwise, + # and honors --linear-backend (#50273); W4A4 → use_a16=False. return init_nvfp4_linear_kernel(use_a16=spec.activation is None) if w.scale.dtype == MXFP8_SCALE_DTYPE: return init_mxfp8_linear_kernel( diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index a8b97abc86..1faf2a88e8 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -531,6 +531,7 @@ class MoeWNA16Method(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) @staticmethod diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index f9c7f37ff3..69ddd0045b 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + import torch from vllm.logger import init_logger @@ -21,6 +23,7 @@ from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput 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, @@ -472,9 +475,67 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) +def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool: + """Route Kimi-K3 weight-only MXFP4 SiTU MoE to AITER A16W4 on gfx950.""" + from vllm.platforms import current_platform + + 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 + ) + + +def _use_k3_situ_int4_gfx942(moe: FusedMoEConfig) -> bool: + """Opt-in gfx942 MXFP4-to-int4 requant for Kimi-K3 SiTU experts.""" + from vllm.platforms import current_platform + + 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_gfx942 + + return ( + rocm_aiter_ops.is_fused_moe_enabled() + and on_gfx942() + and moe.activation == MoEActivation.SITU + and moe.activation_situ_linear_beta is not None + and _moe_weight_override_is_int4() + ) + + +def _moe_weight_override_is_int4() -> bool: + """True when --quantization-config.moe.weight int4_per_group_32 was set.""" + from vllm.config import get_current_vllm_config + + vllm_config = get_current_vllm_config() + if vllm_config is None: + return False + quant_args = getattr(vllm_config.model_config, "quantization_config", None) + moe_spec = getattr(quant_args, "moe", None) + weight = getattr(moe_spec, "weight", None) + if weight is None: + return False + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kInt4Static32, + ) + + return weight == kInt4Static32 + + class Mxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" @@ -482,9 +543,33 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): def __init__(self, moe: FusedMoEConfig): super().__init__(moe) - self.weight_dtype = "mxfp4" - self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe) + self.is_k3_situ_aiter = _use_k3_situ_aiter(moe) + self.is_k3_situ_int4_gfx942 = _use_k3_situ_int4_gfx942(moe) + self.experts_cls: type[mk.FusedMoEExperts] | None + if self.is_k3_situ_aiter or self.is_k3_situ_int4_gfx942: + self.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 + self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0] + if self.is_k3_situ_int4_gfx942: + logger.warning_once( + "Requantizing Kimi-K3 MXFP4 experts to groupwise int4 " + "for gfx942. This conversion is lossy and was explicitly " + "enabled by --quantization-config.moe.weight " + "int4_per_group_32." + ) + else: + 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_enabled(): + # a16w4 does not read this bound. Force 0 so leftover a8w4 + # config cannot keep bf16 activations. + # 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.max_capture_size = moe.max_capture_size @@ -782,7 +867,142 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ) self.moe_kernel.fused_experts.process_weights_after_loading(layer) - def process_weights_after_loading(self, layer): + def _setup_kernel_k3_situ_gfx942(self, layer: RoutedExperts) -> None: + import inspect + + from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import ( + flydsl_a16w4_gemm1, + ) + + # Fail closed if this AITER build would silently run SiLU. + if "act" not in inspect.signature(flydsl_a16w4_gemm1).parameters: + raise RuntimeError( + "This AITER build ignores the SiTUv2 activation on the " + "replacement packed-int4 MoE path and would silently compute " + "SiLU. Rebuild with the replacement A16W4/int4 pipeline." + ) + + from aiter import dtypes as aiter_dtypes + from aiter.ops.quant import per_1x32_i4_quant + from aiter.ops.shuffle import ( + pack_int8_to_packed_int4, + shuffle_scale_for_int4, + shuffle_weight, + ) + from aiter.utility import fp4_utils + + fp4_dtype = torch.float4_e2m1fn_x2 + e8m0_dtype = torch.float8_e8m0fnu + + # mxfp4_to_f32 grows the packed weight 8x and peaks above 20 GiB per + # rank if all experts convert at once. Convert 8 experts at a time. + _CONVERT_CHUNK = 8 + + def convert( + weight: torch.nn.Parameter, + scale: torch.nn.Parameter, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Drop the source before returning. Packed int4 is the same size as + # packed MXFP4, so holding both doubles the footprint. That fits + # under expert parallel (1/ep_size experts per rank) and does not + # fit under pure tensor parallel. + w_all = weight.data.view(fp4_dtype) + s_all = scale.data.view(e8m0_dtype) + num_experts = w_all.shape[0] + + # Write chunks into a preallocated output. torch.cat of a list + # would hold the full result twice at the join. + out_packed: torch.Tensor | None = None + out_scale: torch.Tensor | None = None + scale_stride = 0 + for lo in range(0, num_experts, _CONVERT_CHUNK): + hi = min(lo + _CONVERT_CHUNK, num_experts) + weight_f32 = fp4_utils.mxfp4_to_f32(w_all[lo:hi]) + scale_f32 = fp4_utils.e8m0_to_f32(s_all[lo:hi]) + chunk_experts, output_size, input_size = weight_f32.shape + weight_bf16 = ( + ( + weight_f32.view( + chunk_experts, output_size, input_size // 32, 32 + ) + * scale_f32.view( + chunk_experts, output_size, input_size // 32, 1 + ) + ) + .view(chunk_experts, output_size, input_size) + .to(torch.bfloat16) + ) + del weight_f32, scale_f32 + + weight_int4, weight_scale = per_1x32_i4_quant(weight_bf16) + del weight_bf16 + weight_int4 = weight_int4.view(aiter_dtypes.i4x2).view( + chunk_experts, output_size, input_size + ) + weight_packed = pack_int8_to_packed_int4( + shuffle_weight(weight_int4.view(aiter_dtypes.i8), (16, 16)) + ) + del weight_int4 + chunk_packed = weight_packed.view( + chunk_experts, output_size, input_size // 2 + ).view(aiter_dtypes.i4x2) + chunk_scale = ( + shuffle_scale_for_int4(weight_scale, group_size=32) + .view(-1) + .contiguous() + ) + if out_packed is None: + out_packed = torch.empty( + (num_experts,) + tuple(chunk_packed.shape[1:]), + dtype=chunk_packed.dtype, + device=chunk_packed.device, + ) + scale_stride = chunk_scale.numel() // chunk_experts + out_scale = torch.empty( + num_experts * scale_stride, + dtype=chunk_scale.dtype, + device=chunk_scale.device, + ) + assert out_packed is not None + assert out_scale is not None + out_packed[lo:hi].copy_(chunk_packed) + out_scale[lo * scale_stride : hi * scale_stride].copy_(chunk_scale) + del weight_packed, weight_scale, chunk_packed, chunk_scale + torch.accelerator.empty_cache() + + del w_all, s_all + assert out_packed is not None + assert out_scale is not None + weight.data = torch.empty(0, dtype=torch.uint8, device=out_packed.device) + scale.data = torch.empty(0, dtype=torch.uint8, device=out_packed.device) + torch.accelerator.empty_cache() + return out_packed, out_scale + + w13, w13_scale = convert(layer.w13_weight, layer.w13_weight_scale) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w13_weight_scale", w13_scale) + del w13, w13_scale + torch.accelerator.empty_cache() + + w2, w2_scale = convert(layer.w2_weight, layer.w2_weight_scale) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w2_weight_scale", w2_scale) + del w2, w2_scale + torch.accelerator.empty_cache() + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None and self.experts_cls is not None: + self.moe_kernel = make_mxfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + mxfp4_backend=self.mxfp4_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: if self.mxfp4_backend == Mxfp4MoeBackend.NONE: return @@ -795,6 +1015,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self._build_moe_kernel(layer) return + if self.is_k3_situ_int4_gfx942: + self._setup_kernel_k3_situ_gfx942(layer) + return + w13 = layer.w13_weight w2 = layer.w2_weight w13_scale = layer.w13_weight_scale @@ -900,4 +1124,5 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 41bca8d29d..8ed48b2b64 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -65,6 +65,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static32, kInt8StaticChannelSym, kMxfp4Static, kMxfp8Dynamic, @@ -231,6 +232,9 @@ class OnlineQuantizationConfig(QuantizationConfig): """ if spec is None or spec.weight is None: return None + # Load-time gfx942 requant, not online conversion. Mxfp4MoEMethod owns it. + if spec.weight == kInt4Static32: + return None cls = table.get(spec.weight) if cls is None: raise ValueError( diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 414c4aba5b..6e21bc9092 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -156,6 +156,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): topk_group=layer.topk_group, e_score_correction_bias=layer.e_score_correction_bias, routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) def apply( diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 8da0aba97c..262f9bdb0a 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1976,6 +1976,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + routing_sink=layer.routing_sink, ) diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=1536,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=1536,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..85a2aa9d83 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=1536,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=2048,K=11008,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=2048,K=11008,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..a8bc891dd5 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=2048,K=11008,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=2048,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=2048,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..072db2cbc7 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=2048,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=2112,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=2112,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..8b535371f0 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=2112,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=22016,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=22016,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..323db388eb --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=22016,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..5ce80c0316 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..e1a348c85d --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=24576,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=28672,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=28672,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..eeada919b2 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=28672,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=1536,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=1536,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..9d16cde36f --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=1536,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..5ed9e9d441 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..d35330be1b --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=3072,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=32768,K=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=32768,K=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..dc895d5a29 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=32768,K=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=12288,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=12288,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..ff4536c455 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=12288,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=14336,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=14336,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..0ef06c44e6 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=14336,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..f8b9e8c696 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..fd0318f8b5 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=4096,K=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=4608,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=4608,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..0087fa0de0 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=4608,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=512,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=512,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..d5af629a01 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=512,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=576,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=576,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..bafaf2ad8b --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=576,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=576,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=576,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..998adaa968 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=576,K=7168,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=6144,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=6144,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..6153e4b540 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=6144,K=4096,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=16384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=16384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..a70cd56ba4 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=16384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=18432,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=18432,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..8cc0253e39 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=18432,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..793df30860 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=2048,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=2304,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=2304,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..12246c9700 --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=2304,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 2 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 0000000000..d529d1631a --- /dev/null +++ b/vllm/model_executor/layers/quantization/utils/configs/N=7168,K=256,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,82 @@ +{ + "1": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 16, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 32, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 4c0a8a3ff7..03d4df386a 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -25,6 +25,11 @@ INT4_DTYPE = scalar_types.uint4b8 INT8_DTYPE = scalar_types.uint8b128 +def require_finite_positive_scales(name: str, value: torch.Tensor) -> None: + if not torch.isfinite(value).all() or not (value > 0).all(): + raise ValueError(f"{name} must contain finite positive scales") + + def _dtype_abbr(dtype: torch.dtype | ScalarType) -> str: """Return a stable short name for torch and ScalarType dtypes.""" if isinstance(dtype, ScalarType): diff --git a/vllm/model_executor/models/deepseek_ocr.py b/vllm/model_executor/models/deepseek_ocr.py index 62a46fc498..3aef695d06 100644 --- a/vllm/model_executor/models/deepseek_ocr.py +++ b/vllm/model_executor/models/deepseek_ocr.py @@ -457,8 +457,6 @@ class DeepseekOCRForCausalLM( return None assert isinstance(pixel_values, torch.Tensor) assert images_crop is None or isinstance(images_crop, torch.Tensor) - if torch.sum(pixel_values).item() == 0: - return None # Use actual tensor spatial dim instead of hardcoded # vision_config.image_size (1024). The vision encoders (SAM & CLIP) diff --git a/vllm/model_executor/models/deepseek_ocr2.py b/vllm/model_executor/models/deepseek_ocr2.py index 938598a68b..18de7503c7 100644 --- a/vllm/model_executor/models/deepseek_ocr2.py +++ b/vllm/model_executor/models/deepseek_ocr2.py @@ -318,7 +318,7 @@ class DeepseekOCR2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Support images_spatial_crop = kwargs.pop("images_spatial_crop", None) images_crop = kwargs.pop("images_crop", None) - if pixel_values is None or torch.sum(pixel_values).item() == 0: + if pixel_values is None: return None base_size = self.vision_config.image_size diff --git a/vllm/model_executor/models/gpt_neox.py b/vllm/model_executor/models/gpt_neox.py deleted file mode 100644 index 6e84eebe51..0000000000 --- a/vllm/model_executor/models/gpt_neox.py +++ /dev/null @@ -1,324 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gpt_neox/modeling_gpt_neox.py -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Inference-only GPT-NeoX model compatible with HuggingFace weights.""" - -from collections.abc import Iterable -from itertools import islice - -import torch -from torch import nn -from transformers import GPTNeoXConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import get_act_fn -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - 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.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors - -from .interfaces import SupportsPP -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class GPTNeoXAttention(nn.Module): - def __init__( - self, - config: GPTNeoXConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.total_num_heads = config.num_attention_heads - self.hidden_size = config.hidden_size - self.head_size = self.hidden_size // self.total_num_heads - self.bias = getattr(config, "attention_bias", True) - - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() - assert self.total_num_heads % tensor_model_parallel_world_size == 0 - self.num_heads = self.total_num_heads // tensor_model_parallel_world_size - - self.query_key_value = QKVParallelLinear( - config.hidden_size, - self.head_size, - self.total_num_heads, - bias=self.bias, - quant_config=quant_config, - prefix=f"{prefix}.query_key_value", - ) - self.dense = RowParallelLinear( - config.hidden_size, - config.hidden_size, - bias=self.bias, - quant_config=quant_config, - prefix=f"{prefix}.dense", - ) - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - self.rotary_emb = get_rope( - self.head_size, - max_position=max_position_embeddings, - rope_parameters=config.rope_parameters, - ) - scaling = self.head_size**-0.5 - self.attn = Attention( - self.num_heads, - self.head_size, - scaling, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - qkv, _ = self.query_key_value(hidden_states) - q, k, v = qkv.chunk(chunks=3, dim=-1) - q, k = self.rotary_emb(position_ids, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.dense(attn_output) - return output - - -class GPTNeoXMLP(nn.Module): - def __init__( - self, - config: GPTNeoXConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.dense_h_to_4h = ColumnParallelLinear( - config.hidden_size, - config.intermediate_size, - quant_config=quant_config, - prefix=f"{prefix}.dense_h_to_4h", - ) - self.dense_4h_to_h = RowParallelLinear( - config.intermediate_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.dense_4h_to_h", - ) - self.act = get_act_fn(config.hidden_act) - - def forward(self, hidden_states): - hidden_states, _ = self.dense_h_to_4h(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states, _ = self.dense_4h_to_h(hidden_states) - return hidden_states - - -class GPTNeoXLayer(nn.Module): - def __init__( - self, - config: GPTNeoXConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.use_parallel_residual = config.use_parallel_residual - self.input_layernorm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.post_attention_layernorm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.attention = GPTNeoXAttention( - config, cache_config, quant_config, prefix=f"{prefix}.attention" - ) - self.mlp = GPTNeoXMLP(config, quant_config, prefix=f"{prefix}.mlp") - - def forward( - self, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - attn_input = self.input_layernorm(hidden_states) - attn_output = self.attention( - position_ids=position_ids, - hidden_states=attn_input, - ) - - if self.use_parallel_residual: - # pseudocode: - # x = x + attn(ln1(x)) + mlp(ln2(x)) - mlp_input = self.post_attention_layernorm(hidden_states) - mlp_output = self.mlp(mlp_input) - hidden_states = mlp_output + attn_output + hidden_states - else: - # pseudocode: - # x = x + attn(ln1(x)) - # x = x + mlp(ln2(x)) - attn_output = attn_output + hidden_states - mlp_input = self.post_attention_layernorm(attn_output) - mlp_output = self.mlp(mlp_input) - hidden_states = mlp_output + attn_output - return hidden_states - - -@support_torch_compile -class GPTNeoXModel(nn.Module): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_substr={"attention.bias": None, "attention.masked_bias": None} - ) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - self.config = config - - self.embed_in = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - ) - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: GPTNeoXLayer( - config, cache_config, quant_config, prefix=prefix - ), - prefix=f"{prefix}.layers", - ) - self.final_layer_norm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states"], config.hidden_size - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_in(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - position_ids: 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) - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states = layer(position_ids, hidden_states) - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) - hidden_states = self.final_layer_norm(hidden_states) - return hidden_states - - def _repack_qkv( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> Iterable[tuple[str, torch.Tensor]]: - # GPT-NeoX's fused QKV is laid out as (num_heads * 3 * head_size) on - # its output dim (0), while vLLM expects (3 * num_heads * head_size). - num_heads = self.config.num_attention_heads - for name, loaded_weight in weights: - if "query_key_value" in name: - shape = loaded_weight.shape - loaded_weight = loaded_weight.view((num_heads, 3, -1) + shape[1:]) - loaded_weight = loaded_weight.transpose(0, 1) - loaded_weight = loaded_weight.reshape(shape) - yield name, loaded_weight - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights( - self._repack_qkv(weights), mapper=self.hf_to_vllm_mapper - ) - - -class GPTNeoXForCausalLM(nn.Module, SupportsPP): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.quant_config = quant_config - self.gpt_neox = GPTNeoXModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "gpt_neox") - ) - self.embed_out = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "embed_out"), - ) - if self.config.tie_word_embeddings: - 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 - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.gpt_neox.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.gpt_neox( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.embed_out, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/jais2.py b/vllm/model_executor/models/jais2.py deleted file mode 100644 index 84254e1ee5..0000000000 --- a/vllm/model_executor/models/jais2.py +++ /dev/null @@ -1,441 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Inference-only Jais2 model compatible with HuggingFace weights.""" - -from collections.abc import Iterable - -import torch -from torch import nn -from transformers import Jais2Config - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import ( - get_pp_group, - get_tensor_model_parallel_world_size, -) -from vllm.model_executor.layers.activation import ReLUSquaredActivation -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - 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.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors - -from .interfaces import SupportsLoRA, SupportsPP -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - WeightsMapper, - extract_layer_index, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class Jais2MLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: QuantizationConfig | None = None, - bias: bool = False, - prefix: str = "", - ) -> None: - super().__init__() - self.up_proj = ColumnParallelLinear( - input_size=hidden_size, - output_size=intermediate_size, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.up_proj", - ) - self.down_proj = RowParallelLinear( - input_size=intermediate_size, - output_size=hidden_size, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - self.act_fn = ReLUSquaredActivation() - - def forward(self, x): - x, _ = self.up_proj(x) - x = self.act_fn(x) - x, _ = self.down_proj(x) - return x - - -class Jais2Attention(nn.Module): - def __init__( - self, - config: Jais2Config, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - max_position_embeddings: int = 8192, - quant_config: QuantizationConfig | None = None, - bias: bool = False, - cache_config: CacheConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - layer_idx = extract_layer_index(prefix) - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - # MistralConfig has an optional head_dim introduced by Mistral-Nemo - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self.max_position_embeddings = max_position_embeddings - - self.qkv_proj = QKVParallelLinear( - hidden_size=hidden_size, - head_size=self.head_dim, - total_num_heads=self.total_num_heads, - total_num_kv_heads=self.total_num_kv_heads, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - - self.o_proj = RowParallelLinear( - input_size=self.total_num_heads * self.head_dim, - output_size=hidden_size, - bias=bias, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - is_neox_style = True - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position_embeddings, - rope_parameters=getattr(config, "rope_parameters", None), - is_neox_style=is_neox_style, - ) - - if hasattr(config, "interleaved_sliding_window"): - interleaved_sliding_window = config.interleaved_sliding_window - if isinstance(interleaved_sliding_window, int): - sliding_window = interleaved_sliding_window - elif isinstance(interleaved_sliding_window, list): - sw_idx = layer_idx % len(interleaved_sliding_window) - sliding_window = interleaved_sliding_window[sw_idx] - else: - raise ValueError( - f"{type(interleaved_sliding_window)} is not supported." - ) - else: - sliding_window = None - - 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, - 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) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class Jais2DecoderLayer(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - config: Jais2Config, - prefix: str = "", - ) -> None: - super().__init__() - - config = config or vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = self.get_quant_config(vllm_config) - - self.hidden_size = config.hidden_size - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - # Support abacusai/Smaug-72B-v0.1 with attention_bias - attention_bias = getattr(config, "attention_bias", False) or getattr( - config, "bias", False - ) - self.self_attn = Jais2Attention( - config=config, - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - num_kv_heads=getattr( - config, "num_key_value_heads", config.num_attention_heads - ), - max_position_embeddings=max_position_embeddings, - quant_config=quant_config, - bias=attention_bias, - cache_config=cache_config, - prefix=f"{prefix}.self_attn", - ) - self.mlp = Jais2MLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - bias=getattr(config, "mlp_bias", False), - prefix=f"{prefix}.mlp", - ) - self.input_layernorm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.post_attention_layernorm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = ( - self.input_layernorm(hidden_states + residual), - hidden_states + residual, - ) - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - ) - - # Fully Connected - hidden_states, residual = ( - self.post_attention_layernorm(hidden_states + residual), - hidden_states + residual, - ) - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None: - """Get quantization config for this layer. Override in subclasses.""" - return vllm_config.quant_config - - -@support_torch_compile -class Jais2Model(nn.Module): - def __init__( - self, - vllm_config: VllmConfig, - prefix: str = "", - layer_type: type[nn.Module] = Jais2DecoderLayer, - ): - super().__init__() - - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - self.quant_config = quant_config - - self.vocab_size = config.vocab_size - self.org_vocab_size = config.vocab_size - if get_pp_group().is_first_rank or ( - config.tie_word_embeddings and get_pp_group().is_last_rank - ): - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - org_num_embeddings=config.vocab_size, - quant_config=quant_config, - ) - else: - self.embed_tokens = PPMissingLayer() - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: layer_type( - config=config, - vllm_config=vllm_config, - prefix=prefix, - ), - prefix=f"{prefix}.layers", - ) - if get_pp_group().is_last_rank: - self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) - else: - self.norm = PPMissingLayer() - - 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_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 | tuple[torch.Tensor, list[torch.Tensor]]: - 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"] - - for i in range(self.start_layer, self.end_layer): - layer = self.layers[i] - hidden_states, residual = layer(positions, hidden_states, residual) - - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - - hidden_states, _ = self.norm(hidden_states + residual), residual - return hidden_states - - -class Jais2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - } - ) - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - } - - embedding_modules = { - "embed_tokens": "input_embeddings", - "lm_head": "output_embeddings", - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.model = self._init_model( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - 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.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 - ) - else: - self.lm_head = PPMissingLayer() - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def _init_model(self, vllm_config: VllmConfig, prefix: str = ""): - return Jais2Model(vllm_config=vllm_config, prefix=prefix) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - model_output = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return model_output - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - 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/mistral3.py b/vllm/model_executor/models/mistral3.py index 40af706e3a..f5ffe32aad 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -7,8 +7,6 @@ from typing import Annotated, Any, Literal, cast import torch import torch.nn as nn -import transformers -from packaging.version import Version from transformers import BatchFeature, Mistral3Config, PixtralVisionConfig from transformers.models.pixtral import PixtralProcessor @@ -59,10 +57,6 @@ from .utils import ( maybe_prefix, ) -TRANSFORMERS_SUPPORTS_PIXTRAL_IMAGE_ONLY = Version(transformers.__version__) >= Version( - "5.15.0" -) - class Mistral3ImagePixelInputs(TensorSchema): """Dimensions: @@ -284,13 +278,6 @@ class Mistral3DummyInputsBuilder(BaseDummyInputsBuilder): class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]): - def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str | None: - # PixtralProcessor supports image-only calls starting in transformers 5.15. - # Older releases need dummy text and take the slower tokenization path. - if not TRANSFORMERS_SUPPORTS_PIXTRAL_IMAGE_ONLY: - return self.dummy_inputs.get_dummy_text(mm_counts) - return None - def _postprocess_hf_mm_data( self, mm_data: Mapping[str, object], diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py index 3ba731d2e0..9ac38c0021 100644 --- a/vllm/model_executor/models/nemotron_parse.py +++ b/vllm/model_executor/models/nemotron_parse.py @@ -568,15 +568,6 @@ class NemotronParseForConditionalGeneration(nn.Module, SupportsMultiModal): self.lm_head = ParallelLMHead( config.decoder.vocab_size, config.decoder.d_model, quant_config=quant_config ) - # Some checkpoints (e.g. compact exports) tie the output head to the - # decoder's input embeddings instead of materializing a separate - # lm_head.weight tensor. - tie_word_embeddings = bool( - getattr(config, "tie_word_embeddings", False) - or getattr(config.decoder, "tie_word_embeddings", False) - ) - if tie_word_embeddings: - self.lm_head = self.lm_head.tie_weights(self.decoder.embed_tokens) self.logits_processor = LogitsProcessor( self.vocab_size, config.decoder.vocab_size ) diff --git a/vllm/model_executor/models/phi.py b/vllm/model_executor/models/phi.py deleted file mode 100644 index 61c243aadf..0000000000 --- a/vllm/model_executor/models/phi.py +++ /dev/null @@ -1,325 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from -# https://huggingface.co/microsoft/phi-1_5/blob/main/modeling_phi.py -# Copyright 2023 The vLLM team. -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. -# -# BSD 3-Clause License -# -# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu. -# All rights reserved. -# -# Redistribution and use in source and binary forms, with or without -# modification, are permitted provided that the following conditions are met: -# -# * Redistributions of source code must retain the above copyright notice, this -# list of conditions and the following disclaimer. -# -# * Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation -# and/or other materials provided with the distribution. -# -# * Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from -# this software without specific prior written permission. -# -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE -# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR -# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER -# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, -# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -"""Inference-only Phi-1.5 model compatible with HuggingFace weights.""" - -from collections.abc import Iterable -from itertools import islice - -import torch -from torch import nn -from transformers import PhiConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import get_act_fn -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - 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.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors - -from .interfaces import SupportsLoRA, SupportsPP -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class PhiAttention(nn.Module): - def __init__( - self, - config: PhiConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.hidden_size = config.hidden_size - self.head_size = self.hidden_size // config.num_attention_heads - - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() - assert config.num_attention_heads % tensor_model_parallel_world_size == 0 - self.num_heads = config.num_attention_heads // tensor_model_parallel_world_size - - # pylint: disable=C0103 - self.qkv_proj = QKVParallelLinear( - self.hidden_size, - self.head_size, - config.num_attention_heads, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.dense = RowParallelLinear( - self.hidden_size, - self.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.dense", - ) - - scaling = self.head_size**-0.5 - - max_position_embeddings = getattr(config, "max_position_embeddings", 2048) - self.rotary_emb = get_rope( - self.head_size, - max_position=max_position_embeddings, - rope_parameters=config.rope_parameters, - ) - self.attn = Attention( - self.num_heads, - self.head_size, - scaling, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - - def forward( - self, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.chunk(chunks=3, dim=-1) - q, k = self.rotary_emb(position_ids, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.dense(attn_output) - return output - - -class PhiMLP(nn.Module): - def __init__( - self, - config: PhiConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - - n_inner = getattr(config, "n_inner", None) - n_inner = n_inner if n_inner is not None else 4 * config.hidden_size - - self.fc1 = ColumnParallelLinear( - config.hidden_size, - n_inner, - quant_config=quant_config, - prefix=f"{prefix}.fc1", - ) - self.fc2 = RowParallelLinear( - n_inner, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.fc2", - ) - self.act = get_act_fn(config.hidden_act) - - def forward(self, hidden_states): - hidden_states, _ = self.fc1(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states, _ = self.fc2(hidden_states) - return hidden_states - - -class PhiLayer(nn.Module): - def __init__( - self, - config: PhiConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.input_layernorm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.self_attn = PhiAttention( - config, cache_config, quant_config, prefix=f"{prefix}.self_attn" - ) - self.mlp = PhiMLP(config, quant_config, prefix=f"{prefix}.mlp") - - def forward( - self, - position_ids: torch.Tensor, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - attn_outputs = self.self_attn( - position_ids=position_ids, - hidden_states=hidden_states, - ) - feed_forward_hidden_states = self.mlp(hidden_states) - hidden_states = attn_outputs + feed_forward_hidden_states + residual - return hidden_states - - -@support_torch_compile -class PhiModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - 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 - ) - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: PhiLayer(config, cache_config, quant_config, prefix=prefix), - prefix=f"{prefix}.layers", - ) - self.final_layernorm = nn.LayerNorm( - config.hidden_size, eps=config.layer_norm_eps - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states"], config.hidden_size - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return 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) - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states = layer(positions, hidden_states) - - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) - - hidden_states = self.final_layernorm(hidden_states) - - return hidden_states - - -class PhiForCausalLM(nn.Module, SupportsLoRA, SupportsPP): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - } - ) - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - # lm_head use bias, cannot share word embeddings - assert not config.tie_word_embeddings - - self.quant_config = quant_config - - self.model = PhiModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - bias=True, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states, self.lm_head.bias) - 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/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index f9338b1189..d3510f3e64 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -2260,10 +2260,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) st_idx += text_len - bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(bos_pos) - st_idx += 1 - if modality == "audio": audio_tokens = self._compute_audio_token_count( data["audio_feature_length"] @@ -2276,7 +2272,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(eos_pos) - st = offset + 1 + audio_tokens + 1 + st = offset + audio_tokens + 1 elif modality == "image": grid_t = data["grid_t"] @@ -2294,7 +2290,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(eos_pos) - st = offset + 1 + image_len + 1 + st = offset + image_len + 1 elif modality == "video": grid_t = data["grid_t"] @@ -2313,7 +2309,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(eos_pos) - st = offset + 1 + video_len + 1 + st = offset + video_len + 1 else: audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1)) llm_pos_ids_list.append(audio_bos_pos) @@ -2330,7 +2326,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( audio_len = self._compute_audio_token_count( data["audio_feature_length"] ) - st = offset + 2 + video_len + audio_len + 2 + st = offset + 1 + video_len + audio_len + 2 if st < seq_len: st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 81763fed40..3ff37a2562 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -120,7 +120,6 @@ _TEXT_GENERATION_MODELS = { "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"), "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"), "GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"), - "GPTNeoXForCausalLM": ("gpt_neox", "GPTNeoXForCausalLM"), "GraniteForCausalLM": ("granite", "GraniteForCausalLM"), "GraniteMoeForCausalLM": ("granitemoe", "GraniteMoeForCausalLM"), "GraniteMoeHybridForCausalLM": ("granitemoehybrid", "GraniteMoeHybridForCausalLM"), @@ -136,7 +135,6 @@ _TEXT_GENERATION_MODELS = { "InternLM3ForCausalLM": ("llama", "LlamaForCausalLM"), "IQuestCoderForCausalLM": ("llama", "LlamaForCausalLM"), "IQuestLoopCoderForCausalLM": ("iquest_loopcoder", "IQuestLoopCoderForCausalLM"), - "Jais2ForCausalLM": ("jais2", "Jais2ForCausalLM"), "JambaForCausalLM": ("jamba", "JambaForCausalLM"), "K2HorizonForCausalLM": ("k2_horizon", "K2HorizonForCausalLM"), "KimiLinearForCausalLM": ( @@ -188,7 +186,6 @@ _TEXT_GENERATION_MODELS = { "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"), "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"), "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"), - "PhiForCausalLM": ("phi", "PhiForCausalLM"), "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"), "PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"), "Plamo3ForCausalLM": ("plamo3", "Plamo3ForCausalLM"), @@ -200,7 +197,6 @@ _TEXT_GENERATION_MODELS = { "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"), "SarvamMoEForCausalLM": ("sarvam", "SarvamMoEForCausalLM"), "SarvamMLAForCausalLM": ("sarvam", "SarvamMLAForCausalLM"), - "SeedOssForCausalLM": ("seed_oss", "SeedOssForCausalLM"), "Step1ForCausalLM": ("step1", "Step1ForCausalLM"), "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"), "Step3p5ForCausalLM": ("step3p5", "Step3p5ForCausalLM"), @@ -709,13 +705,17 @@ _TRANSFORMERS_SUPPORTED_MODELS = { "FlexOlmoForCausalLM": ("transformers", "TransformersMoEForCausalLM"), "GlmForCausalLM": ("transformers", "TransformersForCausalLM"), "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"), + "GPTNeoXForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"), "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"), + "Jais2ForCausalLM": ("transformers", "TransformersForCausalLM"), "MellumForCausalLM": ("transformers", "TransformersMoEForCausalLM"), "NanbeigeForCausalLM": ("transformers", "TransformersForCausalLM"), "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"), "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"), + "PhiForCausalLM": ("transformers", "TransformersForCausalLM"), + "SeedOssForCausalLM": ("transformers", "TransformersForCausalLM"), "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"), "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"), "VaultGemmaForCausalLM": ("transformers", "TransformersForCausalLM"), diff --git a/vllm/model_executor/models/seed_oss.py b/vllm/model_executor/models/seed_oss.py deleted file mode 100644 index 976a35a78a..0000000000 --- a/vllm/model_executor/models/seed_oss.py +++ /dev/null @@ -1,432 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Copyright 2025 The Seed team. -# Copyright 2023 The vLLM team. -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Inference-only SeedOss model compatible with HuggingFace weights.""" - -from collections.abc import Iterable -from itertools import islice - -import torch -from torch import nn -from transformers import PreTrainedConfig as SeedOssConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.logger import init_logger -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - 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.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.config import set_default_rope_theta -from vllm.v1.attention.backend import AttentionType - -from .interfaces import SupportsLoRA, SupportsPP -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - WeightsMapper, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - -logger = init_logger(__name__) - - -class SeedOssMLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: 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_act != "silu": - raise ValueError( - f"Unsupported activation: {hidden_act}. Only silu is supported for now." - ) - self.act_fn = SiluAndMul() - - def forward(self, x): - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class SeedOssAttention(nn.Module): - def __init__( - self, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - rope_parameters: dict, - max_position: int = 4096 * 32, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - attn_type: str = AttentionType.DECODER, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - self.head_dim = head_dim - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - self.rotary_emb = get_rope( - self.head_dim, - max_position=max_position, - rope_parameters=rope_parameters, - ) - 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, - attn_type=attn_type, - 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) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class SeedOssDecoderLayer(nn.Module): - def __init__( - self, - config: SeedOssConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - set_default_rope_theta(config, default_theta=1000000) - - # By default, SeedOss uses causal attention as it is a - # decoder-only model. - # You can override the HF config with `is_causal=False` to enable - # bidirectional attention, which is used in some embedding models - if getattr(config, "is_causal", True): - attn_type = AttentionType.DECODER - else: - attn_type = AttentionType.ENCODER_ONLY - - self.self_attn = SeedOssAttention( - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - max_position=config.max_position_embeddings, - num_kv_heads=config.num_key_value_heads, - head_dim=config.head_dim, - cache_config=cache_config, - quant_config=quant_config, - rope_parameters=config.rope_parameters, - prefix=f"{prefix}.self_attn", - attn_type=attn_type, - ) - self.mlp = SeedOssMLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = self.input_layernorm(hidden_states, residual) - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - ) - - # Fully Connected - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - -@support_torch_compile( - dynamic_arg_dims={ - "input_ids": 0, - "positions": -1, - "intermediate_tensors": 0, - "inputs_embeds": 0, - } -) -class SeedOssModel(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - prefix: str = "", - decoder_layer_type: type[nn.Module] = SeedOssDecoderLayer, - ): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - # TODO (@robertgshaw2): see if this can be moved out - if cache_config.sliding_window is not None and hasattr( - config, "max_window_layers" - ): - assert config.max_window_layers == config.num_hidden_layers, ( - "Sliding window for some but all layers is not supported. " - "This model uses sliding window but `max_window_layers` = {} " - "is less than `num_hidden_layers` = {}. Please open an issue " - "to discuss this feature.".format( - config.max_window_layers, - config.num_hidden_layers, - ) - ) - - self.config = config - self.quant_config = quant_config - self.vocab_size = config.vocab_size - - if get_pp_group().is_first_rank or ( - config.tie_word_embeddings and get_pp_group().is_last_rank - ): - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=f"{prefix}.embed_tokens", - ) - else: - self.embed_tokens = PPMissingLayer() - - # Use the provided decoder layer type or default to SeedDecoderLayer - decoder_layer_type = decoder_layer_type or SeedOssDecoderLayer - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: decoder_layer_type( - config=config, - cache_config=cache_config, - quant_config=quant_config, - prefix=prefix, - ), - prefix=f"{prefix}.layers", - ) - - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states", "residual"], config.hidden_size - ) - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - else: - self.norm = PPMissingLayer() - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = 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"] - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer( - positions, - hidden_states, - residual, - ) - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - -class SeedOssForCausalLM(nn.Module, SupportsLoRA, SupportsPP): - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".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"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - - self.config = config - - self.quant_config = quant_config - self.model = SeedOssModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - - 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.lm_head.tie_weights(self.model.embed_tokens) - else: - self.lm_head = PPMissingLayer() - - self.logits_processor = LogitsProcessor(config.vocab_size) - - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - 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/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 4e53f1fe18..9f045173f3 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -32,6 +32,7 @@ from torch import nn from transformers import AutoModel, PreTrainedConfig from transformers.conversion_mapping import ( WeightRenaming, + get_checkpoint_conversion_mapping, get_model_conversion_mapping, ) @@ -184,7 +185,7 @@ class Base( # Substitute remaining layers with vLLM's layers as needed self.recursive_replace() # Create attention instances for KV cache allocation - self._create_attention_instances() + self._create_attention_instances(prefix) # Initialize any parameters that have not had their modules replaced self.init_parameters(self.model) @@ -311,7 +312,14 @@ class Base( orig_to_new_renaming: list[WeightRenaming] = [] orig_to_new_regex: dict[re.Pattern, str | None] = {} - for mapping in get_model_conversion_mapping(self.model): + mappings = get_model_conversion_mapping(self.model) + # Ensure layers that are part of the model, but not AutoModel, get their mapping + instantiated = {type(m).__name__ for m in self.model.modules()} + for arch in self.config.architectures or []: + if arch not in instantiated: + mappings.extend(get_checkpoint_conversion_mapping(arch) or []) + + for mapping in mappings: # Handle weights which have been renamed in Transformers if isinstance(mapping, WeightRenaming): orig_to_new_renaming.append(mapping) @@ -579,7 +587,7 @@ class Base( self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked) - def _create_attention_instances(self): + def _create_attention_instances(self, prefix: str = ""): """Create `Attention` instances to inform KV cache allocation.""" text_config = self.text_config attn_cls = self._get_attn_cls() @@ -613,8 +621,8 @@ class Base( f"Layer {layer} does not dispatch through the Transformers " "attention interface and vLLM has no other way to handle it." ) - prefix, attn_fuser = self.attention_fusers[i] - attn_module = self.get_submodule(prefix) + attn_prefix, attn_fuser = self.attention_fusers[i] + attn_module = self.get_submodule(attn_prefix) # `[i]` is the whole-model config unless the checkpoint is # heterogeneous, in which case it is this layer's own geometry. @@ -632,17 +640,18 @@ class Base( scale=scale, cache_config=self.cache_config, quant_config=self.quant_config, - prefix=f"{i}.attn", + prefix=maybe_prefix(prefix, f"{attn_prefix}.{VLLM_ATTN_ATTR}"), ) if attn_cls is MLAAttention: mla_fuser = next( - (f for f in self.fusers[prefix] if isinstance(f, MLAFuser)), None + (f for f in self.fusers[attn_prefix] if isinstance(f, MLAFuser)), + None, ) if mla_fuser is None: raise ValueError( - f"Layer {i} ({prefix}) did not fuse as MLA, but the rest of " - "the model did, so its latent KV cache cannot be allocated." + f"Layer {i} ({attn_prefix}) did not fuse as MLA, but the rest " + "of the model did, so its latent KV cache cannot be allocated." ) dims = get_mla_dims(self.model_config) kwargs.update( @@ -690,10 +699,8 @@ class Base( # vLLM does not support encoder-decoder models, so if any encoder layer is # found in a text only model, we assume the whole model is an encoder model if has_encoder(self.model) and not is_multimodal(self.config): - self.check_version("5.0.0", "encoder models support") return EncoderOnlyAttention if self.model_config.use_mla: - self.check_version("5.15.0.dev0", "optimized MLA support") if any( isinstance(fuser, MLAFuser) for fusers in self.fusers.values() @@ -819,7 +826,6 @@ class Base( ) def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: - self.check_version("5.2.0", "Eagle3 support") from transformers.utils.output_capturing import ( OutputRecorder, maybe_install_capturing_hooks, diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index e7dba82b0f..9124784d36 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -19,6 +19,8 @@ from collections.abc import Iterable from typing import TYPE_CHECKING +import torch + from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -30,8 +32,6 @@ from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix from .base import Base if TYPE_CHECKING: - import torch - from vllm.config import VllmConfig @@ -68,11 +68,15 @@ class CausalMixin(VllmModelForTextGeneration, Base): def load_weights(self, weights: Iterable[tuple[str, "torch.Tensor"]]) -> set[str]: """A thin wrapper around `Base.load_weights` to handle the lm_head bias.""" lm_head_bias = set() + load_device = self.vllm_config.load_config.device or self.device_config.device def auto_load_lm_head_bias(weights): for name, weight in weights: if name.endswith("lm_head.bias") and self.pp_group.is_last_rank: - self.lm_head._register_bias() + # load_weights is called outside vLLM's device context, + # so we must enter it again when registering the bias. + with torch.device(load_device): + self.lm_head._register_bias() self.lm_head.bias.weight_loader(self.lm_head.bias, weight) lm_head_bias.add(name) else: diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 9f63e73336..71f5bd066c 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -156,7 +156,6 @@ class MoEMixin(MixtureOfExperts, Base): num_local_physical_experts: int def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): - self.check_version("5.0.0", "MoE models support") # Skip MixtureOfExperts.__init__ and call the next class in MRO super(MixtureOfExperts, self).__init__(vllm_config=vllm_config, prefix=prefix) diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index cecb4f564a..9887619858 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -23,17 +23,11 @@ from typing import TYPE_CHECKING, Any import torch import transformers -from packaging.version import Version from transformers.utils.generic import ModelOutput from vllm.compilation.decorators import should_torch_compile_mm_encoder from vllm.config.utils import getattr_iter -from vllm.inputs import ( - MultiModalDataBuiltins, - MultiModalDataDict, - MultiModalInput, - mm_input, -) +from vllm.inputs import MultiModalDataBuiltins, MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, @@ -42,26 +36,15 @@ from vllm.model_executor.models.interfaces import ( ) from vllm.model_executor.models.module_mapping import MultiModelKeys from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems -from vllm.multimodal.inputs import ( - MultiModalFeatureSpec, - MultiModalFieldConfig, - PlaceholderRange, -) -from vllm.multimodal.parse import ( - ImageProcessorItems, - ImageSize, - MultiModalDataItems, - MultiModalDataParser, -) +from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalFieldConfig +from vllm.multimodal.parse import ImageSize, MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, BaseMultiModalProcessor, BaseProcessingInfo, - ProcessorInputs, PromptReplacement, PromptUpdate, PromptUpdateDetails, - TimingContext, cached_encode, ) from vllm.sequence import IntermediateTensors @@ -219,11 +202,13 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn return data -class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo]): - """Processing common to both Transformers backend processors: calling the HF - processor, sizing images, and attributing its outputs to a modality. +class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]): + """Locates placeholders from the `text_replacement_offsets` the HF processor + reports, expressing each one as a `PromptUpdate`. - Subclasses add the strategy for locating placeholders in the prompt. + Stating the expansion as an update is what lets it be rebuilt from an + unexpanded prompt, so this processor takes the base class's processing path + and with it the multi-modal processor cache. """ def _get_modality_field_names(self, modality: str) -> set[str]: @@ -415,301 +400,6 @@ class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo] for image, (height, width) in zip(pixel_values, image_sizes.tolist()) ] - -class LegacyMultiModalProcessor(_MultiModalProcessorBase): - """Locates placeholders by searching the prompt the HF processor has already - expanded for the tokens of each modality. - - Serves transformers versions with no `return_text_replacement_offsets`. - Placeholders found this way cannot be rebuilt from an unexpanded prompt, so - this processor overrides `apply` and gets no multi-modal processor cache. - Remove it once `requirements/common.txt` requires `transformers>=5.15.0`. - """ - - def _get_prompt_updates( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - out_mm_kwargs: MultiModalKwargsItems, - ) -> Sequence[PromptUpdate]: - """Empty, because `apply` writes the placeholder ranges itself rather than - deriving them from updates.""" - return [] - - def _get_mm_token_ids(self, modality: str) -> list[int]: - """Token ids marking where `modality` sits in the prompt, which for some - processors differ from the placeholder written into it. - - The expanded prompt is all this path has to go on, so it takes the - processor at its word about which tokens belong to the modality. - """ - info = self.info - processor = info.get_hf_processor() - declared = getattr(processor, f"{modality}_token_ids", None) or () - if ids := [token_id for token_id in declared if token_id is not None]: - return ids - config = info.get_hf_config() - names = (f"{modality}_token_id", f"{modality}_token_index") - token_id = getattr(processor, names[0], getattr_iter(config, names)) - if token_id is None: - token = getattr(processor, f"{modality}_token", None) - if token is not None: - token_id = info.get_tokenizer().get_vocab().get(token) - if token_id is None: - raise ValueError( - f"Cannot find {modality}_token_id on processor or model config" - ) - return [token_id] - - def _apply_audio( - self, - prompt_ids: list[int], - processed_data: "BatchFeature", - num_audios: int, - ) -> dict[str, list[PlaceholderRange]]: - """Take each contiguous run of the audio token as one item's placeholder, - and record how many tokens the run holds.""" - audio_token_ids = self._get_mm_token_ids("audio") - prompt_tensor = torch.tensor(prompt_ids) - is_audio = torch.isin(prompt_tensor, torch.tensor(audio_token_ids)) - - if not is_audio.any(): - raise ValueError( - f"{num_audios} audio item(s) were passed but the prompt " - "contains no audio token. Add one placeholder per audio item." - ) - - padded = torch.cat([torch.tensor([False]), is_audio, torch.tensor([False])]) - transitions = padded.int().diff() - offsets = torch.where(transitions == 1)[0] - lengths = torch.where(transitions == -1)[0] - offsets - - if len(offsets) != num_audios: - raise ValueError( - f"Found {len(offsets)} run(s) of the audio token in the prompt but " - f"{num_audios} audio item(s) were passed. The Transformers backend " - "locates audio placeholders by finding contiguous runs of the audio " - "token, so placeholders with no text between them cannot yet be told " - "apart. Separate them in the prompt to work around this." - ) - - ranges = [ - PlaceholderRange(offset=offset.item(), length=length.item()) - for offset, length in zip(offsets, lengths) - ] - processed_data["num_audio_tokens"] = lengths - return {"audio": ranges} - - def _get_num_multimodal_tokens( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> Mapping[str, list[int]]: - """Ask the HF processor how many tokens and patches each image expands to.""" - processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - images = mm_items.get_items("image", ImageProcessorItems) - image_sizes = [ - (size.height, size.width) - for size in map(images.get_image_size, range(len(images))) - ] - return processor._get_num_multimodal_tokens( - image_sizes=image_sizes, - **self.info.ctx.get_merged_mm_kwargs({}), - ) - - def _apply_vision( - self, - prompt_ids: list[int], - processed_data: "BatchFeature", - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - mm_token_type_ids: torch.Tensor | None, - ) -> dict[str, list[PlaceholderRange]]: - """Split the positions the processor marks as image into one placeholder per - item, sized by the token count it reports for each image.""" - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - if mm_token_type_ids is None: - raise ValueError( - f"{type(hf_processor).__name__} returned no `mm_token_type_ids`, so " - "the Transformers modeling backend cannot locate the placeholder of " - "each image." - ) - - # These mark the tokens structuring an image as well as the image itself, - # which is what the counts below are over - mm_positions = torch.where(mm_token_type_ids[0] == 1)[0] - mm_tokens_per_modality = self._get_num_multimodal_tokens( - mm_items, hf_processor_mm_kwargs - ) - - mm_placeholders: dict[str, list[PlaceholderRange]] = {} - split_sizes = mm_tokens_per_modality["num_image_tokens"] - if sum(split_sizes) != len(mm_positions): - raise ValueError( - f"The expanded prompt holds {len(mm_positions)} image token(s) but " - f"{type(hf_processor).__name__} accounts for {sum(split_sizes)} " - f"across {mm_items.get_count('image')} image item(s)." - ) - - if split_sizes: - image_token_ids = torch.tensor(self._get_mm_token_ids("image")) - mm_tokens = torch.tensor(prompt_ids)[mm_positions] - ranges = [ - PlaceholderRange( - offset=positions[0].item(), - length=positions.shape[0], - # Only some of the span carries embeddings - is_embed=torch.isin(tokens, image_token_ids), - ) - for positions, tokens in zip( - torch.split(mm_positions, split_sizes), - torch.split(mm_tokens, split_sizes), - ) - ] - mm_placeholders = {"image": ranges} - - processed_data["num_image_patches"] = torch.tensor( - mm_tokens_per_modality["num_image_patches"] - ) - return mm_placeholders - - def apply( - self, - inputs: ProcessorInputs, - timing_ctx: TimingContext, - ) -> MultiModalInput: - """Process the prompt and every multi-modal item in one HF processor call, - then read the placeholder ranges out of the token ids it returns.""" - prompt = inputs.prompt - mm_items = inputs.mm_data_items - hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs - - with timing_ctx.record("apply_hf_processor"): - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - prompt_text = hf_processor.decode(prompt) - - # Bypass cached processor and always apply to the full set of mm inputs - # NOTE: we can't just set caching=False because base class method - # transforms outputs to `MultiModalKwargs` which is not going to - # work for Transformers. The vision path has logic tied to - # `mm_tokens_per_modality` in _apply_vision() - hf_processor_mm_kwargs = { - # vLLM needs the untruncated sequence to keep placeholder - # tokens aligned. Note that the text inputs are just dummy - # text, not the original prompt. The original prompt is - # already tokenized by the renderer. - "truncation": False, - # HF processors only accept text, and the decoded string already - # contains any special tokens, so don't let them be added again - # (the HF processor call disables `add_special_tokens`). - "add_special_tokens": False, - **hf_processor_mm_kwargs, - } - - processor_data, _, passthrough_data = self._get_hf_mm_inputs( - mm_items, hf_processor_mm_kwargs - ) - - # The real prompt is fed to the processor below and the placeholder - # ranges are read back out of its output, so the dummy text meant for - # media-only callers is discarded here. - processor_data.pop("text", None) - has_mm_data = any(processor_data.values()) - - # Ask which modality owns each token of the expanded prompt, which - # is the only marking of the tokens an expansion adds around an item - if has_mm_data: - processor_data = {**processor_data, "return_mm_token_type_ids": True} - - try: - processed_data = self.info.ctx.call_hf_processor( - self.info.get_hf_processor(**hf_processor_mm_kwargs), - dict(text=prompt_text, **processor_data), - hf_processor_mm_kwargs, - ) - except ValueError: - if has_mm_data: - raise - # Some processors reject a prompt holding placeholders with - # no data to go with them, so tokenize it without them - prompt_ids = self.info.get_tokenizer().encode(prompt_text) - processed_data = transformers.BatchFeature( - dict(input_ids=[prompt_ids]), tensor_type="pt" - ) - self._unpad_images(processed_data) - self._unpad_audios(processed_data, processor_data, hf_processor_mm_kwargs) - processed_data.update(passthrough_data) - - input_ids = processed_data.pop("input_ids") - if not isinstance(input_ids, list): - input_ids = input_ids.tolist() - - (prompt_ids,) = input_ids - - # Use overrides if provided; fallback to data-dependent hashing. - with timing_ctx.record("get_mm_hashes"): - mm_hashes = inputs.get_mm_hashes( - self.info.model_id, - self.info.ctx.get_mm_config().mm_hasher_algorithm, - ) - - # Gemma3 reports them under the key the model uses for its own token types - mm_token_type_ids = processed_data.pop("token_type_ids", None) - mm_token_type_ids = processed_data.pop("mm_token_type_ids", mm_token_type_ids) - - mm_placeholders: dict[str, list[PlaceholderRange]] = {} - if num_audios := mm_items.get_count("audio", strict=False): - mm_placeholders.update( - self._apply_audio(prompt_ids, processed_data, num_audios) - ) - if mm_items.get_count("image", strict=False): - mm_placeholders.update( - self._apply_vision( - prompt_ids, - processed_data, - mm_items, - hf_processor_mm_kwargs, - mm_token_type_ids, - ) - ) - - mm_kwargs = MultiModalKwargsItems.from_hf_inputs( - processed_data, - self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs), - ) - - # Bypassing `_maybe_apply_prompt_updates` also bypasses its validation. - # `_validate_mm_placeholders` can't be reused because it is typed for the - # `PlaceholderFeaturesInfo` the prompt update machinery produces. - mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - for modality, item_count in mm_item_counts.items(): - num_placeholders = len(mm_placeholders.get(modality, [])) - if num_placeholders != item_count: - raise RuntimeError( - f"Expected there to be {item_count} prompt placeholders " - f"corresponding to {item_count} {modality} items, but instead " - f"found {num_placeholders} prompt placeholders! Make sure the " - "prompt contains a placeholder token for each item." - ) - - return mm_input( - prompt_token_ids=prompt_ids, - mm_kwargs=mm_kwargs, - mm_hashes=mm_hashes, - mm_placeholders=mm_placeholders, - ) - - -class OffsetsMultiModalProcessor(_MultiModalProcessorBase): - """Locates placeholders from the `text_replacement_offsets` the HF processor - reports, expressing each one as a `PromptUpdate`. - - Stating the expansion as an update is what lets it be rebuilt from an - unexpanded prompt, so this processor takes the base class's processing path - and with it the multi-modal processor cache. - """ - def _get_prompt_updates( self, mm_items: MultiModalDataItems, @@ -800,9 +490,6 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase): except (AttributeError, KeyError, TypeError): return None - def _get_hf_mm_text(self, mm_counts: Mapping[str, int]) -> str: - return self.dummy_inputs.get_dummy_text(mm_counts) - def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, @@ -856,8 +543,7 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase): "to reach `ProcessorMixin.get_text_with_replacements` with one " "replacement per item, which usually means implementing " "`replace__token`. Please report this to transformers " - "so it can be fixed, and install `transformers<5.15.0` in the " - "meantime to locate placeholders in the expanded prompt instead." + "so it can be fixed." ) hf_inputs.pop("input_ids", None) return self._finalize_hf_mm_data( @@ -896,15 +582,6 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase): ) -# From this version on, a processor reporting no offsets is an error rather than a -# fallback to searching the expanded prompt -MultiModalProcessor = ( - LegacyMultiModalProcessor - if Version(transformers.__version__) < Version("5.15.0") - else OffsetsMultiModalProcessor -) - - class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Skip SupportsMRoPE.__init__ and call the next class in MRO @@ -982,7 +659,6 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base): super()._decorate_for_torch_compile() # Decorate the encoder model classes to support torch compile if needed if self.compilation_config.compile_mm_encoder: - self.check_version("5.0.0", "multimodal encoder compilation support") encoder_classes = self._pre_trained_model_classes.encoders if not encoder_classes: raise ValueError( @@ -1113,7 +789,6 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base): if input_features is None: return None - self.check_version("5.13.0", "audio models support") num_audio_tokens = kwargs.pop("num_audio_tokens") kwargs.pop("token_type_ids", None) kwargs.pop("mm_token_type_ids", None) diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index c3f249fa80..e0c69c0328 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -11,8 +11,9 @@ ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for tilelang / triton / torch) instead of calling the tilelang kernels directly, and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (True when AITER or TileLang fused MHC is available; False only on the torch fallback); - * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` / - ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py). + * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping``, + ``use_native_mega_moe`` and ``finalize_mega_moe_weights`` do not exist in + amd/model.py). Everything else — the semi-autoregressive drafting hooks, the Markov head, the sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 987909702c..54250185c8 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -17,7 +17,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -99,7 +99,9 @@ class DSparkDeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 @@ -240,7 +242,7 @@ class DSparkDeepseekV4Model(nn.Module): inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: mega_gate_metadata = prepare_mega_gate_routing_metadata( input_ids, has_hash_routing=False, @@ -437,8 +439,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ first_layer = self.model.layers[0] - use_mega_moe = first_layer.ffn.use_mega_moe - if use_mega_moe: + use_native_mega_moe = first_layer.ffn.use_native_mega_moe + if use_native_mega_moe: expert_mapping = make_deepseek_v4_expert_params_mapping( self.config.n_routed_experts ) diff --git a/vllm/models/deepseek_v4/nvidia/fi_moe.py b/vllm/models/deepseek_v4/nvidia/fi_moe.py deleted file mode 100644 index 2e94e75e7d..0000000000 --- a/vllm/models/deepseek_v4/nvidia/fi_moe.py +++ /dev/null @@ -1,397 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FlashInfer ``moe_ep`` expert module for DeepSeek V4.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -import torch -import torch.nn as nn - -from vllm.model_executor.utils import set_weight_attrs -from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts -from vllm.utils.flashinfer_moe_ep import ( - build_fi_mega_layer, - ensure_fi_moe_ep_runtime, - fi_moe_ep_backend_spec, - make_fi_moe_ep_bootstrap, - mega_moe_weight_pack_from_params, -) - -if TYPE_CHECKING: - from flashinfer.moe_ep import MoEEpMegaLayer - - from vllm.config import VllmConfig - from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MLP - -_MOE_SKIP_PADDING: bool | None = None - - -def ckpt_uses_nvfp4_experts(vllm_config: VllmConfig) -> bool: - """True when the loaded checkpoint quantizes experts with modelopt NVFP4 - (e2m1 + fp8-e4m3 per-16 block scales + per-tensor weight_scale_2), i.e. - the recipe the nvfp4_cutedsl prequantized-weights path consumes verbatim. - - DeepSeek-V4 NVFP4-expert checkpoints declare ``moe_quant_algo: NVFP4`` in - their ``quantization_config``, which DeepseekV4FP8Config surfaces as - ``moe_quant_algo``. - """ - return getattr(vllm_config.quant_config, "moe_quant_algo", "") == "NVFP4" - - -def resolve_mega_moe_is_padding(num_tokens: int) -> torch.Tensor | None: - from vllm.forward_context import get_forward_context, is_forward_context_available - - global _MOE_SKIP_PADDING - if _MOE_SKIP_PADDING is None: - import vllm.envs as envs - - _MOE_SKIP_PADDING = bool(envs.VLLM_MOE_SKIP_PADDING) - if not _MOE_SKIP_PADDING or not is_forward_context_available(): - return None - is_padding = get_forward_context().is_padding - if is_padding is None: - return None - return is_padding[:num_tokens] - - -def apply_mega_moe_routing_preprocess( - topk_ids: torch.Tensor, - *, - is_padding: torch.Tensor | None = None, -) -> torch.Tensor: - """Padding-only routing preprocess (EPLB hooks go here later).""" - if is_padding is not None: - topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids) - return topk_ids - - -def nvfp4_prequant_pack_and_alphas( - w13_weight: torch.Tensor, - w13_weight_scale: torch.Tensor, - w13_weight_scale_2: torch.Tensor, # (E_local, 2) fp32: [:,0]=gate(w1), [:,1]=up(w3) - w2_weight: torch.Tensor, - w2_weight_scale: torch.Tensor, - w2_weight_scale_2: torch.Tensor, # (E_local,) fp32 - *, - intermediate_size: int, -): - """NVFP4 checkpoint params -> (MoEWeightPack, fc1_alpha, fc2_alpha). - - Activation quantization is fully dynamic, so the checkpoint's static - ``input_scale`` drops out and each GEMM's epilogue alpha reduces to the - weight's per-tensor ``weight_scale_2``. - - fc1_alpha is one scalar per expert, but gate (w1) and up (w3) carry their - own ``weight_scale_2``. When they differ, the ratio is folded into the up - half's e4m3 block scales — only if the fold round-trips exactly, since a - lossy rescale would silently change the model. - """ - from flashinfer.moe_ep import MoEWeightPack - - inter = intermediate_size - gate_s2 = w13_weight_scale_2[:, 0].float() - up_s2 = w13_weight_scale_2[:, 1].float() - if (gate_s2 <= 0).any() or (up_s2 <= 0).any() or (w2_weight_scale_2 <= 0).any(): - raise ValueError( - "nvfp4 prequant: non-positive weight_scale_2 loaded — checkpoint " - "scale tensors missing or loader routed them wrong." - ) - - w13_scale = w13_weight_scale - if not torch.equal(gate_s2, up_s2): - ratio = up_s2 / gate_s2 # (E_local,) - up_sf = w13_scale[:, inter:, :].float() - folded = up_sf * ratio[:, None, None] - folded_e4m3 = folded.to(torch.float8_e4m3fn) - if not torch.equal(folded_e4m3.float(), folded): - raise ValueError( - "NVFP4 MoE cannot merge the gate and up weight_scale_2 values " - "because their ratio is not exactly representable in " - "float8_e4m3fn. Use a checkpoint quantized with a shared " - "gate/up weight_scale_2." - ) - w13_scale = w13_scale.clone() - w13_scale[:, inter:, :] = folded_e4m3 - - fc1_alpha = gate_s2.clone().contiguous() - fc2_alpha = w2_weight_scale_2.float().clone().contiguous() - pack = MoEWeightPack( - w13=w13_weight, - w2=w2_weight, - w13_scale=w13_scale, - w2_scale=w2_weight_scale, - ) - return pack, fc1_alpha, fc2_alpha - - -class DeepseekV4MegaMoEExpertsFI(DeepseekV4MegaMoEExperts): - """Same weight layout/loader as the native mega experts, FI compute path.""" - - def __init__( - self, - vllm_config: VllmConfig, - *, - activation_clamp: float | None = None, - **kwargs: Any, - ) -> None: - super().__init__(vllm_config, **kwargs) - self._vllm_config = vllm_config - self._activation_clamp = activation_clamp - self._mega_layer: MoEEpMegaLayer | None = None - self._fast_ctx: tuple[Any, Any, Any, int, bool] | None = None - self._epilogue_alphas: tuple[torch.Tensor, torch.Tensor] | None = None - self._nvfp4_prequant = ckpt_uses_nvfp4_experts(vllm_config) - if self._nvfp4_prequant: - megakernel = fi_moe_ep_backend_spec( - vllm_config.kernel_config.moe_backend - ).megakernel - if megakernel != "nvfp4_cutedsl": - raise ValueError( - "NVFP4-quantized expert checkpoint requires " - "moe_backend=flashinfer_moe_ep_mega_cutedsl, got a " - f"backend using megakernel {megakernel!r} " - "(deep_gemm consumes the MXFP4 checkpoint instead)." - ) - self._realloc_nvfp4_params() - - def _realloc_nvfp4_params(self) -> None: - """Swap the mx-recipe scale params for the NVFP4 checkpoint's: - fp8-e4m3 per-16 block scales plus the per-tensor second-level - scales (weight_scale_2, input_scale) modelopt exports.""" - n_e = self.num_local_experts - inter = self.intermediate_size - hidden = self.hidden_size - attrs = {"weight_loader": self.weight_loader} - - def _param(shape: tuple, dtype: torch.dtype) -> nn.Parameter: - p = nn.Parameter(torch.zeros(*shape, dtype=dtype), requires_grad=False) - set_weight_attrs(p, attrs) - return p - - self.w13_weight_scale = _param( - (n_e, 2 * inter, hidden // 16), torch.float8_e4m3fn - ) - self.w13_weight_scale.quant_method = "block" - self.w2_weight_scale = _param((n_e, hidden, inter // 16), torch.float8_e4m3fn) - self.w2_weight_scale.quant_method = "block" - # (E, 2): column 0 = gate (w1), column 1 = up (w3). - self.w13_weight_scale_2 = _param((n_e, 2), torch.float32) - self.w2_weight_scale_2 = _param((n_e,), torch.float32) - # Loaded for completeness; unused — activation quant is dynamic - # (see nvfp4_prequant_pack_and_alphas). - self.w13_input_scale = _param((n_e, 2), torch.float32) - self.w2_input_scale = _param((n_e,), torch.float32) - - def weight_loader( - self, - param: nn.Parameter, - loaded_weight: torch.Tensor, - weight_name: str, - shard_id: str, - expert_id: int, - return_success: bool = False, - ) -> bool | None: - # NVFP4 checkpoint second-level scalars route here; everything - # else (packed weights, block scales) matches the base layout. - if "weight_scale_2" in weight_name or "input_scale" in weight_name: - local_expert_ids = self._map_global_expert_id(expert_id) - if not local_expert_ids: - return False if return_success else None - value = loaded_weight.reshape(()).to(torch.float32) - for local_expert_id in local_expert_ids: - if shard_id in ("w1", "w3"): - if "w13_" not in weight_name: - return False if return_success else None - param.data[local_expert_id, 0 if shard_id == "w1" else 1] = value - elif shard_id == "w2": - if "w2_" not in weight_name: - return False if return_success else None - param.data[local_expert_id] = value - else: - raise ValueError(f"Unsupported expert shard id: {shard_id}") - return True if return_success else None - return super().weight_loader( - param, - loaded_weight, - weight_name, - shard_id, - expert_id, - return_success, - ) - - def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: - # The FlashInfer megakernel has no shared-expert fusion; the caller's - # serial shared MLP path handles shared_experts. - if self._mega_layer is not None: - return - if self.w13_weight is None: - return - - self._check_runtime_supported() - ensure_fi_moe_ep_runtime(self._vllm_config) - - if self._nvfp4_prequant: - # NVFP4 checkpoint: hand the packed weights + both scale - # planes straight to the backend (no dequant->requant); - # per-expert globals become fc1/fc2 epilogue alphas staged - # at every forward via MoEEpTensors. - weights, fc1_alpha, fc2_alpha = nvfp4_prequant_pack_and_alphas( - self.w13_weight.data, - self.w13_weight_scale.data, - self.w13_weight_scale_2.data, - self.w2_weight.data, - self.w2_weight_scale.data, - self.w2_weight_scale_2.data, - intermediate_size=self.intermediate_size, - ) - self._epilogue_alphas = (fc1_alpha, fc2_alpha) - else: - weights = mega_moe_weight_pack_from_params( - self.w13_weight, - self.w13_weight_scale, - self.w2_weight, - self.w2_weight_scale, - megakernel=fi_moe_ep_backend_spec( - self._vllm_config.kernel_config.moe_backend - ).megakernel, - ) - self._mega_layer = build_fi_mega_layer( - make_fi_moe_ep_bootstrap(), - vllm_config=self._vllm_config, - num_experts=self.num_experts, - max_tokens_per_rank=self.max_num_tokens, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - top_k=self.top_k, - activation_clamp=self._activation_clamp, - weights=weights, - ) - del weights - # Allocate (or attach to) the pooled workspace before first - # forward so warmup/capture never hits the lazy path. - self._mega_layer._ensure_workspace() - self.w13_weight = None - self.w13_weight_scale = None - self.w2_weight = None - self.w2_weight_scale = None - if self._nvfp4_prequant: - self.w13_weight_scale_2 = None - self.w2_weight_scale_2 = None - self.w13_input_scale = None - self.w2_input_scale = None - - # EPLB is rejected for these backends in validate_fi_moe_ep_config, so - # nothing constructs an EplbState and none of these run. They are kept - # as explicit errors rather than deleted because inheriting the native - # implementations would be worse: set_eplb_state would record a map - # this forward path never consults, silently splitting routing from - # weights, and get_expert_weights reads transformed-weight attributes - # the flashinfer path releases after preprocess. - _EPLB_UNSUPPORTED = ( - "EPLB is not supported with the flashinfer moe_ep backends; " - "validate_fi_moe_ep_config should have rejected this " - "configuration at startup." - ) - - def set_eplb_state( - self, - moe_layer_idx: int, - expert_load_view: torch.Tensor, - logical_to_physical_map: torch.Tensor, - logical_replica_count: torch.Tensor, - ) -> None: - raise NotImplementedError(self._EPLB_UNSUPPORTED) - - def get_expert_weights(self) -> list[torch.Tensor]: - raise NotImplementedError(self._EPLB_UNSUPPORTED) - - def update_expert_map(self) -> None: - raise NotImplementedError(self._EPLB_UNSUPPORTED) - - def forward( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - *, - activation_clamp: float | None, - fast_math: bool = True, - ) -> torch.Tensor: - # fast_math is a native deep_gemm knob; the FI kernels have no - # equivalent toggle, so it is accepted for signature parity only. - if hidden_states.shape[0] > self.max_num_tokens: - raise ValueError( - f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " - f"but the symmetric buffer was sized for {self.max_num_tokens}." - ) - - from flashinfer.moe_ep import MoEEpTensors - - num_tokens = hidden_states.shape[0] - is_padding = resolve_mega_moe_is_padding(num_tokens) - topk_ids = apply_mega_moe_routing_preprocess( - topk_ids, - is_padding=is_padding, - ) - - # Fast path: after the first successful full forward the layer is - # immutable, so skip MoEEpMegaLayer.forward()'s per-call validation - # and go straight to the kernel backend's stage_inputs + compute. - alphas = self._epilogue_alphas - fc1_alpha = alphas[0] if alphas is not None else None - fc2_alpha = alphas[1] if alphas is not None else None - - fast = self._fast_ctx - if fast is not None: - kernel, workspace, transformed, hidden_size, zero_copy = fast - t = MoEEpTensors( - hidden_states=hidden_states, - topk_ids=topk_ids, - topk_weights=topk_weights, - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - ) - kernel.stage_inputs(t, workspace, quantize_input=True) - if zero_copy: - # Zero-copy (cutedsl backends): consume the workspace [:n] - # view directly — valid under stream ordering until the next - # MoE layer's launch on the shared workspace. - return kernel.compute(workspace, transformed, output=None) - # deep_gemm_mega's compute() requires a real output tensor. - out = torch.empty( - num_tokens, - hidden_size, - dtype=torch.bfloat16, - device=hidden_states.device, - ) - return kernel.compute(workspace, transformed, output=out) - - ensure_fi_moe_ep_runtime(self._vllm_config) - self.finalize_weights() - assert self._mega_layer is not None - - y = self._mega_layer.forward( - MoEEpTensors( - hidden_states=hidden_states, - topk_ids=topk_ids, - topk_weights=topk_weights, - fc1_alpha=fc1_alpha, - fc2_alpha=fc2_alpha, - ) - ) - layer = self._mega_layer - if hidden_states.dtype == torch.bfloat16: - self._fast_ctx = ( - layer._kernel, - layer._ensure_workspace(), - layer._transformed, - layer._fleet_params.token_hidden_size, - # zero-copy output views are a cutedsl-backend contract - layer._kernel.kernel_name() != "deep_gemm_mega", - ) - return y - - -DeepseekV4MegaMoEExpertsFI.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 7edff8ae02..a004120804 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -11,7 +11,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_ep_group, get_pp_group, @@ -85,10 +85,6 @@ from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors -from vllm.utils.flashinfer_moe_ep import ( - is_fi_moe_ep_backend, - validate_fi_moe_ep_config, -) from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id @@ -807,10 +803,11 @@ class DeepseekV4MoE(nn.Module): self.use_sequence_parallel = use_sequence_parallel self.reduce_results = reduce_results moe_backend = vllm_config.kernel_config.moe_backend - validate_fi_moe_ep_config(vllm_config) - self.use_mega_moe = moe_backend in MEGA_MOE_BACKENDS - self.use_fi_mega_moe = is_fi_moe_ep_backend(moe_backend) - if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + self.use_native_mega_moe = moe_backend in NATIVE_MEGA_MOE_BACKENDS + if ( + self.use_native_mega_moe + and not vllm_config.parallel_config.enable_expert_parallel + ): raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " "Enable it with --enable-expert-parallel, or pick a different " @@ -832,15 +829,15 @@ class DeepseekV4MoE(nn.Module): self.swiglu_limit = config.swiglu_limit self.renormalize = config.norm_topk_prob self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") - if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": + if self.use_native_mega_moe and self.scoring_func != "sqrtsoftplus": raise NotImplementedError( "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." ) - if self.use_mega_moe and not self.renormalize: + if self.use_native_mega_moe and not self.renormalize: raise NotImplementedError( "DeepSeek V4 MegaMoE requires normalized top-k probabilities." ) - if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": + if self.use_native_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4": raise NotImplementedError( "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype=" f"{config.expert_dtype!r}. Drop --kernel-config moe_backend=" @@ -864,7 +861,9 @@ class DeepseekV4MoE(nn.Module): image_sentinel_lo if getattr(config, "vision_n_layers", 0) > 0 else 0 ) is_hash_moe = extract_layer_index(prefix) < num_hash_layers - self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 + self.hash_indices_dtype = ( + torch.int64 if self.use_native_mega_moe else torch.int32 + ) if is_hash_moe: # hash MoE doesn't use e_score_correction_bias # Use randint instead of empty to avoid garbage values causing @@ -908,12 +907,12 @@ class DeepseekV4MoE(nn.Module): hidden_act=config.hidden_act, swiglu_limit=self.swiglu_limit, quant_config=quant_config, - reduce_results=self.use_mega_moe, + reduce_results=self.use_native_mega_moe, is_sequence_parallel=use_sequence_parallel, prefix=f"{prefix}.shared_experts", ) - if self.use_mega_moe: + if self.use_native_mega_moe: self._init_mega_moe_experts(vllm_config, config, prefix) else: self._init_fused_moe_experts(vllm_config, config, quant_config, prefix) @@ -950,28 +949,13 @@ class DeepseekV4MoE(nn.Module): # Native DeepGEMM fusion requires each EP rank to own the complete # shared MLP. Sequence parallel replicates those weights while sharding # tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP - # remains tensor-sharded, so retain the serial path. The FlashInfer - # megakernel has no shared-expert fusion, so it keeps the serial path. + # remains tensor-sharded, so retain the serial path. fuse_shared_experts = bool( self.shared_experts is not None and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION and (self.use_sequence_parallel or self.tp_size == 1) - and not self.use_fi_mega_moe ) - activation_clamp = ( - float(self.swiglu_limit) if self.swiglu_limit is not None else None - ) - if self.use_fi_mega_moe: - # Deferred: fi_moe subclasses DeepseekV4MegaMoEExperts, so a - # module-level import here would be circular. - from vllm.models.deepseek_v4.nvidia.fi_moe import ( - DeepseekV4MegaMoEExpertsFI, - ) - - experts_cls: type[DeepseekV4MegaMoEExperts] = DeepseekV4MegaMoEExpertsFI - else: - experts_cls = DeepseekV4MegaMoEExperts expert_kwargs: dict[str, typing.Any] = dict( num_experts=self.n_physical_experts, num_local_experts=self.n_local_physical_experts, @@ -983,9 +967,7 @@ class DeepseekV4MoE(nn.Module): num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0), prefix=f"{prefix}.experts", ) - if self.use_fi_mega_moe: - expert_kwargs["activation_clamp"] = activation_clamp - self.experts = experts_cls(vllm_config, **expert_kwargs) + self.experts = DeepseekV4MegaMoEExperts(vllm_config, **expert_kwargs) def _init_fused_moe_experts( self, @@ -1048,7 +1030,7 @@ class DeepseekV4MoE(nn.Module): if bias_vl is not None and input_ids is None: raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") - if not self.use_mega_moe: + if not self.use_native_mega_moe: return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape @@ -1139,7 +1121,7 @@ class DeepseekV4MoE(nn.Module): return final_hidden_states.view(org_shape) def finalize_mega_moe_weights(self) -> None: - if self.use_mega_moe: + if self.use_native_mega_moe: self.experts.finalize_weights(self.shared_experts) @@ -1179,12 +1161,15 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: def _use_sequence_parallel(vllm_config: VllmConfig) -> bool: parallel_config = vllm_config.parallel_config - use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + moe_backend = vllm_config.kernel_config.moe_backend + moe_needs_token_sharded_input = ( + moe_backend in MEGA_MOE_BACKENDS or parallel_config.data_parallel_size > 1 + ) return ( parallel_config.pipeline_parallel_size == 1 and parallel_config.enable_expert_parallel and parallel_config.tensor_parallel_size > 1 - and (use_mega_moe or parallel_config.data_parallel_size > 1) + and moe_needs_token_sharded_input ) @@ -1425,9 +1410,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): self.config = config self.quant_config = quant_config self.parallel_config = vllm_config.parallel_config - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) - if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + if ( + self.use_native_mega_moe + and not vllm_config.parallel_config.enable_expert_parallel + ): raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " "Enable it with --enable-expert-parallel, or pick a different " @@ -1564,7 +1554,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - if self.use_mega_moe: + if self.use_native_mega_moe: input_ids = input_ids.to(torch.int64) full_num_tokens = positions.shape[0] @@ -1578,7 +1568,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): input_ids = sp_shard(input_ids) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: mega_gate_metadata = prepare_mega_gate_routing_metadata( input_ids, has_hash_routing=self.has_local_hash_moe, @@ -1802,7 +1792,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) - if first_layer.ffn.use_mega_moe: + if first_layer.ffn.use_native_mega_moe: return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 50886b0afa..64547af038 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -20,7 +20,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -95,7 +95,9 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): self.config = config quant_config = vllm_config.quant_config self.rms_norm_eps = config.rms_norm_eps - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -186,7 +188,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): inputs_embeds ).unsqueeze(-2) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: routing_input_ids = input_ids if self.mtp_block.use_sequence_parallel: routing_input_ids = sp_shard(routing_input_ids) @@ -384,7 +386,7 @@ class DeepSeekV4MTP(nn.Module): # Pre-compute expert mapping ONCE. first_layer = next(iter(self.model.layers.values())) - if first_layer.mtp_block.ffn.use_mega_moe: + if first_layer.mtp_block.ffn.use_native_mega_moe: expert_mapping = make_deepseek_v4_expert_params_mapping( self.config.n_routed_experts ) diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index a477d1b5ec..d1534d5248 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -18,7 +18,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -81,7 +81,9 @@ class DSparkDeepseekV4Model(nn.Module): self.rms_norm_eps = config.rms_norm_eps self.num_hidden_layers = config.num_hidden_layers self.target_layer_ids = tuple(config.dspark_target_layer_ids) - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) self.num_dspark_layers = ( @@ -214,7 +216,7 @@ class DSparkDeepseekV4Model(nn.Module): inputs_embeds = sp_shard(inputs_embeds) input_ids = sp_shard(input_ids) mega_gate_metadata = None - if self.use_mega_moe: + if self.use_native_mega_moe: mega_gate_metadata = prepare_mega_gate_routing_metadata( input_ids, has_hash_routing=False, @@ -371,7 +373,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. """ first_layer = self.model.layers[0] - use_mega_moe = first_layer.ffn.use_mega_moe + use_native_mega_moe = first_layer.ffn.use_native_mega_moe # Draft MoE layers use the dspark_* expert counts, not the # backbone's (see DeepseekV4MoE and the reference # ModelArgs.get_moe_config). @@ -379,7 +381,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): getattr(self.config, "dspark_n_routed_experts", 0) or self.config.n_routed_experts ) - if use_mega_moe: + if use_native_mega_moe: expert_mapping = make_deepseek_v4_expert_params_mapping(n_draft_experts) else: expert_mapping = fused_moe_make_expert_params_mapping( diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index b8fa7754b7..bd91e8f12b 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -11,7 +11,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import VllmConfig -from vllm.config.kernel import MEGA_MOE_BACKENDS +from vllm.config.kernel import MEGA_MOE_BACKENDS, NATIVE_MEGA_MOE_BACKENDS from vllm.distributed import ( get_engram_dp_size, get_pp_group, @@ -174,7 +174,7 @@ class DeepseekV4MoE(DeepseekV4MoEBase): def defers_finalize(self, num_tokens: int) -> bool: return ( - not self.use_mega_moe + not self.use_native_mega_moe and self.experts.moe_config.should_defer_moe_finalize(num_tokens) ) @@ -244,12 +244,14 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: def _use_sequence_parallel(vllm_config: VllmConfig) -> bool: parallel_config = vllm_config.parallel_config - use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + moe_needs_token_sharded_input = ( + vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + ) return ( parallel_config.pipeline_parallel_size == 1 and parallel_config.enable_expert_parallel and parallel_config.tensor_parallel_size > 1 - and (use_mega_moe or parallel_config.data_parallel_size > 1) + and (moe_needs_token_sharded_input or parallel_config.data_parallel_size > 1) ) @@ -631,9 +633,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): self.config = config self.quant_config = quant_config self.parallel_config = vllm_config.parallel_config - self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + self.use_native_mega_moe = ( + vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS + ) self.use_sequence_parallel = _use_sequence_parallel(vllm_config) - if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + if ( + self.use_native_mega_moe + and not vllm_config.parallel_config.enable_expert_parallel + ): raise NotImplementedError( "DeepSeek V4 MegaMoE currently requires expert parallel. " "Enable it with --enable-expert-parallel, or pick a different " @@ -865,7 +872,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] - if self.use_mega_moe: + if self.use_native_mega_moe: input_ids = input_ids.to(torch.int64) # Engram n-gram hashes for the whole (flattened) batch, computed once @@ -1015,7 +1022,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): def _mega_gate_metadata( self, input_ids: torch.Tensor | None ) -> MegaGateRoutingMetadata | None: - if not self.use_mega_moe: + if not self.use_native_mega_moe: return None assert input_ids is not None return prepare_mega_gate_routing_metadata( @@ -1331,7 +1338,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) - if first_layer.ffn.use_mega_moe: + if first_layer.ffn.use_native_mega_moe: return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index 761f79d95d..c357608c47 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -8,10 +8,11 @@ import vllm.envs as envs from vllm import _custom_ops # noqa: F401 # registers the torch.ops._C kernels from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import get_current_vllm_config_or_none +from vllm.config import CUDAGraphMode, get_current_vllm_config_or_none from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.indexer_topk import get_indexer_topk from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops from vllm.models.glm5next.common.sparse_indexer import ( RADIX_TOPK_WORKSPACE_SIZE, @@ -24,6 +25,7 @@ from vllm.models.glm5next.common.sparse_indexer import ( kv_cache_as_quant_view, ) from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -37,6 +39,9 @@ from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) +_GFX950_KPOOL_AITER_MAX_POOLS = 256 * 1024 # (1M ctx) +_GFX950_KPOOL_AITER_MIN_POOLS = 16 * 1024 + # kpool write helper: form pools from the current token batch and compress them # into the index K cache via the fused Triton kernel. @@ -87,6 +92,51 @@ def _kpool_compress_insert( ) +def _kpool_decode_topk_backend( + configured: str, + *, + num_rows: int, + max_valid_seq_len: int, + select_k: int, + index_kpool: int, + full_cudagraph: bool, +) -> str: + """Select the topk backend for decodes. + + Heuristic based on ctx lengths: + - <16k pools: in-tree hip kernel + - 16-256k pools: use aiter + - >256k pools: in-tree because not measured on >1m ctx + + Since under FULL cudagraphs we cannot access the context length, + and the aiter kernel does not always outperform the in-tree kernel, + we default to in-tree kernel under FULL cudagraphs. Users can opt in + if they know their context length is long enough to benefit from aiter. + """ + if ( + configured != "auto" + or full_cudagraph + or not rocm_aiter_ops.is_indexer_top_k_enabled() + ): + return configured + if rocm_aiter_ops.is_indexer_top_k_supported( + is_prefill=False, + compress_ratio=index_kpool, + num_rows=num_rows, + max_valid_seq_len=max_valid_seq_len, + ): + return "aiter" + if ( + index_kpool > 1 + and select_k == 512 + and _GFX950_KPOOL_AITER_MIN_POOLS + <= max_valid_seq_len + <= _GFX950_KPOOL_AITER_MAX_POOLS + ): + return "aiter" + return configured + + @eager_break_during_capture def sparse_attn_indexer_kpool( hidden_states: torch.Tensor, @@ -117,6 +167,7 @@ def sparse_attn_indexer_kpool( # path and when the tail cache is disabled. tail_kv_cache: torch.Tensor | None = None, tail_prefix: str | None = None, + topk_backend: str = "auto", ) -> torch.Tensor: # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata @@ -580,15 +631,29 @@ def sparse_attn_indexer_kpool( else: topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - torch.ops._C.top_k_per_row_decode( + full_cudagraph = ( + get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL + ) + if full_cudagraph: + topk_max_seq_len = max_pool_len * index_kpool + else: + topk_max_seq_len = attn_metadata_narrowed.max_seq_len + + resolved_backend = _kpool_decode_topk_backend( + topk_backend, + num_rows=num_rows, + max_valid_seq_len=cdiv(topk_max_seq_len, index_kpool), + select_k=select_k, + index_kpool=index_kpool, + full_cudagraph=full_cudagraph, + ) + get_indexer_topk(resolved_backend)( logits, - next_n, seq_lens, + next_n, topk_dst, - num_rows, - logits.stride(0), - logits.stride(1), select_k, + topk_max_seq_len, ) # Resolve to token-level indices in the output buffer. @@ -664,6 +729,10 @@ class SparseAttnIndexerKpool(CustomOp): self.topk_indices_buffer = topk_indices_buffer self.skip_k_cache_insert = skip_k_cache_insert self.use_fp4_cache = use_fp4_cache + cfg = get_current_vllm_config_or_none() + self.topk_backend = ( + cfg.kernel_config.sparse_indexer_topk_backend if cfg is not None else "auto" + ) def forward_hip( self, @@ -726,6 +795,7 @@ class SparseAttnIndexerKpool(CustomOp): positions, self.tail_cache.kv_cache if self.tail_cache is not None else None, self.tail_cache.prefix if self.tail_cache is not None else None, + self.topk_backend, ) def forward_native( diff --git a/vllm/models/glm5next/common/multimodal.py b/vllm/models/glm5next/common/multimodal.py index 75921b697f..23c90cce73 100644 --- a/vllm/models/glm5next/common/multimodal.py +++ b/vllm/models/glm5next/common/multimodal.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GLM-5.3-Flash vision tower and multimodal processor.""" +import math from collections.abc import Mapping from functools import partial from typing import Any @@ -45,6 +46,7 @@ from vllm.model_executor.models.vision import ( ) from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -394,7 +396,7 @@ class Glm5NextVisionTransformer(nn.Module): head_dim = self.hidden_size // self.num_heads self.rotary_pos_emb = get_rope( head_size=head_dim, - max_position=8192, + max_position=text_config.max_position_embeddings, is_neox_style=True, rope_parameters={"partial_rotary_factor": 0.5}, ) @@ -478,7 +480,7 @@ class Glm5NextVisionTransformer(nn.Module): cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size) - pos_ids = pos_ids.to(cos.device, non_blocking=True) + pos_ids = async_tensor_h2d(pos_ids, device=cos.device) cos_combined = cos[pos_ids].flatten(1) sin_combined = sin[pos_ids].flatten(1) return cos_combined, sin_combined, pos_ids @@ -589,7 +591,7 @@ class Glm5NextVisionTransformer(nn.Module): grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] ).cumsum(dim=0, dtype=torch.int32) cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens]) - cu_seqlens = cu_seqlens.to(self.device, non_blocking=True) + cu_seqlens = async_tensor_h2d(cu_seqlens, device=self.device) max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens) # transformers @@ -672,6 +674,29 @@ class Glm5NextProcessingInfo(Glm4vProcessingInfo): return int(override) return self._processor_pixel_budget(self.get_hf_processor().video_processor)[1] + def get_image_size_with_most_features(self) -> ImageSize: + # The inherited square probe strands budget whenever the token + # ceiling is not a perfect square: with max_image_tokens=8000 the + # square refits to 2492x2492 (89x89 = 7921 tokens) while a + # 2240x2800 canvas reaches 80x100 = 8000, so the encoder cache came + # up short of the processor's own maximum and refused valid images + # (#59539). Factor the ceiling exactly instead of probing a square. + vision_config = self.get_hf_config().vision_config + factor = ( + vision_config.patch_size + * vision_config.spatial_merge_size + * self.get_hf_processor().image_processor.patch_expand_factor + ) + pixels_per_token = vision_config.temporal_patch_size * factor * factor + max_tokens = max(1, self._get_image_max_pixels() // pixels_per_token) + short_side = math.isqrt(max_tokens) + while max_tokens % short_side: + short_side -= 1 + return ImageSize( + width=(max_tokens // short_side) * factor, + height=short_side * factor, + ) + def _get_video_second_idx_glm46v( self, metadata: dict[str, Any], total_frames: int ) -> list[int]: diff --git a/vllm/models/minimax_m3/common/ops/index_topk.py b/vllm/models/minimax_m3/common/ops/index_topk.py index 94b74363e7..025c3dd428 100644 --- a/vllm/models/minimax_m3/common/ops/index_topk.py +++ b/vllm/models/minimax_m3/common/ops/index_topk.py @@ -124,9 +124,11 @@ def _index_block_score_kernel( order=(1, 0), ) q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") - # FP8 has no mixed-dtype dot with bf16/fp32. Leave bf16 and fp32 loads - # in their stored dtype so fp32 pipeline tests keep full precision. - if q.dtype.is_fp8(): + # A matched fp8 pair (e.g. e4m3 x e4m3) lowers to a native FP8 MMA, so + # upcast only when the operands differ: there is no mixed-dtype fp8 MMA, + # including across fp8 flavours (fp8e4nv vs fp8e4b8). bf16 and fp32 loads + # always stay in their stored dtype. + if q.dtype.is_fp8() and q.dtype != ik_cache_ptr.dtype.element_ty: q = q.to(tl.bfloat16) q_start = prefix_len + pid_q * BLOCK_SIZE_Q @@ -162,7 +164,7 @@ def _index_block_score_kernel( + off_k[None, :] * stride_ik_pos + off_d[:, None] * stride_ik_d, ) - if k.dtype.is_fp8(): + if k.dtype.is_fp8() and k.dtype != q.dtype: k = k.to(tl.bfloat16) qk = tl.dot(q, k, out_dtype=tl.float32) # apply causal mask as needed @@ -378,7 +380,7 @@ def _decode_index_score_kernel( mask=q_mask[None, :], other=0.0, ) # [D,HQ] - if q.dtype.is_fp8(): + if q.dtype.is_fp8() and q.dtype != ik_cache_ptr.dtype.element_ty: q = q.to(tl.bfloat16) for blk in tl.range(chunk_start_block, chunk_end_block): page = tl.load(bt_row + blk).to(tl.int64) @@ -393,10 +395,10 @@ def _decode_index_score_kernel( + off_k[:, None] * stride_ik_pos + off_d * stride_ik_d, ) # [N,D] - # Upcast only FP8 cache loads. BF16 operands keep FP8 dot - # instructions out of the fallback path; FP32 accumulation - # preserves score accuracy. BF16/FP32 loads stay as stored. - if k.dtype.is_fp8(): + # Matched fp8 operands keep the native FP8 MMA; any mismatch upcasts, + # since no mixed-dtype fp8 MMA exists. FP32 accumulation preserves + # score accuracy either way. BF16/FP32 loads stay as stored. + if k.dtype.is_fp8() and k.dtype != q.dtype: k = k.to(tl.bfloat16) kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ] kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf")) diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py index f34f0b6335..fa01c75698 100644 --- a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -207,6 +207,110 @@ QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = } +# DGX Spark (GB10) plans, TP=1. +QWEN4_EXP_SM121_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = { + # Shared-expert gate. + (1, 2560): { + 1: SkinnyGemmConfig(1, 64, 1, k_unroll=2, static_k=2560), + 2: SkinnyGemmConfig(2, 64, 1, k_unroll=2, vector_width=4, static_k=2560), + 4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560), + 16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560), + }, + # GDN fused B/A projection. + (96, 2560): { + 1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=2, static_k=2560), + 2: SkinnyGemmConfig(2, 64, 4, vector_width=2, static_k=2560), + 4: SkinnyGemmConfig(4, 128, 4, vector_width=2, static_k=2560), + 8: SkinnyGemmConfig(8, 128, 4, k_unroll=2, vector_width=4, static_k=2560), + 16: SkinnyGemmConfig(16, 64, 1), + }, + # HC merged down/injection projection. + (336, 10240): { + 1: SkinnyGemmConfig(1, 256, 1, vector_width=4, static_k=10240), + 2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, vector_width=4, static_k=10240), + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, static_k=10240), + 8: SkinnyGemmConfig(8, 128, 4, k_unroll=4, static_k=10240), + 16: SkinnyGemmConfig(16, 128, 4, k_unroll=4, static_k=10240), + }, + # Router. + (512, 2560): { + 1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, vector_width=2, static_k=2560), + 2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, vector_width=2, static_k=2560), + 4: SkinnyGemmConfig(4, 64, 1, k_unroll=2, static_k=2560), + 8: SkinnyGemmConfig(8, 64, 2, static_k=2560), + 16: SkinnyGemmConfig(16, 64, 2, k_unroll=2, static_k=2560), + }, + # QSA indexer Q/K projection. + (640, 2560): { + 1: SkinnyGemmConfig(1, 256, 1, vector_width=2, static_k=2560), + 2: SkinnyGemmConfig(2, 256, 1, vector_width=2, static_k=2560), + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=2, vector_width=4, static_k=2560), + 8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560), + 16: SkinnyGemmConfig(16, 64, 1, static_k=2560), + }, + # Shared-expert fused gate/up projection. + (1280, 2560): { + 1: SkinnyGemmConfig(1, 32, 1, k_unroll=2, static_k=2560), + 2: SkinnyGemmConfig(2, 64, 1, vector_width=4, static_k=2560), + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=2, static_k=2560), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=2, vector_width=2, static_k=2560), + 16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, static_k=2560), + }, + # Shared-expert down projection. + (2560, 640): { + 1: SkinnyGemmConfig(1, 64, 1, k_unroll=4, vector_width=2, static_k=640), + 2: SkinnyGemmConfig(2, 64, 1, vector_width=2, static_k=640), + 4: SkinnyGemmConfig(4, 64, 1, vector_width=2, static_k=640), + 8: SkinnyGemmConfig(8, 32, 1, k_unroll=4, vector_width=4, static_k=640), + 16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, vector_width=4, static_k=640), + }, + # GDN and QSA output projections. + (2560, 6144): { + 1: SkinnyGemmConfig(1, 64, 2, vector_width=4, static_k=6144), + 2: SkinnyGemmConfig(2, 64, 1, k_unroll=4, vector_width=4, static_k=6144), + 4: SkinnyGemmConfig(4, 64, 1, k_unroll=4, vector_width=4, static_k=6144), + 8: SkinnyGemmConfig(8, 128, 1, vector_width=2, static_k=6144), + 16: SkinnyGemmConfig(16, 32, 1, static_k=6144), + }, + # HC up projection. + (10240, 320): { + 1: SkinnyGemmConfig(1, 32, 1, k_unroll=2, vector_width=2, static_k=320), + 2: SkinnyGemmConfig(2, 32, 1, vector_width=2, static_k=320), + 4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=320), + 8: SkinnyGemmConfig(8, 32, 1, vector_width=2, static_k=320), + }, + # QSA fused QKV/gate projection. + (13312, 2560): { + 1: SkinnyGemmConfig(1, 32, 1, vector_width=4, static_k=2560), + 2: SkinnyGemmConfig(2, 32, 1, vector_width=4, static_k=2560), + 4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560), + 8: SkinnyGemmConfig(8, 64, 1, vector_width=2, static_k=2560), + 16: SkinnyGemmConfig(16, 32, 1, vector_width=4, static_k=2560), + }, + # GDN fused QKVZ projection. + (16384, 2560): { + 1: SkinnyGemmConfig(1, 32, 1, k_unroll=4, vector_width=2, static_k=2560), + 2: SkinnyGemmConfig(2, 32, 1, vector_width=2, static_k=2560), + 4: SkinnyGemmConfig(4, 32, 1, k_unroll=2, vector_width=2, static_k=2560), + 8: SkinnyGemmConfig(8, 64, 1, vector_width=2, static_k=2560), + 16: SkinnyGemmConfig(16, 32, 1, k_unroll=2, vector_width=2, static_k=2560), + }, + # LM head. + (248320, 2560): { + 1: SkinnyGemmConfig(1, 128, 1, k_unroll=2, vector_width=4, static_k=2560), + 2: SkinnyGemmConfig(2, 128, 1, k_unroll=2, vector_width=4, static_k=2560), + 4: SkinnyGemmConfig(4, 32, 1, vector_width=4, static_k=2560), + 8: SkinnyGemmConfig(8, 64, 1, k_unroll=2, vector_width=4, static_k=2560), + 16: SkinnyGemmConfig(16, 128, 1, k_unroll=2, vector_width=4, static_k=2560), + }, +} + + +def _is_sm121() -> bool: + return current_platform.is_device_capability((12, 1)) + + def _is_sm100() -> bool: return current_platform.is_device_capability((10, 0)) @@ -226,6 +330,8 @@ def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]: return QWEN4_EXP_SM100_GEMM_PLANS if _is_sm90(): return QWEN4_EXP_SM90_GEMM_PLANS + if _is_sm121(): + return QWEN4_EXP_SM121_GEMM_PLANS return {} diff --git a/vllm/multimodal/cache/base.py b/vllm/multimodal/cache/base.py index d2aaeb50e3..f613325f85 100644 --- a/vllm/multimodal/cache/base.py +++ b/vllm/multimodal/cache/base.py @@ -377,6 +377,10 @@ class BaseMultiModalProcessorCache( """ raise NotImplementedError + def release_sender_touches(self) -> None: + """Release the items that were touched but not updated afterwards.""" + return None + @abstractmethod def make_stats(self, *, delta: bool = False) -> CacheInfo: """Get (and reset) the multi-modal cache stats. diff --git a/vllm/multimodal/cache/shm.py b/vllm/multimodal/cache/shm.py index ab84dcf6af..2fd8ad47e3 100644 --- a/vllm/multimodal/cache/shm.py +++ b/vllm/multimodal/cache/shm.py @@ -170,10 +170,13 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache): @override def touch_sender_cache_item(self, mm_hash: str) -> None: - """Touch the item in shared memory cache to prevent eviction. - Increments writer_flag on sender side.""" + """Touch the item in shared memory cache to prevent eviction.""" self._shm_cache.touch(mm_hash) + @override + def release_sender_touches(self) -> None: + self._shm_cache.release_touches() + @override def validate_input_item( self, @@ -305,8 +308,7 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache): mm_hash: str, mm_item: MultiModalKwargsItem | None = None, ) -> None: - """Touch the item in shared memory cache to prevent eviction. - Increments reader_count on receiver side.""" + """Validate the item's handle in shared memory cache.""" assert mm_item is not None if (handle := _get_shm_handle(mm_item)) is not None: address, monotonic_id, signature = handle diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index eb53023f0f..7267020353 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -1534,13 +1534,16 @@ class BaseMultiModalProcessor(ABC, Generic[_I]): ) with timing_ctx.record("merge_mm_kwargs"): - mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs( - cache, - mm_hashes=mm_hashes, - mm_is_cached=mm_is_cached, - mm_missing_kwargs=mm_missing_kwargs, - mm_missing_prompt_updates=mm_missing_prompt_updates, - ) + try: + mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs( + cache, + mm_hashes=mm_hashes, + mm_is_cached=mm_is_cached, + mm_missing_kwargs=mm_missing_kwargs, + mm_missing_prompt_updates=mm_missing_prompt_updates, + ) + finally: + cache.release_sender_touches() return MultiModalProcessingResult( prompt_ids=self._postprocess_prompt(inputs.prompt), diff --git a/vllm/outputs.py b/vllm/outputs.py index 8bbbbdf40c..35f71ec304 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -18,6 +18,15 @@ from vllm.v1.metrics.stats import RequestSpecDecodeMetrics, RequestStateStats logger = init_logger(__name__) +@dataclass(frozen=True) +class RequestError: + """A request-level error that does not mark the engine as unhealthy.""" + + code: str + message: str + retryable: bool = False + + @dataclass class SamplingMask: """Per-token sampling support sets aligned with completion token IDs. @@ -118,9 +127,9 @@ class RequestOutput: decoder input prompt token ids. prompt_logprobs: The log probabilities to return per prompt token. prompt_token_id_logprobs: Logprobs of prompt_logprob_token_ids, shaped - [max(prompt_len - 1 - prompt_logprob_start, 0), - len(prompt_logprob_token_ids)]; row i scores them as predictions of - prompt token prompt_logprob_start + i + 1. + [prompt_len - 1 - prompt_logprob_start, num_ids]; row i holds the + logprobs of its own IDs as predictions of prompt token + prompt_logprob_start + i + 1, with -inf in padded columns. outputs: The output sequences of the request. finished: Whether the whole request is finished. metrics: Metrics associated with the request. @@ -259,6 +268,7 @@ class PoolingRequestOutput(Generic[_O]): prompt_token_ids (list[int]): A list of token IDs used in the prompt. num_cached_tokens: The number of tokens with prefix cache hit. finished (bool): A flag indicating whether the pooling is completed. + error: Structured request-level error information, if the request failed. """ @@ -269,20 +279,25 @@ class PoolingRequestOutput(Generic[_O]): prompt_token_ids: list[int], num_cached_tokens: int, finished: bool, + *, + error: RequestError | None = None, ): self.request_id = request_id self.prompt_token_ids = prompt_token_ids self.num_cached_tokens = num_cached_tokens self.finished = finished self.outputs = outputs + self.error = error def __repr__(self) -> str: + error = f", error={self.error!r}" if self.error is not None else "" return ( f"{type(self).__name__}(request_id={self.request_id!r}, " f"outputs={self.outputs!r}, " f"prompt_token_ids={self.prompt_token_ids}, " f"num_cached_tokens={self.num_cached_tokens}, " - f"finished={self.finished})" + f"finished={self.finished}" + f"{error})" ) @@ -319,12 +334,18 @@ class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]): def from_base( request_output: PoolingRequestOutput, ) -> "EmbeddingRequestOutput": + outputs = ( + EmbeddingOutput([]) + if request_output.error is not None + else EmbeddingOutput.from_base(request_output.outputs) + ) return EmbeddingRequestOutput( request_id=request_output.request_id, - outputs=EmbeddingOutput.from_base(request_output.outputs), + outputs=outputs, prompt_token_ids=request_output.prompt_token_ids, num_cached_tokens=request_output.num_cached_tokens, finished=request_output.finished, + error=request_output.error, ) @@ -362,12 +383,18 @@ class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]): def from_base( request_output: PoolingRequestOutput, ) -> "ClassificationRequestOutput": + outputs = ( + ClassificationOutput([]) + if request_output.error is not None + else ClassificationOutput.from_base(request_output.outputs) + ) return ClassificationRequestOutput( request_id=request_output.request_id, - outputs=ClassificationOutput.from_base(request_output.outputs), + outputs=outputs, prompt_token_ids=request_output.prompt_token_ids, num_cached_tokens=request_output.num_cached_tokens, finished=request_output.finished, + error=request_output.error, ) @@ -402,10 +429,16 @@ class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]): def from_base( request_output: PoolingRequestOutput, ) -> "ScoringRequestOutput": + outputs = ( + ScoringOutput(0.0) + if request_output.error is not None + else ScoringOutput.from_base(request_output.outputs) + ) return ScoringRequestOutput( request_id=request_output.request_id, - outputs=ScoringOutput.from_base(request_output.outputs), + outputs=outputs, prompt_token_ids=request_output.prompt_token_ids, num_cached_tokens=request_output.num_cached_tokens, finished=request_output.finished, + error=request_output.error, ) diff --git a/vllm/parser/CLAUDE.md b/vllm/parser/CLAUDE.md deleted file mode 100644 index 43c994c2d3..0000000000 --- a/vllm/parser/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -@AGENTS.md diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index eea957f90c..bee090e7f6 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -506,6 +506,11 @@ class DelegatingParser(Parser): is_auto = request.tool_choice == "auto" single_call = request.parallel_tool_calls is False strict_level = self.tool_strict_level + if ( + strict_level == ToolStrictLevel.AUTO + and tool_parser.default_tool_strict_level is not None + ): + strict_level = tool_parser.default_tool_strict_level if single_call and not (is_auto and structured_outputs is not None): # Limiting the call count needs the call envelope in the grammar. # Auto with structured outputs keeps its format-only grammar, so the diff --git a/vllm/reasoning/CLAUDE.md b/vllm/reasoning/CLAUDE.md deleted file mode 100644 index 43c994c2d3..0000000000 --- a/vllm/reasoning/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -@AGENTS.md diff --git a/vllm/reasoning/hy_v4_reasoning_parser.py b/vllm/reasoning/hy_v4_reasoning_parser.py index 65035ed53a..eabe6c31ae 100644 --- a/vllm/reasoning/hy_v4_reasoning_parser.py +++ b/vllm/reasoning/hy_v4_reasoning_parser.py @@ -48,21 +48,18 @@ def detect_token_suffix(tokenizer: TokenizerLike) -> str: round-trips. """ - import transformers - - if int(transformers.__version__.split(".")[0]) >= 5: - init_kwargs = getattr(tokenizer, "init_kwargs", None) or {} - think_begin_as_special = init_kwargs.get( - "model_specific_special_tokens", {} - ).get("think_begin_token", "") - if think_begin_as_special: - raise RuntimeError( - "This checkpoint declares HYV4 structural tokens (think_begin_token" - "/toolcalls_begin_token/argkey_begin_token) in " - "tokenizer_config.json, which transformers 5 no longer supports. " - "Remove those fields and keep the tokens in the tokenizer's own " - "token definitions so the suffix can be read from the vocab." - ) + init_kwargs = getattr(tokenizer, "init_kwargs", None) or {} + think_begin_as_special = init_kwargs.get("model_specific_special_tokens", {}).get( + "think_begin_token", "" + ) + if think_begin_as_special: + raise RuntimeError( + "This checkpoint declares HYV4 structural tokens (think_begin_token" + "/toolcalls_begin_token/argkey_begin_token) in " + "tokenizer_config.json, which transformers 5 no longer supports. " + "Remove those fields and keep the tokens in the tokenizer's own " + "token definitions so the suffix can be read from the vocab." + ) structural_token_re = re.compile( r"<(?:think|tool_calls|tool_call|arg_key|arg_value)(:[^\s>]+)?>" diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 0261401dd4..d221b32a8b 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -28,8 +28,8 @@ from vllm.entrypoints.openai.completion.protocol import ( ) from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( DerenderStreamState, - GenerateResponse, - GenerateStreamResponse, + GenerateTokensResponse, + GenerateTokensStreamResponse, ) from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm.entrypoints.serve.utils.request_logger import RequestLogger @@ -118,7 +118,7 @@ class OnlineDerenderer: async def derender_chat( self, - generate_response: GenerateResponse, + generate_response: GenerateTokensResponse, chat_request: ChatCompletionRequest | None = None, prompt_token_ids: list[int] | None = None, ) -> list[ChatCompletionResponseChoice]: @@ -128,7 +128,7 @@ class OnlineDerenderer: def _derender_chat( self, - generate_response: GenerateResponse, + generate_response: GenerateTokensResponse, chat_request: ChatCompletionRequest | None = None, prompt_token_ids: list[int] | None = None, ) -> list[ChatCompletionResponseChoice]: @@ -323,7 +323,7 @@ class OnlineDerenderer: async def derender_chat_stream( self, model: str, - generate_chunk: GenerateStreamResponse, + generate_chunk: GenerateTokensStreamResponse, state: DerenderStreamState | None = None, chat_request: ChatCompletionRequest | None = None, prompt_tokens: int | None = None, @@ -472,7 +472,7 @@ class OnlineDerenderer: self, parser_cls: type[Parser], model: str, - generate_chunk: GenerateStreamResponse, + generate_chunk: GenerateTokensStreamResponse, state: DerenderStreamState, chat_request: ChatCompletionRequest, prompt_tokens: int | None, @@ -682,7 +682,7 @@ class OnlineDerenderer: async def derender_completion( self, - generate_responses: list[GenerateResponse], + generate_responses: list[GenerateTokensResponse], prompt_tokens: list[int] | None = None, completion_request: CompletionRequest | None = None, prompt_token_ids: list[list[int] | None] | None = None, @@ -693,7 +693,7 @@ class OnlineDerenderer: def _derender_completion( self, - generate_responses: list[GenerateResponse], + generate_responses: list[GenerateTokensResponse], prompt_tokens: list[int] | None = None, completion_request: CompletionRequest | None = None, prompt_token_ids: list[list[int] | None] | None = None, @@ -761,7 +761,7 @@ class OnlineDerenderer: async def derender_completion_stream( self, model: str, - generate_chunk: GenerateStreamResponse, + generate_chunk: GenerateTokensStreamResponse, state: DerenderStreamState | None = None, prompt_tokens: int | None = None, completion_request: CompletionRequest | None = None, diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 15657997ab..246e2ce90c 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -11,7 +11,8 @@ from functools import cached_property from typing import Annotated, Any import msgspec -from pydantic import BeforeValidator +import numpy as np +from pydantic import BeforeValidator, GetPydanticSchema, StrictInt from pydantic.dataclasses import dataclass import vllm.envs as envs @@ -20,6 +21,7 @@ from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike from vllm.utils.mistral import is_mistral_tokenizer +from vllm.utils.torch_utils import make_ndarray_with_pad from vllm.v1.serial_utils import PydanticMsgspecMixin logger = init_logger(__name__) @@ -289,10 +291,14 @@ class SamplingParams( prompt_logprobs: int | None = None """Number of log probabilities to return per prompt token. When set to -1, return all `vocab_size` log probabilities.""" - prompt_logprob_token_ids: list[int] | None = None - """Token IDs to score at each scored causal prompt row, where row i scores - them as predictions of prompt token i + 1. The last prompt row is excluded, - so the result has `max(prompt_len - 1 - prompt_logprob_start, 0)` rows.""" + prompt_logprob_token_ids: ( + Annotated[np.ndarray, GetPydanticSchema(lambda _, h: h(list[list[StrictInt]]))] + | None + ) = None + """Token IDs to score per causal prompt row: an integer array of shape + [num_rows, num_ids] or nested lists; -1 entries (padding) score -inf. + Row i scores its IDs as predictions of prompt token + prompt_logprob_start + i + 1, so the last prompt row is excluded.""" prompt_logprob_start: int | None = None """First causal prompt row to score; defaults to the first row.""" logprob_token_ids: list[int] | None = None @@ -404,7 +410,7 @@ class SamplingParams( min_tokens: int = 0, logprobs: int | None = None, prompt_logprobs: int | None = None, - prompt_logprob_token_ids: list[int] | None = None, + prompt_logprob_token_ids: np.ndarray | list[list[int]] | None = None, prompt_logprob_start: int | None = None, detokenize: bool = True, skip_special_tokens: bool = True, @@ -916,42 +922,58 @@ class SamplingParams( ) # Validate prompt_logprob_token_ids. - if self.prompt_logprob_token_ids is not None: - n = len(self.prompt_logprob_token_ids) - if n == 0: + ids = self.prompt_logprob_token_ids + if ids is not None: + shape: tuple[int, ...] = () + if isinstance(ids, list) and all(isinstance(row, list) for row in ids): + shape = (len(ids), max(map(len, ids), default=0)) + elif isinstance(ids, np.ndarray) and ids.dtype.kind in "iu": + shape = ids.shape + if len(shape) != 2 or 0 in shape: raise VLLMValidationError( - "prompt_logprob_token_ids must not be empty.", + "prompt_logprob_token_ids must be a non-empty integer array " + "of shape [num_rows, num_ids].", parameter="prompt_logprob_token_ids", - value=n, + value=getattr(ids, "shape", type(ids).__name__), ) - if n > max_logprobs: + max_rows = model_config.max_model_len - 1 + if shape[0] > max_rows: raise VLLMValidationError( - f"Requested prompt_logprob_token_ids of length {n}, " - f"which is greater than max allowed: {max_logprobs}. " - f"Set max_logprobs (--max-logprobs) to at least {n}.", + f"prompt_logprob_token_ids has {shape[0]} rows, but a prompt " + f"has at most {max_rows} scored rows (max_model_len - 1).", parameter="prompt_logprob_token_ids", - value=n, + value=shape[0], ) + num_ids = shape[1] + if num_ids > max_logprobs: + raise VLLMValidationError( + f"Requested {num_ids} token ids per row in " + f"prompt_logprob_token_ids, which is greater than max allowed: " + f"{max_logprobs}. " + f"Set max_logprobs (--max-logprobs) to at least {num_ids}.", + parameter="prompt_logprob_token_ids", + value=num_ids, + ) + if isinstance(ids, list): + # Pad only after the checks above bound the allocation. + try: + ids = make_ndarray_with_pad(ids, -1, np.int64, max_len=num_ids) + except (TypeError, ValueError, OverflowError) as e: + raise VLLMValidationError( + "prompt_logprob_token_ids must contain integer token ids.", + parameter="prompt_logprob_token_ids", + ) from e vocab_size = model_config.get_vocab_size() - invalid_token_ids = [ - token_id - for token_id in self.prompt_logprob_token_ids - if token_id < 0 or token_id >= vocab_size - ] - if invalid_token_ids: + lo, hi = int(ids.min()), int(ids.max()) + if lo < -1 or hi >= vocab_size: raise VLLMValidationError( - f"token_id(s) {invalid_token_ids} in " - f"prompt_logprob_token_ids contain out-of-vocab token ids. " - f"Vocabulary size: {vocab_size}", + "prompt_logprob_token_ids contain out-of-vocab token ids " + f"(-1 pads a row). Vocabulary size: {vocab_size}", parameter="prompt_logprob_token_ids", - value=invalid_token_ids, - ) - if len(set(self.prompt_logprob_token_ids)) != n: - raise VLLMValidationError( - "prompt_logprob_token_ids must not contain duplicates.", - parameter="prompt_logprob_token_ids", - value=self.prompt_logprob_token_ids, + value=[lo, hi], ) + # In-vocab IDs fit int32, like the runner's other token ID buffers. + self.prompt_logprob_token_ids = np.ascontiguousarray(ids, dtype=np.int32) if self.prompt_logprob_start is not None and self.prompt_logprob_start < 0: raise VLLMValidationError( "prompt_logprob_start must be non-negative.", diff --git a/vllm/snapshot/runtime.py b/vllm/snapshot/runtime.py index 582bc9a7d9..fdced5116c 100644 --- a/vllm/snapshot/runtime.py +++ b/vllm/snapshot/runtime.py @@ -31,6 +31,7 @@ from typing import NoReturn import regex as re from pydantic import ValidationError +from vllm.logger import init_logger from vllm.snapshot.manifest import ( ReadyMarker, SnapshotManifest, @@ -43,6 +44,8 @@ from vllm.snapshot.manifest import ( ) from vllm.snapshot.types import Oracle +logger = init_logger(__name__) + class SnapshotCreateError(RuntimeError): """Snapshot creation did not produce a complete artifact.""" @@ -59,6 +62,7 @@ _COMMON_CRIU_OPTIONS = ( "--link-remap", "--file-locks", ) +_CRIU_DIAGNOSTIC_LIMIT = 8192 _TCP_TABLES = ( ("AF_INET", Path("/proc/net/tcp")), ("AF_INET6", Path("/proc/net/tcp6")), @@ -254,6 +258,7 @@ class LocalSnapshotTools: ] child_environment = os.environ.copy() child_environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + child_environment.setdefault("NCCL_IB_DISABLE", "1") with (workdir / "child.log").open("wb") as log_file: process = subprocess.Popen( command, @@ -357,6 +362,13 @@ class LocalSnapshotTools: if "anon_inode:[io_uring]" in targets: io_uring_pids.append(pid) for target in targets: + if target.startswith("/dev/infiniband/"): + raise SnapshotCreateError( + "snapshot process has open InfiniBand/RDMA state that " + f"CRIU cannot capture: pid {pid}, {target}. If NCCL owns " + "it, use NCCL_IB_DISABLE=1. Close other RDMA clients " + "before snapshot creation." + ) if target.startswith("socket:[") and target.endswith("]"): # Shared descriptors name one holder; any of them locates # the connection for the operator. @@ -495,7 +507,49 @@ class LocalSnapshotTools: str(self.plugin_dir), *arguments, ] - self._run(command) + try: + self._run(command) + except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error: + self._log_criu_diagnostics(error, action, artifact) + raise + + def _log_criu_diagnostics( + self, + error: subprocess.CalledProcessError | subprocess.TimeoutExpired, + action: str, + artifact: Path, + ) -> None: + def bounded(value: str | bytes) -> str: + if isinstance(value, bytes): + return value[-_CRIU_DIAGNOSTIC_LIMIT:].decode(errors="replace") + return value[-_CRIU_DIAGNOSTIC_LIMIT:] + + for name in ("stdout", "stderr"): + value = getattr(error, name, None) + if value: + logger.error("CRIU %s (tail):\n%s", name, bounded(value)) + + log_path = artifact / "images" / f"{action}.log" + try: + result = self._run( + [ + *self._privileged(), + "tail", + "-c", + str(_CRIU_DIAGNOSTIC_LIMIT), + "--", + str(log_path), + ], + timeout=5, + ) + if result.stdout: + logger.error("CRIU %s.log (tail):\n%s", action, bounded(result.stdout)) + except BaseException as diagnostic_error: + logger.error( + "CRIU %s.log unavailable: %s", + action, + bounded(_error_detail(diagnostic_error)), + ) def _record_child_log_size(self, artifact: Path) -> None: child_log = artifact / "child.log" @@ -512,10 +566,18 @@ class LocalSnapshotTools: os.fsync(output.fileno()) def _link_remap_names(self) -> set[str]: + try: + entries = tuple(self.shm_dir.iterdir()) + except OSError as error: + raise SnapshotCreateError( + f"could not inventory CRIU link remaps in {self.shm_dir}: " + f"{_error_detail(error)}" + ) from error return { path.name - for path in self.shm_dir.glob("link_remap.*") - if path.name.removeprefix("link_remap.").isdigit() + for path in entries + if path.name.startswith("link_remap.") + and path.name.removeprefix("link_remap.").isdigit() } def _capture_link_remaps(self, artifact: Path, names: set[str]) -> None: @@ -565,7 +627,8 @@ class LocalSnapshotTools: try: for source in sorted(source_dir.iterdir()): if ( - not source.name.removeprefix("link_remap.").isdigit() + not source.name.startswith("link_remap.") + or not source.name.removeprefix("link_remap.").isdigit() or source.is_symlink() or not source.is_file() ): @@ -662,8 +725,16 @@ class LocalSnapshotTools: workdir, self._link_remap_names() - remaps_before, ) - except BaseException: - for name in self._link_remap_names() - remaps_before: + except BaseException as primary: + try: + remaps_after = self._link_remap_names() + except BaseException as cleanup_error: + raise SnapshotCreateError( + f"snapshot dump failed: {_error_detail(primary)}; " + "remap cleanup inventory failed: " + f"{_error_detail(cleanup_error)}" + ) from primary + for name in remaps_after - remaps_before: (self.shm_dir / name).unlink(missing_ok=True) raise self._record_child_log_size(workdir) @@ -747,6 +818,7 @@ class LocalSnapshotTools: prefixes = ("VLLM_", "CUDA_", "NCCL_", "TORCH_", "TRITON_") environment = os.environ.copy() environment.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + environment.setdefault("NCCL_IB_DISABLE", "1") selected = tuple( sorted( (key, hashlib.sha256(key.encode() + b"\0" + value.encode()).hexdigest()) diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index 09a715c117..e7064fc252 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -7,12 +7,6 @@ from pathlib import Path from typing import TYPE_CHECKING import huggingface_hub -import transformers -from packaging.version import Version -from transformers.models.auto.tokenization_auto import ( - TOKENIZER_MAPPING_NAMES as _HF_TOKENIZER_MAPPING_NAMES, -) -from transformers.utils.import_utils import is_mistral_common_available from typing_extensions import TypeVar, assert_never from vllm import envs @@ -32,8 +26,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -_MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL = Version("5.15.0") - # Model types whose hub tokenizer_class is incorrect and should be overridden with # TokenizersBackend (the generic fast tokenizer). Adding a model type here is always a @@ -243,26 +235,6 @@ def get_tokenizer( token=kwargs.get("token"), ) - # TODO: delete when Transformers version dependency is bumped >= 5.15.0 - if ( - tokenizer_mode == "hf" - and Version(transformers.__version__) < _MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL - and is_mistral_common_available() - and _HF_TOKENIZER_MAPPING_NAMES.get(getattr(config, "model_type", None)) - == "MistralCommonBackend" - and any_pattern_in_repo_files( - model_name_or_path=str(tokenizer_name), - allow_patterns=["tekken.json", "tokenizer.model.v*"], - revision=revision, - token=kwargs.get("token"), - ) - ): - raise ValueError( - "Loading Mistral models with tokenizer_mode='hf' requires " - f"transformers>={_MIN_TRANSFORMERS_VERSION_FOR_HF_MISTRAL}. " - "Please upgrade transformers or delete Mistral tokenizer files." - ) - # Some models have an incorrect tokenizer_class on the hub. # For these model types, bypass AutoTokenizer and use TokenizersBackend directly. model_type = getattr(config, "model_type", None) if config else None diff --git a/vllm/tool_parsers/CLAUDE.md b/vllm/tool_parsers/CLAUDE.md deleted file mode 100644 index 43c994c2d3..0000000000 --- a/vllm/tool_parsers/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -@AGENTS.md diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index bc5abea3a2..3d3f6c1f09 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/vllm/tool_parsers/abstract_tool_parser.py @@ -59,6 +59,10 @@ class ToolParser: # xgrammar builtin structural tag model key. Subclasses set this when # their parsed tool-call syntax matches a builtin xgrammar format. structural_tag_model: str | None = None + # Strict level applied when the operator leaves ``--tool-strict-level`` + # at auto. Operator flags override it, and + # ``VLLM_ENFORCE_STRICT_TOOL_CALLING=0`` still disables structural tags. + default_tool_strict_level: ToolStrictLevel | None = None engine_based_streaming: bool = False def __init_subclass__(cls, **kwargs: Any) -> None: diff --git a/vllm/tool_parsers/glm47_moe_tool_parser.py b/vllm/tool_parsers/glm47_moe_tool_parser.py index 70275a6ac0..6beb9185fb 100644 --- a/vllm/tool_parsers/glm47_moe_tool_parser.py +++ b/vllm/tool_parsers/glm47_moe_tool_parser.py @@ -4,8 +4,11 @@ from __future__ import annotations from vllm.parser.engine.registered_adapters import Glm47MoeParserToolAdapter +from vllm.tool_parsers.tool_strict_level import ToolStrictLevel class Glm47MoeModelToolParser(Glm47MoeParserToolAdapter): # type: ignore[valid-type, misc] supports_required_and_named = False structural_tag_model = "glm_4_7" + # Constrain auto tool calls with the shallow tag by default. + default_tool_strict_level = ToolStrictLevel.FUNCTION diff --git a/vllm/tool_parsers/hy_v4_tool_parser.py b/vllm/tool_parsers/hy_v4_tool_parser.py index b9c857dbd9..2ce0381a01 100644 --- a/vllm/tool_parsers/hy_v4_tool_parser.py +++ b/vllm/tool_parsers/hy_v4_tool_parser.py @@ -104,21 +104,18 @@ def detect_token_suffix(tokenizer: TokenizerLike) -> str: round-trips. """ - import transformers - - if int(transformers.__version__.split(".")[0]) >= 5: - init_kwargs = getattr(tokenizer, "init_kwargs", None) or {} - think_begin_as_special = init_kwargs.get( - "model_specific_special_tokens", {} - ).get("think_begin_token", "") - if think_begin_as_special: - raise RuntimeError( - "This checkpoint declares HYV4 structural tokens (think_begin_token" - "/toolcalls_begin_token/argkey_begin_token) in " - "tokenizer_config.json, which transformers 5 no longer supports. " - "Remove those fields and keep the tokens in the tokenizer's own " - "token definitions so the suffix can be read from the vocab." - ) + init_kwargs = getattr(tokenizer, "init_kwargs", None) or {} + think_begin_as_special = init_kwargs.get("model_specific_special_tokens", {}).get( + "think_begin_token", "" + ) + if think_begin_as_special: + raise RuntimeError( + "This checkpoint declares HYV4 structural tokens (think_begin_token" + "/toolcalls_begin_token/argkey_begin_token) in " + "tokenizer_config.json, which transformers 5 no longer supports. " + "Remove those fields and keep the tokens in the tokenizer's own " + "token definitions so the suffix can be read from the vocab." + ) structural_token_re = re.compile( r"<(?:think|tool_calls|tool_call|arg_key|arg_value)(:[^\s>]+)?>" diff --git a/vllm/tool_parsers/structural_tag_registry.py b/vllm/tool_parsers/structural_tag_registry.py index ef57e85fab..a1ee263dc2 100644 --- a/vllm/tool_parsers/structural_tag_registry.py +++ b/vllm/tool_parsers/structural_tag_registry.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json from collections.abc import Callable, Sequence from typing import Any, Literal, TypeAlias @@ -71,12 +73,13 @@ XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset( "mimo", "qwen_3", "deepseek_v3_2", - "glm_4_7", "deepseek_v4", "deepseek_v4_1", } ) -VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset({"hermes", "hy_v4", "kimi_k3", "plamo3"}) +VLLM_BUILTIN_STRUCTURAL_TAG_MODELS = frozenset( + {"glm_4_7", "hermes", "hy_v4", "kimi_k3", "plamo3"} +) SUPPORTED_STRUCTURAL_TAG_MODELS = ( XGRAMMAR_BUILTIN_STRUCTURAL_TAG_MODELS | VLLM_BUILTIN_STRUCTURAL_TAG_MODELS ) @@ -903,3 +906,182 @@ def get_hy_v4_structural_tag( # the ``...`` prefix. prefix_tag = TagFormat(begin="", content=AnyTextFormat(), end=think_end) return StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) + + +# --------------------------------------------------------------------------- +# GLM-4.7 tool calls (NAMEKV) +# --------------------------------------------------------------------------- +# Assistant output after the reasoning gate (````): +# [] +# NAMEK1V1... +# Values are raw JSON scalars: strings verbatim (unquoted), other types as +# bare JSON, matching how the chat template serializes tool-call history. + +_GLM_4_7_VALUE_EXCLUDES = [ + "", + "", + "", + "", + "", + "", + "", + "", +] + + +def _glm_4_7_value_format(prop: dict[str, Any]) -> Any: + if "enum" in prop: + values = [ + ConstStringFormat(value=v if isinstance(v, str) else json.dumps(v)) + for v in prop["enum"] + ] + return values[0] if len(values) == 1 else OrFormat(elements=values) + prop_type = prop.get("type") + if isinstance(prop_type, list): + formats = [ + _glm_4_7_value_format({"type": t}) for t in prop_type if t != "string" + ] + if "string" in prop_type: + formats.append(AnyTextFormat(excludes=_GLM_4_7_VALUE_EXCLUDES)) + return formats[0] if len(formats) == 1 else OrFormat(elements=formats) + if prop_type in (None, "string"): + return AnyTextFormat(excludes=_GLM_4_7_VALUE_EXCLUDES) + return JSONSchemaFormat(json_schema={"type": prop_type}) + + +def _glm_4_7_plain_properties(params: dict[str, Any]) -> dict[str, Any] | None: + """Return ``properties`` when the schema is a plain object schema.""" + properties = params.get("properties") + if ( + not isinstance(properties, dict) + or not properties + or params.get("additionalProperties") is True + or any( + key in params + for key in ("$ref", "allOf", "anyOf", "oneOf", "patternProperties") + ) + ): + return None + return properties + + +def _glm_4_7_tool_tag(tool: FunctionToolParam) -> TagFormat: + """Shallow tag for a non-strict tool: keys and basic value types pinned.""" + properties = _glm_4_7_plain_properties(tool.function.parameters or {}) + if properties is None: + return TagFormat( + begin=f"{tool.function.name}", + content=JSONSchemaFormat( + json_schema=True, style="glm_xml", any_order=False + ), + end="", + ) + pairs: list[Any] = [ + SequenceFormat( + elements=[ + ConstStringFormat(value=f"{key}"), + _glm_4_7_value_format(prop), + ConstStringFormat(value=""), + ] + ) + for key, prop in properties.items() + ] + content: Any = ( + StarFormat(content=pairs[0] if len(pairs) == 1 else OrFormat(elements=pairs)) + if pairs + else ConstStringFormat(value="") + ) + return TagFormat( + begin=f"{tool.function.name}", + content=content, + end="", + ) + + +def _glm_4_7_strict_tool_tag(tool: FunctionToolParam) -> TagFormat: + """Schema-pinned tag matching the xgrammar builtin's strict content.""" + return TagFormat( + begin=f"{tool.function.name}", + content=JSONSchemaFormat( + json_schema=tool.function.parameters + if tool.function.parameters is not None + else True, + style="glm_xml", + any_order=False, + ), + end="", + ) + + +@register_vllm_structural_tag("glm_4_7") +def get_glm_4_7_structural_tag( + function_tools: list[FunctionToolParam], + builtin_tools: list[BuiltinToolParam], + tool_choice: Any, + reasoning: bool, + token_suffix: str = "", +) -> StructuralTag: + """Build the GLM-4.7 structural tag, dispatching per tool on strictness. + + Args: + function_tools: Normalized function tools with ``strict`` pinned by + ``resolve_tool_strictness``. + builtin_tools: Unused; GLM-4.7 has no builtin tools. + tool_choice: Simplified tool choice (``auto`` / ``required`` / + ``forced``). + reasoning: Whether the grammar also covers the reasoning phase. + token_suffix: Unused; GLM-4.7 structural tokens are fixed. + + """ + del builtin_tools, token_suffix + + if function_tools and all(tool.function.strict for tool in function_tools): + # Strict-only requests stay byte-identical to the xgrammar builtin. + dumped_choice: Any = ( + {"type": "function", "function": {"name": function_tools[0].function.name}} + if tool_choice == "forced" + else tool_choice + ) + dumped_tools = [] + for tool in function_tools: + function: dict[str, Any] = {"name": tool.function.name, "strict": True} + if tool.function.description is not None: + function["description"] = tool.function.description + if tool.function.parameters is not None: + function["parameters"] = tool.function.parameters + dumped_tools.append({"type": "function", "function": function}) + return get_xgrammar_model_structural_tag( + model="glm_4_7", + tools=dumped_tools, + tool_choice=dumped_choice, + reasoning=reasoning, + ) + + tags = [ + _glm_4_7_strict_tool_tag(tool) + if tool.function.strict + else _glm_4_7_tool_tag(tool) + for tool in function_tools + ] + if tool_choice == "auto": + # Keep the trigger itself out of the excludes so tag dispatch works. + suffix_tag: Any = TriggeredTagsFormat( + triggers=[""], + tags=tags, + excludes=[e for e in _GLM_4_7_VALUE_EXCLUDES if e != ""], + ) + else: + suffix_tag = TagsWithSeparatorFormat( + tags=tags, + separator="", + at_least_one=True, + stop_after_first=tool_choice == "forced", + ) + if reasoning: + prefix_tag = TagFormat( + begin="", + content=AnyTextFormat(excludes=_GLM_4_7_VALUE_EXCLUDES), + end="", + ) + return StructuralTag(format=SequenceFormat(elements=[prefix_tag, suffix_tag])) + return StructuralTag(format=suffix_tag) diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 273c504f9d..0766f9b9ba 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -7,7 +7,6 @@ from collections.abc import Callable, Iterator, Sequence from contextlib import contextmanager from dataclasses import asdict from functools import cache, partial, wraps -from importlib.metadata import version from pathlib import Path from typing import Any, Literal, TypeAlias @@ -15,7 +14,6 @@ import huggingface_hub import torch import transformers.configuration_utils as hf_configuration_utils from huggingface_hub import constants -from packaging.version import Version from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE from transformers import GenerationConfig, PreTrainedConfig from transformers.configuration_utils import ALLOWED_LAYER_TYPES @@ -69,12 +67,6 @@ _DENSE_MODULE_TYPES = { "pylate.models.Dense.Dense", } -if Version(version("transformers")) < Version("5.0.0"): - raise ImportError( - "Support for Transformers v4 is deprecated and was removed in vLLM v0.24.0. " - "Please upgrade to Transformers v5: pip install --upgrade transformers" - ) - class LazyConfigDict(dict): def __getitem__(self, key): diff --git a/vllm/transformers_utils/processors/voxtral.py b/vllm/transformers_utils/processors/voxtral.py index 4cc23c4775..86c2aaf206 100644 --- a/vllm/transformers_utils/processors/voxtral.py +++ b/vllm/transformers_utils/processors/voxtral.py @@ -61,10 +61,8 @@ class MistralCommonFeatureExtractor(SequenceFeatureExtractor): """HF-compatible duck-typed ``fetch_audio``. Mirrors :meth:`transformers.SequenceFeatureExtractor.fetch_audio` so - :class:`transformers.ProcessorMixin.prepare_inputs_layout` (added in - transformers 5.10) works on this duck-typed feature extractor. Older - transformers versions never invoke this method, so the addition is a - no-op there. + :class:`transformers.ProcessorMixin.prepare_inputs_layout` works on this + duck-typed feature extractor. Accepts the same shapes as ``SequenceFeatureExtractor.fetch_audio``: @@ -77,11 +75,6 @@ class MistralCommonFeatureExtractor(SequenceFeatureExtractor): ``ProcessorMixin.prepare_inputs_layout`` always passes already-decoded audio (numpy array or torch tensor), so the str / list-of-str branches exist only to keep the contract identical to the upstream method. - - The semantics of ``transformers.audio_utils.is_valid_audio`` differ - between transformers versions (5.9 only accepts ndarray/tensor; 5.10 - also accepts ``list[float]``). We detect ``list[float]`` explicitly to - keep behavior identical across versions. """ from transformers.audio_utils import is_valid_audio @@ -89,9 +82,6 @@ class MistralCommonFeatureExtractor(SequenceFeatureExtractor): if is_valid_audio(audio_url_or_urls): return audio_url_or_urls if isinstance(audio_url_or_urls, (list, tuple)): - if audio_url_or_urls and isinstance(audio_url_or_urls[0], float): - # A single audio represented as ``list[float]``. - return audio_url_or_urls return [ self.fetch_audio(x, sampling_rate=sampling_rate) for x in audio_url_or_urls diff --git a/vllm/utils/flashinfer_moe_ep.py b/vllm/utils/flashinfer_moe_ep.py deleted file mode 100644 index 5bdd5cbf01..0000000000 --- a/vllm/utils/flashinfer_moe_ep.py +++ /dev/null @@ -1,339 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FlashInfer ``moe_ep`` helpers for DeepSeek V4 vLLM integration.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -import torch -import torch.nn as nn - -from vllm.config.kernel import FLASHINFER_MOE_EP_BACKENDS -from vllm.distributed import get_ep_group -from vllm.platforms import current_platform - -if TYPE_CHECKING: - from flashinfer.moe_ep import BootstrapConfig, MoEEpMegaLayer - - from vllm.config import VllmConfig - - -# Every mega kernel is Blackwell-only, so the arch is a property of the family -# rather than of a particular backend: validate it against the live device -# instead of encoding it in the backend name. An explicit allowlist so new -# archs (SM110/SM120, Rubin) fail loudly until flashinfer supports them. -FI_MOE_EP_SUPPORTED_CAPABILITIES = frozenset({(10, 0), (10, 3)}) - - -@dataclass(frozen=True) -class FiMoeEpBackendSpec: - """Static properties of one ``flashinfer_moe_ep_*`` backend string. - - The backend names the kernel *family*; the arch comes from the device and - the weight handling from the checkpoint, so this only has to carry which - megakernel to build and whether it needs NVSHMEM in the runtime set. - """ - - megakernel: str - needs_nvshmem: bool - - -FI_MOE_EP_BACKEND_SPECS: dict[str, FiMoeEpBackendSpec] = { - # Consumes an MXFP4 checkpoint verbatim (e2m1 weights + E8M0 per-32 - # scales) -- the same recipe the native deep_gemm mega path uses. - "flashinfer_moe_ep_mega_deep_gemm": FiMoeEpBackendSpec( - megakernel="deep_gemm_mega", - needs_nvshmem=False, - ), - # The checkpoint picks the weight path, not the kernel: an NVFP4 - # checkpoint is consumed prequantized (no round trip), while MXFP4 weights - # are dequantized to bf16 and requantized. See - # models/deepseek_v4/nvidia/fi_moe.py:ckpt_uses_nvfp4_experts. - "flashinfer_moe_ep_mega_cutedsl": FiMoeEpBackendSpec( - megakernel="nvfp4_cutedsl", - needs_nvshmem=True, - ), -} - -assert set(FI_MOE_EP_BACKEND_SPECS) == FLASHINFER_MOE_EP_BACKENDS - -_FI_RUNTIME_HANDLE: Any = None - - -def is_fi_moe_ep_backend(moe_backend: str) -> bool: - return moe_backend in FI_MOE_EP_BACKEND_SPECS - - -def fi_moe_ep_backend_spec(moe_backend: str) -> FiMoeEpBackendSpec: - try: - return FI_MOE_EP_BACKEND_SPECS[moe_backend] - except KeyError: - raise ValueError( - f"{moe_backend!r} is not a flashinfer moe_ep backend; expected " - f"one of {sorted(FI_MOE_EP_BACKEND_SPECS)}" - ) from None - - -def validate_fi_moe_ep_config(vllm_config: VllmConfig) -> None: - """Config-time checks for the mega-MoE backends, native and flashinfer.""" - moe_backend = vllm_config.kernel_config.moe_backend - if not is_fi_moe_ep_backend(moe_backend): - return - - # flashinfer validates the arch too, but not until the layer constructor - # runs during weight load; check here so the error names the flag the user - # actually typed. - capability = current_platform.get_device_capability() - if capability is not None: - cc = (capability.major, capability.minor) - if cc not in FI_MOE_EP_SUPPORTED_CAPABILITIES: - supported = ", ".join( - f"{m}.{n}" for m, n in sorted(FI_MOE_EP_SUPPORTED_CAPABILITIES) - ) - raise ValueError( - f"moe_backend={moe_backend!r} is only supported on compute " - f"capability {supported} (SM100/SM103), but this device is " - f"{cc[0]}.{cc[1]}." - ) - - if vllm_config.parallel_config.enable_eplb: - raise NotImplementedError( - f"EPLB is not supported with moe_backend={moe_backend!r}: the " - "flashinfer moe_ep experts neither apply the logical-to-physical " - "expert map nor report per-expert load, so rebalancing would move " - "weights without moving routing. Use " - "moe_backend=deep_gemm_mega_moe to run the mega path with EPLB." - ) - - -def make_fi_moe_ep_bootstrap() -> BootstrapConfig: - from flashinfer.moe_ep import BootstrapConfig - - ep = get_ep_group() - return BootstrapConfig( - world_size=ep.world_size, - rank=ep.rank_in_group, - process_group=ep.device_group, - auto_bootstrap=False, - # vLLM already bound this worker's (possibly remapped) device; - # without this the runtime would rebind to cuda:LOCAL_RANK|rank and - # launch weight transforms against another device's pointers. - device=torch.accelerator.current_device_index(), - ) - - -def megakernel_runtime_requirements(spec: FiMoeEpBackendSpec) -> frozenset[str]: - from flashinfer.moe_ep.core.runtime import NVSHMEM, TORCH_DIST - - if spec.needs_nvshmem: - return frozenset({TORCH_DIST, NVSHMEM}) - return frozenset({TORCH_DIST}) - - -def ensure_fi_moe_ep_runtime(vllm_config: VllmConfig) -> None: - """Acquire the process-wide flashinfer moe_ep runtime once per worker.""" - global _FI_RUNTIME_HANDLE - if _FI_RUNTIME_HANDLE is not None: - return - - from flashinfer.moe_ep import bootstrap_moe_ep_runtime - - bootstrap = make_fi_moe_ep_bootstrap() - spec = fi_moe_ep_backend_spec(vllm_config.kernel_config.moe_backend) - _FI_RUNTIME_HANDLE = bootstrap_moe_ep_runtime( - bootstrap, - megakernel_runtime_requirements(spec), - ) - - -def finalize_fi_moe_ep_runtime() -> None: - """Release the process-wide flashinfer moe_ep runtime.""" - global _FI_RUNTIME_HANDLE - if _FI_RUNTIME_HANDLE is None: - return - - from flashinfer.moe_ep import finalize_moe_ep_runtime - - finalize_moe_ep_runtime(_FI_RUNTIME_HANDLE) - _FI_RUNTIME_HANDLE = None - - -_E2M1_LUT = ( - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, -) - - -def _dequant_fp4_ue8m0_gran32( - packed: torch.Tensor, sf_ue8m0: torch.Tensor -) -> torch.Tensor: - """[rows, K//2] packed e2m1 + [rows, K//32] ue8m0-uint8 scales -> bf16 [rows, K].""" - raw = packed.view(torch.uint8) - lut = torch.tensor(_E2M1_LUT, dtype=torch.float32, device=raw.device) - vals = torch.empty( - raw.shape[0], raw.shape[1] * 2, dtype=torch.float32, device=raw.device - ) - vals[:, ::2] = lut[(raw & 0x0F).to(torch.int64)] - vals[:, 1::2] = lut[(raw >> 4).to(torch.int64)] - sf = (sf_ue8m0.to(torch.int32) << 23).view(torch.float32) - return (vals * sf.repeat_interleave(32, dim=-1)).to(torch.bfloat16) - - -def _dequant_expert_weights_to_bf16( - weight: torch.Tensor, scale: torch.Tensor -) -> torch.Tensor: - """[E, N, K//2] fp4 + [E, N, K//32] ue8m0 -> [E, N, K] bf16 (expert loop).""" - num_experts, n, k_half = weight.shape - out = torch.empty( - num_experts, n, k_half * 2, dtype=torch.bfloat16, device=weight.device - ) - for e in range(num_experts): - out[e] = _dequant_fp4_ue8m0_gran32(weight[e], scale[e]) - return out - - -def mega_moe_weight_pack_from_params( - w13_weight: nn.Parameter, - w13_weight_scale: nn.Parameter, - w2_weight: nn.Parameter, - w2_weight_scale: nn.Parameter, - *, - megakernel: str = "deep_gemm_mega", -): - from flashinfer.moe_ep import MoEWeightPack - - if megakernel == "deep_gemm_mega": - # Same fp4-e2m1 + ue8m0-per-32 recipe as the native path: pass verbatim, - # flashinfer runs the identical deep_gemm transform. - return MoEWeightPack( - w13=w13_weight.data, - w2=w2_weight.data, - w13_scale=w13_weight_scale.data, - w2_scale=w2_weight_scale.data, - ) - # The cutedsl kernel quantizes with its own recipe (nvfp4 - # e2m1+e4m3-per-16): dequantize the checkpoint fp4 to bf16 and let the - # backend preprocess requantize. Double quantization: outputs are close to - # but not bit-identical with the native path. - return MoEWeightPack( - w13=_dequant_expert_weights_to_bf16(w13_weight.data, w13_weight_scale.data), - w2=_dequant_expert_weights_to_bf16(w2_weight.data, w2_weight_scale.data), - ) - - -def build_fi_mega_config( - *, - intermediate_size: int, - top_k: int, - activation_clamp: float | None, - megakernel: str, -): - from flashinfer.moe_ep import ( - DeepGemmMegaMoeConfig, - MegaConfig, - Nvfp4CutedslMegaMoeConfig, - ) - - # fast_math selects approximate exp/rcp in DeepGEMM's fused SwiGLU - # epilogue; the cutedsl kernels accept it for API parity only. - if megakernel == "deep_gemm_mega": - mk = DeepGemmMegaMoeConfig( - intermediate_size=intermediate_size, - top_k=top_k, - activation_clamp=activation_clamp, - fast_math=True, - ) - elif megakernel == "nvfp4_cutedsl": - mk = Nvfp4CutedslMegaMoeConfig( - intermediate_size=intermediate_size, - top_k=top_k, - activation_clamp=activation_clamp, - fast_math=True, - ) - else: - raise ValueError(f"Unsupported fi_moe_ep megakernel {megakernel!r}") - - return MegaConfig( - megakernel=mk, - preprocess_weights=True, - quantize_input=True, - ) - - -# All MoE layers share one symmetric workspace, like the native path's -# class-level DeepseekV4MegaMoEExperts._symm_buffer_cache. Without this the -# fi path allocates one symm buffer PER LAYER (43x memory + cold working -# sets); the workspace is stateless across forwards (kernel tail-cleans) and -# layers execute sequentially on one stream, so sharing is safe. -def build_fi_mega_layer( - bootstrap: BootstrapConfig, - *, - vllm_config: VllmConfig, - num_experts: int, - max_tokens_per_rank: int, - hidden_size: int, - intermediate_size: int, - top_k: int, - activation_clamp: float | None, - weights, -) -> MoEEpMegaLayer: - from flashinfer.moe_ep import FleetParams, MoEEpLayer - - megakernel = fi_moe_ep_backend_spec( - vllm_config.kernel_config.moe_backend - ).megakernel - mega_config = build_fi_mega_config( - intermediate_size=intermediate_size, - top_k=top_k, - activation_clamp=activation_clamp, - megakernel=megakernel, - ) - layer = MoEEpLayer( - bootstrap=bootstrap, - fleet_params=FleetParams( - num_experts=num_experts, - max_tokens_per_rank=max_tokens_per_rank, - token_hidden_size=hidden_size, - ), - weights=weights, - backend=mega_config, - ) - from flashinfer.moe_ep import MoEEpMegaLayer - - if not isinstance(layer, MoEEpMegaLayer): - raise TypeError( - f"fi_moe_ep expected MoEEpMegaLayer, got {type(layer).__name__}" - ) - return layer - - -__all__ = [ - "FI_MOE_EP_BACKEND_SPECS", - "FiMoeEpBackendSpec", - "build_fi_mega_config", - "build_fi_mega_layer", - "ensure_fi_moe_ep_runtime", - "fi_moe_ep_backend_spec", - "finalize_fi_moe_ep_runtime", - "is_fi_moe_ep_backend", - "make_fi_moe_ep_bootstrap", - "mega_moe_weight_pack_from_params", - "megakernel_runtime_requirements", - "validate_fi_moe_ep_config", -] diff --git a/vllm/utils/numa_wrapper.sh b/vllm/utils/numa_wrapper.sh index 3c8b38919c..4f97ebf242 100755 --- a/vllm/utils/numa_wrapper.sh +++ b/vllm/utils/numa_wrapper.sh @@ -15,9 +15,8 @@ if ! command -v numactl >/dev/null 2>&1; then exit 1 fi -# shellcheck disable=SC1001 # verified equivalent behavior; TODO: revisit escaping in a follow-up cleanup PR case "${_VLLM_INTERNAL_NUMACTL_ARGS}" in - *[![:alnum:]\ \-\_=,./]*) + *[![:alnum:]\ _=,./-]*) echo "Invalid characters in _VLLM_INTERNAL_NUMACTL_ARGS" >&2 exit 1 ;; diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 57f4462759..3ca88992aa 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py @@ -130,6 +130,10 @@ class FlashAttnMLASparseMetadataBuilder( ) threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256) self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True) + self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 + + def update_draft_decode_metadata(self, _metadata: SparseMLACommonMetadata) -> None: + pass class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): @@ -275,15 +279,27 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]): kv_rows, block_stride_rows = flat_kv_row_view( kv_c_and_k_pe_cache, attn_metadata.block_size ) - topk_indices, valid_counts = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token[:num_actual_toks], - attn_metadata.block_table, - topk_indices, - BLOCK_SIZE=attn_metadata.block_size, - BLOCK_STRIDE_ROWS=block_stride_rows, - NUM_TOPK_TOKENS=topk_indices.shape[1], - return_valid_counts=True, - ) + if ( + index_group is not None + and index_group.num_layers > 1 + and attn_metadata.num_decode_tokens == num_actual_toks + ): + topk_indices, valid_counts = self._convert_logical_to_physical_topk( + topk_indices, + attn_metadata, + block_stride_rows=block_stride_rows, + return_valid_counts=True, + ) + else: + topk_indices, valid_counts = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_actual_toks], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + BLOCK_STRIDE_ROWS=block_stride_rows, + NUM_TOPK_TOKENS=topk_indices.shape[1], + return_valid_counts=True, + ) return ( self._run_mqa_kernel( q_nope, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 20a52af82b..5cfaf6ca5f 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -669,6 +669,7 @@ class DeepseekV32IndexerMetadata: decode: DeepSeekV32IndexerDecodeMetadata | None = None prefill: DeepseekV32IndexerPrefillMetadata | None = None + positions: torch.Tensor | None = None @triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) @@ -758,6 +759,7 @@ class KpoolTailMetadataBuilder(AttentionMetadataBuilder): _cudagraph_support = AttentionCGSupport.ALWAYS supports_update_block_table = False + supports_draft_decode_metadata_update = True reorder_batch_threshold = None def __init__( @@ -807,6 +809,21 @@ class KpoolTailMetadataBuilder(AttentionMetadataBuilder): num_decode_tokens=num_decode_tokens, num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, + positions=positions, + ) + + def update_draft_decode_metadata( + self, + metadata: DeepseekV32IndexerMetadata, + ) -> None: + assert metadata.positions is not None + slot_mapping = metadata.slot_mapping[: metadata.num_decode_tokens] + slot_mapping.div_(self.kv_cache_spec.block_size, rounding_mode="floor") + slot_mapping.mul_(self.kv_cache_spec.block_size) + slot_mapping.add_( + metadata.positions[: metadata.num_decode_tokens].remainder( + self.kv_cache_spec.block_size + ) ) @@ -1016,6 +1033,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) self.indexer_decode_block_table_buffer: torch.Tensor | None = None self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens + # DCP not supported yet + self.supports_draft_decode_metadata_update = self.dcp_world_size == 1 def _dcp_localize_decode_seq_lens( self, @@ -1668,6 +1687,50 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): return attn_metadata + def update_draft_decode_metadata( + self, + metadata: DeepseekV32IndexerMetadata, + ) -> None: + decode = metadata.decode + if decode is None or metadata.num_decode_tokens == 0: + return + + assert metadata.num_prefills == 0 + assert metadata.num_decodes == metadata.num_decode_tokens + assert decode.seq_lens.numel() == metadata.num_decode_tokens + assert self.dcp_world_size == 1 + + if self.compress_ratio > 1: + get_compressed_slot_mapping( + metadata.num_decode_tokens, + self.arange_buffer[: metadata.num_decode_tokens], + self.arange_buffer[: metadata.num_decode_tokens + 1], + metadata.seq_lens, + decode.block_table, + self.kv_cache_spec.num_states, + self.compress_ratio, + out=metadata.slot_mapping, + ) + torch.div( + metadata.seq_lens, + self.compress_ratio, + rounding_mode="floor", + out=decode.seq_lens.view(-1), + ) + else: + decode.seq_lens.view(-1).copy_(metadata.seq_lens) + decode.decode_lens.fill_(1) + + if current_platform.is_cuda() and has_deep_gemm(): + schedule_metadata = get_paged_mqa_logits_metadata( + decode.seq_lens, + self.kv_cache_spec.num_states, + self.num_sms, + indices=decode.indices, + ) + assert schedule_metadata.shape == decode.schedule_metadata.shape + decode.schedule_metadata.copy_(schedule_metadata) + def build_prefill_chunk_metadata( start_idx: int, diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index e676576fb7..a3412aa835 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -32,18 +32,6 @@ logger = init_logger(__name__) FP8_DTYPE = current_platform.fp8_dtype() -@functools.cache -def _get_aiter_topk_ops() -> tuple[Callable[..., None], Callable[..., None]] | None: - try: - from aiter.ops.topk import ( - top_k_per_row_decode, - top_k_per_row_prefill, - ) - except ImportError: - return None - return top_k_per_row_prefill, top_k_per_row_decode - - @functools.cache def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: from vllm._aiter_ops import rocm_aiter_ops @@ -58,9 +46,6 @@ def _get_aiter_sparse_prefill_opus() -> Callable[..., torch.Tensor] | None: return pa_sparse_prefill_opus -_GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN = 64 * 1024 -_GFX950_C4A_NATIVE_MAX_ROWS = 256 -_GFX950_DSV4_NATIVE_MAX_COLUMNS = 1024 * 1024 # Conservative perf gate, not a correctness bound: OPUS is correct for any query # count, but Triton stays faster below this measured crossover. _GFX950_AITER_SPARSE_PREFILL_OPUS_MIN_QUERIES = 1024 @@ -71,43 +56,6 @@ def _indexer_k_is_c4a_block_flat(compress_ratio: int) -> bool: return compress_ratio == 4 -def _get_aiter_top_k_kernel( - *, - is_prefill: bool, - compress_ratio: int, - num_rows: int, - max_valid_seq_len: int | None = None, - num_columns: int | None = None, - topk_tokens: int = 1024, - on_gfx950: bool = _ON_GFX950, -) -> Callable[..., None] | None: - if compress_ratio <= 1 or not on_gfx950: - return None - - if not is_prefill: - assert max_valid_seq_len is not None - if ( - topk_tokens == 512 - and 0 < num_rows <= 384 - and num_columns is not None - and num_columns <= _GFX950_DSV4_NATIVE_MAX_COLUMNS - ): - return None - # AITER v0.1.19 decode is one-block only. This measured gfx950 - # FP32/k=1024 compressed-row boundary is independent of the native - # split-count boundary in sampler.cu. - if ( - num_rows <= _GFX950_C4A_NATIVE_MAX_ROWS - and max_valid_seq_len > _GFX950_C4A_AITER_MAX_COMPRESSED_SEQ_LEN - ): - return None - - topk_ops = _get_aiter_topk_ops() - if topk_ops is None: - return None - return topk_ops[0] if is_prefill else topk_ops[1] - - @triton.jit def _localize_aiter_prefill_topk_kernel( indices_ptr, @@ -145,46 +93,20 @@ def _localize_aiter_prefill_topk( def _launch_aiter_top_k_per_row_prefill( - top_k_per_row_prefill: Callable[..., None], logits: torch.Tensor, row_starts: torch.Tensor, row_ends: torch.Tensor, indices: torch.Tensor, topk_tokens: int, ) -> None: - top_k_per_row_prefill( - logits, - row_starts, - row_ends, - indices, - None, - logits.shape[0], - logits.stride(0), - logits.stride(1), - k=topk_tokens, + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.indexer_top_k_prefill( + logits, row_starts, row_ends, indices, topk_tokens ) _localize_aiter_prefill_topk(indices, row_starts) -def _launch_aiter_top_k_per_row_decode( - top_k_per_row_decode: Callable[..., None], - logits: torch.Tensor, - seq_lens: torch.Tensor, - indices: torch.Tensor, - topk_tokens: int, -) -> None: - top_k_per_row_decode( - logits, - 1, - seq_lens.reshape(-1), - indices, - logits.shape[0], - logits.stride(0), - logits.stride(1), - k=topk_tokens, - ) - - @triton.jit def _indexer_k_quant_and_cache_kernel( k_ptr, # [num_tokens, head_dim] @@ -1195,6 +1117,8 @@ def rocm_aiter_sparse_attn_indexer( candidate_block_size: int = 0, candidate_write: bool = False, ) -> torch.Tensor: + from vllm._aiter_ops import rocm_aiter_ops + # careful! this will be None in dummy run forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata @@ -1347,14 +1271,12 @@ def rocm_aiter_sparse_attn_indexer( num_rows = logits.shape[0] - aiter_topk_kernel = _get_aiter_top_k_kernel( + if rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=True, compress_ratio=compress_ratio, num_rows=num_rows, - ) - if aiter_topk_kernel is not None: + ): _launch_aiter_top_k_per_row_prefill( - aiter_topk_kernel, logits, chunk.cu_seqlen_ks, chunk.cu_seqlen_ke, @@ -1448,19 +1370,20 @@ def rocm_aiter_sparse_attn_indexer( max_compressed_seq_len = max_model_len else: max_compressed_seq_len = layer_attn_metadata.max_seq_len // compress_ratio - aiter_topk_kernel = _get_aiter_top_k_kernel( + if rocm_aiter_ops.is_indexer_top_k_supported( is_prefill=False, compress_ratio=compress_ratio, num_rows=num_rows, max_valid_seq_len=max_compressed_seq_len, + ) and not rocm_aiter_ops.dsv4_indexer_prefers_native_top_k( + num_rows=num_rows, num_columns=logits.shape[1], topk_tokens=topk_tokens, - ) - if aiter_topk_kernel is not None: - _launch_aiter_top_k_per_row_decode( - aiter_topk_kernel, + ): + rocm_aiter_ops.indexer_top_k_decode( logits, - decode_metadata.seq_lens, + 1, + decode_metadata.seq_lens.reshape(-1), topk_indices, topk_tokens, ) @@ -2190,6 +2113,7 @@ def _sparse_attn_prefill_ragged_kernel( num_kv, scale, HAS_ATTN_SINK: tl.constexpr, + OUT_DV: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_K: tl.constexpr, @@ -2280,7 +2204,7 @@ def _sparse_attn_prefill_ragged_kernel( + head_offsets[:, None] * out_stride_h + dim_offsets[None, :] * out_stride_d, out, - mask=head_mask[:, None] & dim_mask[None, :], + mask=head_mask[:, None] & (dim_offsets[None, :] < OUT_DV), ) @@ -3442,6 +3366,7 @@ def _rocm_sparse_attn_prefill_ragged_triton( attn_sink: torch.Tensor | None, nope_head_dim: int, rope_head_dim: int, + out: torch.Tensor | None = None, ) -> torch.Tensor: assert q.ndim == 3, f"expected q=[sq,h,d], got {q.shape}" assert kv.ndim == 2, f"expected kv=[skv,d], got {kv.shape}" @@ -3472,7 +3397,15 @@ def _rocm_sparse_attn_prefill_ragged_triton( block_d = triton.next_power_of_2(head_dim) block_k = 16 if head_dim >= 256 else 32 num_warps = 4 - out = torch.empty_like(q) + if out is None: + out = torch.empty_like(q) + else: + assert out.shape[:2] == q.shape[:2], ( + f"expected out=[{num_queries},{num_heads},*], got {out.shape}" + ) + assert out.shape[-1] in (nope_head_dim, head_dim), ( + f"expected out width {nope_head_dim} or {head_dim}, got {out.shape[-1]}" + ) _sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))]( q, kv, @@ -3493,6 +3426,7 @@ def _rocm_sparse_attn_prefill_ragged_triton( kv.shape[0], float(scale), HAS_ATTN_SINK=has_attn_sink, + OUT_DV=out.shape[-1], BLOCK_H=block_h, BLOCK_D=block_d, BLOCK_K=block_k, @@ -3510,6 +3444,7 @@ def _rocm_sparse_attn_prefill_triton( nope_head_dim: int, rope_head_dim: int, topk_length: torch.Tensor | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: ragged_indices, ragged_indptr = build_ragged_indices_from_dense( indices, @@ -3527,6 +3462,7 @@ def _rocm_sparse_attn_prefill_triton( attn_sink=attn_sink, nope_head_dim=nope_head_dim, rope_head_dim=rope_head_dim, + out=out, ) @@ -4180,7 +4116,7 @@ def rocm_sparse_attn_prefill( return if ragged_indices is not None and ragged_indptr is not None: - output_chunk = _rocm_sparse_attn_prefill_ragged_triton( + _rocm_sparse_attn_prefill_ragged_triton( q=q, kv=kv.squeeze(1), indices=ragged_indices, @@ -4189,11 +4125,12 @@ def rocm_sparse_attn_prefill( attn_sink=None if attn_sink is None else attn_sink[: q.shape[1]], nope_head_dim=nope_head_dim, rope_head_dim=rope_head_dim, + out=output, ) else: assert indices is not None indices_2d = indices.reshape(indices.shape[0], -1) - output_chunk = _rocm_sparse_attn_prefill_triton( + _rocm_sparse_attn_prefill_triton( q=q, kv=kv.squeeze(1), indices=indices_2d, @@ -4202,8 +4139,8 @@ def rocm_sparse_attn_prefill( nope_head_dim=nope_head_dim, rope_head_dim=rope_head_dim, topk_length=topk_length, + out=output, ) - output.copy_(output_chunk[..., : output.shape[-1]].to(output.dtype)) def rocm_sparse_attn_decode( diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index daf0cd12c0..9b91bb1920 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -241,6 +241,7 @@ class KVCacheCoordinator(ABC): new_computed_blocks: tuple[Sequence[KVCacheBlock], ...], num_local_computed_tokens: int, num_external_computed_tokens: int, + skip_zeroing_group_ids: tuple[int, ...] = (), ) -> None: """Add the new computed blocks to the request. Optionally allocate new blocks for external computed tokens (if any). @@ -251,6 +252,8 @@ class KVCacheCoordinator(ABC): prefix cache. num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. + skip_zeroing_group_ids: Groups whose external-token blocks are + written by an async load and must not be zeroed. """ # A running request is already tracked in num_cached_block and won't @@ -274,11 +277,12 @@ class KVCacheCoordinator(ABC): num_external_computed_tokens, ) if num_external_computed_tokens > 0: - for manager in self.single_type_managers: + for i, manager in enumerate(self.single_type_managers): manager.allocate_external_computed_blocks( request_id, num_local_computed_tokens, num_external_computed_tokens, + record_for_zeroing=i not in skip_zeroing_group_ids, ) def allocate_new_blocks( diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 2d39a38651..37ceb4e474 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -381,6 +381,7 @@ class KVCacheManager: full_sequence_must_fit: bool = False, reserved_blocks: int = 0, has_scheduled_reqs: bool = True, + skip_zeroing_group_ids: tuple[int, ...] = (), ) -> KVCacheBlocks | None: """Add slots for a request with new tokens to append. @@ -413,6 +414,8 @@ class KVCacheManager: blocks an already in-flight (prefilling) sequence is relying on. has_scheduled_reqs: Whether any requests are already scheduled to run this step, controls whether watermark is applied. + skip_zeroing_group_ids: Groups whose external-token blocks will be + written by an async load and must not be zeroed concurrently. Blocks layout: ``` @@ -586,6 +589,7 @@ class KVCacheManager: new_computed_blocks=new_computed_block_list, num_local_computed_tokens=num_local_computed_tokens, num_external_computed_tokens=num_external_computed_tokens, + skip_zeroing_group_ids=skip_zeroing_group_ids, ) new_blocks = self.coordinator.allocate_new_blocks( @@ -863,20 +867,6 @@ class KVCacheManager: ids.extend(mgr.take_new_block_ids()) return ids - def get_zeroing_block_ids_in_range( - self, request_id: str, start_token: int, end_token: int - ) -> list[int]: - """The request's block ids covering [start_token, end_token), from - the groups whose new blocks are zeroed by the worker.""" - ids: list[int] = [] - for mgr in self.coordinator.single_type_managers: - if mgr.records_new_block_ids: - start_idx = start_token // mgr.block_size - end_idx = cdiv(end_token, mgr.block_size) - blocks = mgr.req_to_blocks[request_id] - ids.extend(blk.block_id for blk in blocks[start_idx:end_idx]) - return ids - def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None: """Re-record the request's blocks from start_token onwards for zeroing, e.g. blocks a failed async KV load left unwritten. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5585c9c70c..b390f1aadb 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -353,9 +353,6 @@ class Scheduler(SchedulerInterface): self.has_mamba_layers = kv_cache_config.has_mamba_layers self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing - # Blocks that async KV loads will overwrite this step, skipped from - # zeroing since the zeroing could race the out-of-band write. - self._skip_zero_block_ids: set[int] = set() self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) @@ -1230,6 +1227,11 @@ class Scheduler(SchedulerInterface): num_lookahead_tokens=effective_lookahead_tokens, num_external_computed_tokens=num_external_computed_tokens, delay_cache_blocks=load_kv_async, + skip_zeroing_group_ids=( + self.connector.get_loaded_kv_cache_group_ids(request) + if load_kv_async and self.connector is not None + else () + ), num_encoder_tokens=num_encoder_tokens, full_sequence_must_fit=self.scheduler_reserve_full_isl, reserved_blocks=reserved_blocks, @@ -1291,16 +1293,6 @@ class Scheduler(SchedulerInterface): request.num_computed_tokens = num_computed_tokens self._inflight_prefills.add(request) self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS) - if self.needs_kv_cache_zeroing: - # Skip zeroing of the blocks the async load will - # overwrite; the zeroing could race the write. - self._skip_zero_block_ids.update( - self.kv_cache_manager.get_zeroing_block_ids_in_range( - request_id, - num_new_local_computed_tokens, - num_computed_tokens, - ) - ) skip_request(request_queue) continue @@ -1545,12 +1537,6 @@ class Scheduler(SchedulerInterface): new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() if not self.needs_kv_cache_zeroing: return None - - if self._skip_zero_block_ids: - skip = self._skip_zero_block_ids - new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip] - skip.clear() - return new_block_ids_to_zero or None def _preempt_request( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2622b3555c..7f23a886ed 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -331,6 +331,7 @@ class SingleTypeKVCacheManager(ABC): request_id: str, num_local_computed_tokens: int, num_external_computed_tokens: int, + record_for_zeroing: bool = True, ) -> None: """Allocate new blocks for external (KV-connector) computed tokens. @@ -342,6 +343,8 @@ class SingleTypeKVCacheManager(ABC): request_id: The request ID. num_local_computed_tokens: The number of local computed tokens. num_external_computed_tokens: The number of external computed tokens. + record_for_zeroing: Whether the new blocks need zeroing. False when + the load writes them after this step, which zeroing would race. """ num_total_computed_tokens = ( @@ -363,7 +366,7 @@ class SingleTypeKVCacheManager(ABC): ) allocated_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(allocated_blocks) - if self._record_new_block_ids: + if self._record_new_block_ids and record_for_zeroing: self.new_block_ids.extend(b.block_id for b in allocated_blocks) def allocate_new_blocks( @@ -1201,13 +1204,15 @@ class CircularBufferManager(FullAttentionManager): supports_fine_grained_hash_lookup: ClassVar[bool] = False - def _claim_ring_block(self, request_id: str) -> list[KVCacheBlock]: + def _claim_ring_block( + self, request_id: str, record_for_zeroing: bool = True + ) -> list[KVCacheBlock]: req_blocks = self.req_to_blocks[request_id] if req_blocks: return [] new_blocks = self.block_pool.get_new_blocks(1) req_blocks.extend(new_blocks) - if self._record_new_block_ids: + if self._record_new_block_ids and record_for_zeroing: self.new_block_ids.extend(block.block_id for block in new_blocks) return new_blocks @@ -1234,8 +1239,9 @@ class CircularBufferManager(FullAttentionManager): request_id: str, num_local_computed_tokens: int, num_external_computed_tokens: int, + record_for_zeroing: bool = True, ) -> None: - self._claim_ring_block(request_id) + self._claim_ring_block(request_id, record_for_zeroing) @classmethod def find_longest_cache_hit( @@ -2175,6 +2181,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager): request_id: str, num_local_computed_tokens: int, num_external_computed_tokens: int, + record_for_zeroing: bool = True, ) -> None: # Cross-attention does not use prefix caching / external KV loads. return @@ -2320,6 +2327,7 @@ class HiSparseSourceManager(FullAttentionManager): request_id: str, num_local_computed_tokens: int, num_external_computed_tokens: int, + record_for_zeroing: bool = True, ) -> None: if num_external_computed_tokens <= 0: return @@ -2330,7 +2338,10 @@ class HiSparseSourceManager(FullAttentionManager): request_id, num_local_computed_tokens + num_external_computed_tokens ) super().allocate_external_computed_blocks( - request_id, num_local_computed_tokens, num_external_computed_tokens + request_id, + num_local_computed_tokens, + num_external_computed_tokens, + record_for_zeroing, ) def cache_blocks( @@ -2465,6 +2476,7 @@ class HiSparseHotManager(_HiSparseAuxiliaryManager): request_id: str, num_local_computed_tokens: int, num_external_computed_tokens: int, + record_for_zeroing: bool = True, ) -> None: self.require_hot(request_id) @@ -2527,6 +2539,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager): request_id: str, num_local_computed_tokens: int, num_external_computed_tokens: int, + record_for_zeroing: bool = True, ) -> None: """Reserve a writable final page; earlier imported pages stay on host.""" assert num_external_computed_tokens > 0 diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index b1cc2edbc4..4a7a91cb8e 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -599,6 +599,8 @@ class AsyncLLM(EngineClient): async def handle_inputs(): cancelled = False + errored = False + any_added = False try: async for input_chunk in input_stream: sp = input_chunk.sampling_params @@ -623,18 +625,23 @@ class AsyncLLM(EngineClient): self.model_config, input_chunk.prompt ) await self._add_request(req, prompt_text, None, 0, queue) + any_added = True except (asyncio.CancelledError, GeneratorExit): cancelled = True except Exception as error: # Wrap in InputStreamError so generate() can propagate it # without wrapping in EngineGenerateError. queue.put(InputStreamError(error)) + errored = True finally: queue._input_stream_task = None if not cancelled: - # Send empty final request to indicate that inputs have - # finished. Don't send if cancelled (session was aborted). - await self._add_request(final_req, None, None, 0, queue) + if any_added: + # Send empty final request to indicate that inputs have + # finished. Don't send if cancelled (session was aborted). + await self._add_request(final_req, None, None, 0, queue) + elif not errored: + queue.put(STREAM_FINISHED) # Ensure output handler is running. self._run_output_handler() diff --git a/vllm/v1/engine/input_processor.py b/vllm/v1/engine/input_processor.py index fa68eb8ec2..8b5a69bcf4 100644 --- a/vllm/v1/engine/input_processor.py +++ b/vllm/v1/engine/input_processor.py @@ -423,6 +423,18 @@ class InputProcessor: parameter="routed_experts_prompt_start", value=sampling_params.routed_experts_prompt_start, ) + rows = sampling_params.prompt_logprob_token_ids + if rows is not None: + start = sampling_params.prompt_logprob_start or 0 + num_rows = max(prompt_len - 1 - start, 0) + if len(rows) != num_rows: + raise VLLMValidationError( + f"prompt_logprob_token_ids has {len(rows)} rows, but the " + f"prompt has {num_rows} scored rows " + f"(prompt_len - 1 - prompt_logprob_start).", + parameter="prompt_logprob_token_ids", + value=len(rows), + ) # If unset max tokens, then generate up to the max_model_len. if sampling_params.max_tokens is None: sampling_params.max_tokens = ( diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index d14b3b9755..c3fb3efe7e 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -323,6 +323,14 @@ class LLMEngine: engine_core_timestamp=outputs.timestamp, iteration_stats=iteration_stats, ) + + mm_processor_cache = self.renderer.mm_processor_cache + if mm_processor_cache is not None: + for engine_core_output in outputs.outputs: + if engine_core_output.mm_cache_miss_hashes: + for mm_hash in engine_core_output.mm_cache_miss_hashes: + mm_processor_cache.invalidate(mm_hash) + self.output_processor.update_scheduler_stats(outputs.scheduler_stats) # 3) Abort any reqs that finished due to stop strings. diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 89b3a9cf09..c22c5474c7 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -16,6 +16,7 @@ from vllm.outputs import ( CompletionOutput, PoolingOutput, PoolingRequestOutput, + RequestError, RequestOutput, SamplingMask, ) @@ -304,6 +305,7 @@ class RequestState: stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, ec_transfer_params: dict[str, Any] | None = None, + error: RequestError | None = None, ) -> RequestOutput | PoolingRequestOutput | None: finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY @@ -342,6 +344,7 @@ class RequestState: external_req_id, [self._new_pooling_output(pooling_output)], finished, + error=error, ) output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) @@ -369,6 +372,7 @@ class RequestState: finished: bool, kv_transfer_params: dict[str, Any] | None = None, ec_transfer_params: dict[str, Any] | None = None, + error: RequestError | None = None, ) -> RequestOutput | PoolingRequestOutput: # If prompt embeds were used, put placeholder prompt token ids prompt_token_ids = self.prompt_token_ids @@ -385,6 +389,7 @@ class RequestState: num_cached_tokens=self.num_cached_tokens, prompt_token_ids=prompt_token_ids, finished=finished, + error=error, ) assert self.logprobs_processor is not None if self.output_kind == RequestOutputKind.DELTA: @@ -715,6 +720,19 @@ class OutputProcessor: if engine_core_output.spec_decode_metrics is not None: req_state.spec_decode_metrics = engine_core_output.spec_decode_metrics + request_error = None + if ( + req_state.detokenizer is None + and engine_core_output.mm_cache_miss_hashes + and finish_reason == FinishReason.ERROR + ): + request_error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + pooling_output = EMPTY_CPU_TENSOR + if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None @@ -742,6 +760,7 @@ class OutputProcessor: stop_reason, kv_transfer_params, ec_transfer_params, + error=request_error, ): if req_state.streaming_input: request_output.finished = False diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 9df749f365..203d1b407c 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -942,6 +942,7 @@ class CoreEngineActorManager: runtime_env = RuntimeEnv( env_vars=self.env_vars_dict | {"VLLM_ELASTIC_EP_SCALE_UP_LAUNCH": "1"} ) + new_actors = [] for i, (pg, local_rank) in enumerate(zip(placement_groups, local_dp_ranks)): rank = cur_data_parallel_size + i dp_vllm_config = copy.deepcopy(cur_vllm_config) @@ -985,6 +986,7 @@ class CoreEngineActorManager: local_dp_rank=local_rank, ) ) + new_actors.append(actor) if local_client: self.local_engine_actors.append(actor) @@ -993,14 +995,8 @@ class CoreEngineActorManager: self.created_placement_groups.append(pg) self.placement_group_is_local.append(local_client) - actors = ( - self.local_engine_actors[-new_local_engines:] - if new_local_engines > 0 - else [] - ) + self.remote_engine_actors[-(len(placement_groups) - new_local_engines) :] - - ray.get([actor.wait_for_init.remote() for actor in actors]) - for actor in actors: + ray.get([actor.wait_for_init.remote() for actor in new_actors]) + for actor in new_actors: ref = actor.run.remote() self.run_refs.append(ref) self.actor_run_ref_dict[actor] = ref diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 797abed5ea..a56febf702 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -689,10 +689,10 @@ class SingleDirectionOffloadingHandler: # Stores must wait for the model to finish writing the KV they read. # Loads must wait for pending writes (including zeroing) to their - # destination blocks: the scheduler's _skip_zero_block_ids only edits - # the step being scheduled and cannot retract zeroing shipped in an - # earlier step for a since-reallocated block; with async scheduling - # nothing else orders that zeroing against this copy. + # destination blocks: the scheduler leaves a block unzeroed only in + # the step that allocates it for the load and cannot retract zeroing + # shipped in an earlier step for a since-reallocated block; with async + # scheduling nothing else orders that zeroing against this copy. stream.wait_stream(current_platform.current_stream()) if self._transfers: last_transfer: Transfer = self._transfers[-1] diff --git a/vllm/v1/pool/late_interaction_runner.py b/vllm/v1/pool/late_interaction_runner.py index 7ea39f5754..6e8fc31861 100644 --- a/vllm/v1/pool/late_interaction_runner.py +++ b/vllm/v1/pool/late_interaction_runner.py @@ -42,6 +42,18 @@ class LateInteractionRunner: self._query_uses.clear() self._doc_query_keys.clear() + def release_queries(self, query_keys: Iterable[str]) -> None: + query_keys = set(query_keys) + for query_key in query_keys: + self._query_cache.pop(query_key, None) + self._query_uses.pop(query_key, None) + + self._doc_query_keys = { + req_id: query_key + for req_id, query_key in self._doc_query_keys.items() + if query_key not in query_keys + } + def register_request( self, req_id: str, pooling_params: PoolingParams | None ) -> None: diff --git a/vllm/v1/serial_utils.py b/vllm/v1/serial_utils.py index 882bc44917..13d725410a 100644 --- a/vllm/v1/serial_utils.py +++ b/vllm/v1/serial_utils.py @@ -532,7 +532,7 @@ class PydanticMsgspecMixin: called on every validation. """ msgspec_fields = {f.name: f for f in msgspec.structs.fields(source_type)} - type_hints = get_type_hints(source_type) + type_hints = get_type_hints(source_type, include_extras=True) # Build the Pydantic typed_dict_field for each msgspec field fields = {} diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py index 36d00313ae..3cfe334a90 100644 --- a/vllm/v1/watermarking/spec_decode.py +++ b/vllm/v1/watermarking/spec_decode.py @@ -44,8 +44,8 @@ class DraftWatermarker: dtype=torch.int64, device=device, ) + self.prompt_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) self.all_token_ids: torch.Tensor | None = None - self.prompt_lens: torch.Tensor | None = None self.total_lens: torch.Tensor | None = None def prepare( @@ -60,7 +60,7 @@ class DraftWatermarker: self.contexts[:num_reqs].copy_(contexts) self.enabled[:num_reqs].copy_(enabled) self.all_token_ids = all_token_ids - self.prompt_lens = prompt_lens + self.prompt_lens.copy_(prompt_lens) self.total_lens = total_lens def _sampling_state( @@ -79,7 +79,6 @@ class DraftWatermarker: if self.deduplicate_contexts == "none": return steps, enabled assert self.all_token_ids is not None - assert self.prompt_lens is not None assert self.total_lens is not None enabled = draft_watermarking_mask( self.all_token_ids, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b40fb4eac1..1b9d0608e2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1012,6 +1012,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.pooling_runner is not None: self.pooling_runner.clear() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + if self.pooling_runner is not None: + self.pooling_runner.release_late_interaction_query_cache(query_keys) + @torch.inference_mode() def profile_cudagraph_memory(self) -> int: """Estimate the GPU memory required to capture CUDA graphs.""" diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index 1a81046b17..be27019dd0 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/vllm/v1/worker/gpu/pool/pooling_runner.py @@ -98,6 +98,9 @@ class PoolingRunner: def clear(self) -> None: self.late_interaction_runner.clear() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + self.late_interaction_runner.release_queries(query_keys) + def _get_pooling_metadata( self, input_batch: InputBatch, diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py index 24fdb6188e..0fbdbc4134 100644 --- a/vllm/v1/worker/gpu/sample/prompt_logprob.py +++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py @@ -22,7 +22,8 @@ CHUNK_SIZE = 1024 @dataclass class _TokenIdScores: - token_ids: torch.Tensor # [num_ids], on device + token_ids: torch.Tensor # [num_rows, num_ids], on device + pad: torch.Tensor # [num_rows, num_ids] bool, True at padded columns start: int # first prompt row to score scores: torch.Tensor | None = None # [num_rows, num_ids], filled per chunk @@ -52,11 +53,13 @@ class PromptLogprobsWorker: if uses_prompt_logprobs: self.in_progress_prompt_logprobs[req_id] = [] if sampling_params.prompt_logprob_token_ids is not None: + ids = sampling_params.prompt_logprob_token_ids + # Decoded arrays may be read-only; torch.from_numpy warns on those. + if not ids.flags.writeable: + ids = ids.copy() + ids = async_tensor_h2d(ids, self.device, torch.int64) self.token_id_scores[req_id] = _TokenIdScores( - async_tensor_h2d( - sampling_params.prompt_logprob_token_ids, self.device, torch.int64 - ), - sampling_params.prompt_logprob_start or 0, + ids.clamp_min(0), ids < 0, sampling_params.prompt_logprob_start or 0 ) def remove_request(self, req_id: str) -> None: @@ -85,6 +88,7 @@ class PromptLogprobsWorker: # Skip decode steps and, as compute_prompt_logprobs does, requests # resumed after preemption: their scores were already emitted. if chunk_start >= prompt_len or prompt_len < input_batch.prefill_len_np[i]: + self.token_id_scores.pop(req_id, None) continue if req.scores is None: if chunk_start > req.start: @@ -93,10 +97,7 @@ class PromptLogprobsWorker: # emitting a buffer with unwritten rows. self.token_id_scores.pop(req_id, None) continue - req.scores = hidden_states.new_empty( - (max(prompt_len - 1 - req.start, 0), len(req.token_ids)), - dtype=torch.float32, - ) + req.scores = torch.empty_like(req.token_ids, dtype=torch.float32) # The last prompt row predicts the first decode token; skip it. lo = max(req.start, chunk_start) hi = min(chunk_end, prompt_len - 1) @@ -104,15 +105,16 @@ class PromptLogprobsWorker: for a in range(lo, hi, CHUNK_SIZE): b = min(a + CHUNK_SIZE, hi) logits = logits_fn(hidden_states[base + a : base + b]) - ids = req.token_ids.expand(b - a, -1) + ids = req.token_ids[a - req.start : b - req.start] req.scores[a - req.start : b - req.start] = ( logits.gather(-1, ids) if logits_mode else compute_token_logprobs(logits, ids) ) if chunk_end >= prompt_len: + req.scores.masked_fill_(req.pad, float("-inf")) out[req_id] = req.scores - req.scores = None + self.token_id_scores.pop(req_id, None) return out def compute_prompt_logprobs( diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 0728ddb87c..77a9ece360 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -373,6 +373,7 @@ class DraftModelSpeculator(BaseSpeculator): kv_cache_config=self.kv_cache_config, causal=causal, seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, + positions=self.input_buffers.positions[:num_tokens], is_prefilling=self.draft_is_prefilling[:num_reqs_padded], ) return attn_metadata diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 9a0a4eeda9..047fdc5ff0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -993,6 +993,9 @@ class GPUModelRunner( self.encoder_cache.clear() self.late_interaction_runner.clear() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + self.late_interaction_runner.release_queries(query_keys) + def _get_positions(self, num_tokens: Any): if isinstance(num_tokens, int): if self.uses_mrope: @@ -6531,6 +6534,9 @@ class GPUModelRunner( if hasattr(self, "kv_cache_config"): delattr(self, "kv_cache_config") self.cache_config.num_gpu_blocks = None + # Profiling may have rebuilt the InputBatch outside the worker's + # "runtime" pool; the real initialize_kv_cache rebuilds it. + self._init_block_sizes = [] for layer in self.compilation_config.static_forward_context.values(): if hasattr(layer, "kv_cache"): diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 15e2ee0611..55d56919a4 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -841,10 +841,17 @@ class Worker(WorkerBase): if mem_pool_context is None: mem_pool_context = self._maybe_get_memory_pool_context(tag="kv_cache") - self.model_runner.initialize_kv_cache( - kv_cache_config, - kv_cache_allocation_context=mem_pool_context, + # Offload KV-init state, except on XPU, whose outermost pool would win. + runtime_pool = ( + nullcontext() + if current_platform.is_xpu() + else self._maybe_get_memory_pool_context(tag="runtime") ) + with runtime_pool: + self.model_runner.initialize_kv_cache( + kv_cache_config, + kv_cache_allocation_context=mem_pool_context, + ) # Build KV-zero metadata outside the CuMem pool so the bookkeeping # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 017b1230cd..95a0c3475d 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -143,6 +143,13 @@ class WorkerBase: if callable(reset_fn): reset_fn() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + release_fn = getattr( + self.model_runner, "release_late_interaction_query_cache", None + ) + if callable(release_fn): + release_fn(query_keys) + def get_model(self) -> nn.Module: raise NotImplementedError